<commit_before>package work

import (
	// "fmt"
	"github.com/garyburd/redigo/redis"
	"sort"
	"strconv"
	"strings"
)

type Client struct {
	namespace string // eg, "myapp-work"
	pool      *redis.Pool
}

func NewClient(namespace string, pool *redis.Pool) *Client {
	return &Client{
		namespace: namespace,
		pool:      pool,
	}
}

type WorkerPoolHeartbeat struct {
	WorkerPoolID string
	StartedAt    int64
	HeartbeatAt  int64

	JobNames    []string
	Concurrency uint
	Host        string
	Pid         int

	WorkerIDs []string
}

func (c *Client) WorkerPoolHeartbeats() ([]*WorkerPoolHeartbeat, error) {
	conn := c.pool.Get()
	defer conn.Close()

	workerPoolsKey := redisKeyWorkerPools(c.namespace)

	workerPoolIDs, err := redis.Strings(conn.Do("SMEMBERS", workerPoolsKey))
	if err != nil {
		return nil, err
	}
	sort.Strings(workerPoolIDs)

	for _, wpid := range workerPoolIDs {
		key := redisKeyHeartbeat(c.namespace, wpid)
		conn.Send("HGETALL", key)
	}

	if err := conn.Flush(); err != nil {
		logError("worker_pool_statuses.flush", err)
		return nil, err
	}

	heartbeats := make([]*WorkerPoolHeartbeat, 0, len(workerPoolIDs))

	for _, wpid := range workerPoolIDs {
		vals, err := redis.Strings(conn.Receive())
		if err != nil {
			logError("worker_pool_statuses.receive", err)
			return nil, err
		}

		heartbeat := &WorkerPoolHeartbeat{
			WorkerPoolID: wpid,
		}

		for i := 0; i < len(vals)-1; i += 2 {
			key := vals[i]
			value := vals[i+1]

			var err error
			if key == "heartbeat_at" {
				heartbeat.HeartbeatAt, err = strconv.ParseInt(value, 10, 64)
			} else if key == "started_at" {
				heartbeat.StartedAt, err = strconv.ParseInt(value, 10, 64)
			} else if key == "job_names" {
				heartbeat.JobNames = strings.Split(value, ",")
				sort.Strings(heartbeat.JobNames)
			} else if key == "concurrency" {
				var vv uint64
				vv, err = strconv.ParseUint(value, 10, 0)
				heartbeat.Concurrency = uint(vv)
			} else if key == "host" {
				heartbeat.Host = value
			} else if key == "pid" {
				var vv int64
				vv, err = strconv.ParseInt(value, 10, 0)
				heartbeat.Pid = int(vv)
			} else if key == "worker_ids" {
				heartbeat.WorkerIDs = strings.Split(value, ",")
				sort.Strings(heartbeat.WorkerIDs)
			}
			if err != nil {
				logError("worker_pool_statuses.parse", err)
				return nil, err
			}
		}

		heartbeats = append(heartbeats, heartbeat)
	}

	return heartbeats, nil
}

type WorkerObservation struct {
	WorkerID string
	IsBusy   bool

	// If IsBusy:
	JobName   string
	JobID     string
	StartedAt int64
	ArgsJSON  string
	Checkin   string
	CheckinAt int64
}

func (c *Client) WorkerObservations() ([]*WorkerObservation, error) {
	conn := c.pool.Get()
	defer conn.Close()

	hbs, err := c.WorkerPoolHeartbeats()
	if err != nil {
		logError("worker_observations.worker_pool_heartbeats", err)
		return nil, err
	}

	var workerIDs []string
	for _, hb := range hbs {
		workerIDs = append(workerIDs, hb.WorkerIDs...)
	}

	for _, wid := range workerIDs {
		key := redisKeyWorkerStatus(c.namespace, wid) // TODO: rename this func
		conn.Send("HGETALL", key)
	}

	if err := conn.Flush(); err != nil {
		logError("worker_observations.flush", err)
		return nil, err
	}

	observations := make([]*WorkerObservation, 0, len(workerIDs))

	for _, wid := range workerIDs {
		vals, err := redis.Strings(conn.Receive())
		if err != nil {
			logError("worker_observations.receive", err)
			return nil, err
		}

		ob := &WorkerObservation{
			WorkerID: wid,
		}

		for i := 0; i < len(vals)-1; i += 2 {
			key := vals[i]
			value := vals[i+1]

			ob.IsBusy = true

			var err error
			if key == "job_name" {
				ob.JobName = value
			} else if key == "job_id" {
				ob.JobID = value
			} else if key == "started_at" {
				ob.StartedAt, err = strconv.ParseInt(value, 10, 64)
			} else if key == "args" {
				ob.ArgsJSON = value
			} else if key == "checkin" {
				ob.Checkin = value
			} else if key == "checkin_at" {
				ob.CheckinAt, err = strconv.ParseInt(value, 10, 64)
			}
			if err != nil {
				logError("worker_observations.parse", err)
				return nil, err
			}
		}

		observations = append(observations, ob)
	}

	return observations, nil
}

type Queue struct {
	JobName string
	Count   int64
	Latency int64
}

func (c *Client) Queues() ([]*Queue, error) {
	conn := c.pool.Get()
	defer conn.Close()

	key := redisKeyKnownJobs(c.namespace)
	jobNames, err := redis.Strings(conn.Do("SMEMBERS", key))
	if err != nil {
		return nil, err
	}
	sort.Strings(jobNames)

	for _, jobName := range jobNames {
		conn.Send("LLEN", redisKeyJobs(c.namespace, jobName))
	}

	if err := conn.Flush(); err != nil {
		logError("client.queues.flush", err)
		return nil, err
	}

	queues := make([]*Queue, 0, len(jobNames))

	for _, jobName := range jobNames {
		count, err := redis.Int64(conn.Receive())
		if err != nil {
			logError("client.queues.receive", err)
			return nil, err
		}

		queue := &Queue{
			JobName: jobName,
			Count:   count,
		}

		queues = append(queues, queue)
	}

	for _, s := range queues {
		if s.Count > 0 {
			conn.Send("LINDEX", redisKeyJobs(c.namespace, s.JobName), -1)
		}
	}

	if err := conn.Flush(); err != nil {
		logError("client.queues.flush2", err)
		return nil, err
	}

	now := nowEpochSeconds()

	for _, s := range queues {
		if s.Count > 0 {
			b, err := redis.Bytes(conn.Receive())
			if err != nil {
				logError("client.queues.receive2", err)
				return nil, err
			}

			job, err := newJob(b, nil, nil)
			if err != nil {
				logError("client.queues.new_job", err)
			}
			s.Latency = now - job.EnqueuedAt
		}
	}

	return queues, nil
}

type FailedJob struct {
	RetryAt int64
	Job
}

type ScheduledJob struct {
	RunAt int64
	*Job
}

func (c *Client) ScheduledJobs(page uint) ([]*ScheduledJob, error) {
	conn := c.pool.Get()
	defer conn.Close()

	if page == 0 {
		page = 1
	}

	key := redisKeyScheduled(c.namespace)
	values, err := redis.Values(conn.Do("ZRANGEBYSCORE", key, "-inf", "+inf", "WITHSCORES", "LIMIT", (page-1)*20, 20))
	if err != nil {
		logError("client.scheduled_jobs.values", err)
		return nil, err
	}

	var jobsWithScores []struct {
		JobBytes []byte
		Score    int64
	}

	if err := redis.ScanSlice(values, &jobsWithScores); err != nil {
		logError("client.scheduled_jobs.scan_slice", err)
		return nil, err
	}

	jobs := make([]*ScheduledJob, 0, len(jobsWithScores))

	for _, jws := range jobsWithScores {
		job, err := newJob(jws.JobBytes, nil, nil)
		if err != nil {
			logError("client.scheduled_jobs.new_job", err)
			return nil, err
		}

		jobs = append(jobs, &ScheduledJob{RunAt: jws.Score, Job: job})
	}

	return jobs, nil
}
<commit_msg>Get failed and dead jobs<commit_after>package work

import (
	// "fmt"
	"github.com/garyburd/redigo/redis"
	"sort"
	"strconv"
	"strings"
)

type Client struct {
	namespace string // eg, "myapp-work"
	pool      *redis.Pool
}

func NewClient(namespace string, pool *redis.Pool) *Client {
	return &Client{
		namespace: namespace,
		pool:      pool,
	}
}

type WorkerPoolHeartbeat struct {
	WorkerPoolID string
	StartedAt    int64
	HeartbeatAt  int64

	JobNames    []string
	Concurrency uint
	Host        string
	Pid         int

	WorkerIDs []string
}

func (c *Client) WorkerPoolHeartbeats() ([]*WorkerPoolHeartbeat, error) {
	conn := c.pool.Get()
	defer conn.Close()

	workerPoolsKey := redisKeyWorkerPools(c.namespace)

	workerPoolIDs, err := redis.Strings(conn.Do("SMEMBERS", workerPoolsKey))
	if err != nil {
		return nil, err
	}
	sort.Strings(workerPoolIDs)

	for _, wpid := range workerPoolIDs {
		key := redisKeyHeartbeat(c.namespace, wpid)
		conn.Send("HGETALL", key)
	}

	if err := conn.Flush(); err != nil {
		logError("worker_pool_statuses.flush", err)
		return nil, err
	}

	heartbeats := make([]*WorkerPoolHeartbeat, 0, len(workerPoolIDs))

	for _, wpid := range workerPoolIDs {
		vals, err := redis.Strings(conn.Receive())
		if err != nil {
			logError("worker_pool_statuses.receive", err)
			return nil, err
		}

		heartbeat := &WorkerPoolHeartbeat{
			WorkerPoolID: wpid,
		}

		for i := 0; i < len(vals)-1; i += 2 {
			key := vals[i]
			value := vals[i+1]

			var err error
			if key == "heartbeat_at" {
				heartbeat.HeartbeatAt, err = strconv.ParseInt(value, 10, 64)
			} else if key == "started_at" {
				heartbeat.StartedAt, err = strconv.ParseInt(value, 10, 64)
			} else if key == "job_names" {
				heartbeat.JobNames = strings.Split(value, ",")
				sort.Strings(heartbeat.JobNames)
			} else if key == "concurrency" {
				var vv uint64
				vv, err = strconv.ParseUint(value, 10, 0)
				heartbeat.Concurrency = uint(vv)
			} else if key == "host" {
				heartbeat.Host = value
			} else if key == "pid" {
				var vv int64
				vv, err = strconv.ParseInt(value, 10, 0)
				heartbeat.Pid = int(vv)
			} else if key == "worker_ids" {
				heartbeat.WorkerIDs = strings.Split(value, ",")
				sort.Strings(heartbeat.WorkerIDs)
			}
			if err != nil {
				logError("worker_pool_statuses.parse", err)
				return nil, err
			}
		}

		heartbeats = append(heartbeats, heartbeat)
	}

	return heartbeats, nil
}

type WorkerObservation struct {
	WorkerID string
	IsBusy   bool

	// If IsBusy:
	JobName   string
	JobID     string
	StartedAt int64
	ArgsJSON  string
	Checkin   string
	CheckinAt int64
}

func (c *Client) WorkerObservations() ([]*WorkerObservation, error) {
	conn := c.pool.Get()
	defer conn.Close()

	hbs, err := c.WorkerPoolHeartbeats()
	if err != nil {
		logError("worker_observations.worker_pool_heartbeats", err)
		return nil, err
	}

	var workerIDs []string
	for _, hb := range hbs {
		workerIDs = append(workerIDs, hb.WorkerIDs...)
	}

	for _, wid := range workerIDs {
		key := redisKeyWorkerStatus(c.namespace, wid) // TODO: rename this func
		conn.Send("HGETALL", key)
	}

	if err := conn.Flush(); err != nil {
		logError("worker_observations.flush", err)
		return nil, err
	}

	observations := make([]*WorkerObservation, 0, len(workerIDs))

	for _, wid := range workerIDs {
		vals, err := redis.Strings(conn.Receive())
		if err != nil {
			logError("worker_observations.receive", err)
			return nil, err
		}

		ob := &WorkerObservation{
			WorkerID: wid,
		}

		for i := 0; i < len(vals)-1; i += 2 {
			key := vals[i]
			value := vals[i+1]

			ob.IsBusy = true

			var err error
			if key == "job_name" {
				ob.JobName = value
			} else if key == "job_id" {
				ob.JobID = value
			} else if key == "started_at" {
				ob.StartedAt, err = strconv.ParseInt(value, 10, 64)
			} else if key == "args" {
				ob.ArgsJSON = value
			} else if key == "checkin" {
				ob.Checkin = value
			} else if key == "checkin_at" {
				ob.CheckinAt, err = strconv.ParseInt(value, 10, 64)
			}
			if err != nil {
				logError("worker_observations.parse", err)
				return nil, err
			}
		}

		observations = append(observations, ob)
	}

	return observations, nil
}

type Queue struct {
	JobName string
	Count   int64
	Latency int64
}

func (c *Client) Queues() ([]*Queue, error) {
	conn := c.pool.Get()
	defer conn.Close()

	key := redisKeyKnownJobs(c.namespace)
	jobNames, err := redis.Strings(conn.Do("SMEMBERS", key))
	if err != nil {
		return nil, err
	}
	sort.Strings(jobNames)

	for _, jobName := range jobNames {
		conn.Send("LLEN", redisKeyJobs(c.namespace, jobName))
	}

	if err := conn.Flush(); err != nil {
		logError("client.queues.flush", err)
		return nil, err
	}

	queues := make([]*Queue, 0, len(jobNames))

	for _, jobName := range jobNames {
		count, err := redis.Int64(conn.Receive())
		if err != nil {
			logError("client.queues.receive", err)
			return nil, err
		}

		queue := &Queue{
			JobName: jobName,
			Count:   count,
		}

		queues = append(queues, queue)
	}

	for _, s := range queues {
		if s.Count > 0 {
			conn.Send("LINDEX", redisKeyJobs(c.namespace, s.JobName), -1)
		}
	}

	if err := conn.Flush(); err != nil {
		logError("client.queues.flush2", err)
		return nil, err
	}

	now := nowEpochSeconds()

	for _, s := range queues {
		if s.Count > 0 {
			b, err := redis.Bytes(conn.Receive())
			if err != nil {
				logError("client.queues.receive2", err)
				return nil, err
			}

			job, err := newJob(b, nil, nil)
			if err != nil {
				logError("client.queues.new_job", err)
			}
			s.Latency = now - job.EnqueuedAt
		}
	}

	return queues, nil
}

type RetryJob struct {
	RetryAt int64
	*Job
}

type ScheduledJob struct {
	RunAt int64
	*Job
}

type DeadJob struct {
	DiedAt int64
	*Job
}

func (c *Client) ScheduledJobs(page uint) ([]*ScheduledJob, error) {
	key := redisKeyScheduled(c.namespace)
	jobsWithScores, err := c.getZsetPage(key, page)
	if err != nil {
		logError("client.scheduled_jobs.get_zset_page", err)
		return nil, err
	}

	jobs := make([]*ScheduledJob, 0, len(jobsWithScores))

	for _, jws := range jobsWithScores {
		jobs = append(jobs, &ScheduledJob{RunAt: jws.Score, Job: jws.job})
	}

	return jobs, nil
}

func (c *Client) RetryJobs(page uint) ([]*RetryJob, error) {
	key := redisKeyRetry(c.namespace)
	jobsWithScores, err := c.getZsetPage(key, page)
	if err != nil {
		logError("client.retry_jobs.get_zset_page", err)
		return nil, err
	}

	jobs := make([]*RetryJob, 0, len(jobsWithScores))

	for _, jws := range jobsWithScores {
		jobs = append(jobs, &RetryJob{RetryAt: jws.Score, Job: jws.job})
	}

	return jobs, nil
}

func (c *Client) DeadJobs(page uint) ([]*DeadJob, error) {
	key := redisKeyDead(c.namespace)
	jobsWithScores, err := c.getZsetPage(key, page)
	if err != nil {
		logError("client.dead_jobs.get_zset_page", err)
		return nil, err
	}

	jobs := make([]*DeadJob, 0, len(jobsWithScores))

	for _, jws := range jobsWithScores {
		jobs = append(jobs, &DeadJob{DiedAt: jws.Score, Job: jws.job})
	}

	return jobs, nil
}

type jobScore struct {
	JobBytes []byte
	Score    int64
	job      *Job
}

func (c *Client) getZsetPage(key string, page uint) ([]jobScore, error) {
	conn := c.pool.Get()
	defer conn.Close()

	if page == 0 {
		page = 1
	}

	values, err := redis.Values(conn.Do("ZRANGEBYSCORE", key, "-inf", "+inf", "WITHSCORES", "LIMIT", (page-1)*20, 20))
	if err != nil {
		logError("client.get_zset_page.values", err)
		return nil, err
	}

	var jobsWithScores []jobScore

	if err := redis.ScanSlice(values, &jobsWithScores); err != nil {
		logError("client.get_zset_page.scan_slice", err)
		return nil, err
	}

	for i, jws := range jobsWithScores {
		job, err := newJob(jws.JobBytes, nil, nil)
		if err != nil {
			logError("client.get_zset_page.new_job", err)
			return nil, err
		}

		jobsWithScores[i].job = job
	}

	return jobsWithScores, nil
}
<|endoftext|>


<commit_before>package mdns

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync/atomic"
	"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
	InfoFields []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   int32
	closedCh chan struct{} // TODO(reddaly): This doesn't appear to be used.
}

// 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 {
	if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		// something else already closed it
		return nil
	}

	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 system
// 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 append(resp.Answer, resp.Extra...) {
				// 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.InfoFields = 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 atomic.LoadInt32(&c.closed) == 0 {
		n, err := l.Read(buf)

		if atomic.LoadInt32(&c.closed) == 1 {
			return
		}

		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>Added error handling around WriteToUDP calls (#73)<commit_after>package mdns

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync/atomic"
	"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
	InfoFields []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   int32
	closedCh chan struct{} // TODO(reddaly): This doesn't appear to be used.
}

// 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 {
	if !atomic.CompareAndSwapInt32(&c.closed, 0, 1) {
		// something else already closed it
		return nil
	}

	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 system
// 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 append(resp.Answer, resp.Extra...) {
				// 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.InfoFields = 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 {
		_, err = c.ipv4UnicastConn.WriteToUDP(buf, ipv4Addr)
		if err != nil {
			return err
		}
	}
	if c.ipv6UnicastConn != nil {
		_, err = c.ipv6UnicastConn.WriteToUDP(buf, ipv6Addr)
		if err != nil {
			return err
		}
	}
	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 atomic.LoadInt32(&c.closed) == 0 {
		n, err := l.Read(buf)

		if atomic.LoadInt32(&c.closed) == 1 {
			return
		}

		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 weatherhist

import (
	"log"
	"net/http"
	"net/url"

	"io/ioutil"
	"path"
	"strconv"
	"time"

	"fmt"
	"strings"

	"github.com/pkg/errors"
)

type Client struct {
	URL        *url.URL
	HTTPClient *http.Client
	Logger     *log.Logger
}

const defaultURL = "http://www.data.jma.go.jp/obd/stats/etrn/view"

func NewClient(urlStrp *string, logger *log.Logger) (*Client, error) {
	urlStr := defaultURL
	if urlStrp != nil {
		urlStr = *urlStrp
	}
	url, err := url.Parse(urlStr)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse url")
	}

	l := log.New(ioutil.Discard, "", log.LstdFlags)
	if logger != nil {
		l = logger
	}

	c := http.DefaultClient
	return &Client{
		URL:        url,
		HTTPClient: c,
		Logger:     l,
	}, nil
}

func (c *Client) getFullURL(spath string, s Station, targetDate time.Time) string {
	q := url.Values{}
	q.Add("block_no", string(s.ID))
	q.Add("prec_no", s.GroupNumber)
	q.Add("year", strconv.Itoa(targetDate.Year()))
	q.Add("month", strconv.Itoa(int(targetDate.Month())))
	q.Add("day", strconv.Itoa(targetDate.Day()))
	q.Add("view", "")

	u := *c.URL
	u.RawQuery = q.Encode()

	switch spath {
	case DailyPath:
		spath = fmt.Sprintf(DailyPath, s.Type)
	}
	u.Path = path.Join(c.URL.Path, spath)
	return u.String()
}

func getFloatValueWithQuality(value string) FloatWithQuality {
	fwq := FloatWithQuality{}
	value = strings.TrimRight(value, " ")
	f, err := strconv.ParseFloat(value, 32)
	if err != nil {
		fwq.IsBadQuality = true
		// "2.5 ]" とかをパースやってみる
		v := strings.TrimRight(value, ")]")
		f, err = strconv.ParseFloat(v, 32)
		if err != nil {
			return FloatWithQuality{
				Value: nil,
				IsBadQuality: true,
			}
		}
	}
	f32 := float32(f)
	fwq.Value = &f32

	return fwq
}

const NilValue = "///"

func getStringValueWithQuality(value string) StringWithQuality {
	swq := StringWithQuality{}
	value = strings.TrimRight(value, " ")

	if strings.ContainsAny(value, ")]") {
		swq.IsBadQuality = true
	}

	value = strings.TrimRight(value, ")]")
	swq.Value = &value

	if value == NilValue || value == "×" || value == "#" {
		swq.Value = nil
		swq.IsBadQuality = true
	}

	return swq
}
<commit_msg>fix float string parser<commit_after>package weatherhist

import (
	"log"
	"net/http"
	"net/url"

	"io/ioutil"
	"path"
	"strconv"
	"time"

	"fmt"
	"strings"

	"github.com/pkg/errors"
)

type Client struct {
	URL        *url.URL
	HTTPClient *http.Client
	Logger     *log.Logger
}

const defaultURL = "http://www.data.jma.go.jp/obd/stats/etrn/view"

func NewClient(urlStrp *string, logger *log.Logger) (*Client, error) {
	urlStr := defaultURL
	if urlStrp != nil {
		urlStr = *urlStrp
	}
	url, err := url.Parse(urlStr)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse url")
	}

	l := log.New(ioutil.Discard, "", log.LstdFlags)
	if logger != nil {
		l = logger
	}

	c := http.DefaultClient
	return &Client{
		URL:        url,
		HTTPClient: c,
		Logger:     l,
	}, nil
}

func (c *Client) getFullURL(spath string, s Station, targetDate time.Time) string {
	q := url.Values{}
	q.Add("block_no", string(s.ID))
	q.Add("prec_no", s.GroupNumber)
	q.Add("year", strconv.Itoa(targetDate.Year()))
	q.Add("month", strconv.Itoa(int(targetDate.Month())))
	q.Add("day", strconv.Itoa(targetDate.Day()))
	q.Add("view", "")

	u := *c.URL
	u.RawQuery = q.Encode()

	switch spath {
	case DailyPath:
		spath = fmt.Sprintf(DailyPath, s.Type)
	}
	u.Path = path.Join(c.URL.Path, spath)
	return u.String()
}

func parseFloatFromString(value string) (*float32, bool) {
	f, err := strconv.ParseFloat(value, 32)
	if err == nil {
		f32 := float32(f)
		return &f32, true
	}
	value = strings.TrimRight(value, " )")
	f, err = strconv.ParseFloat(value, 32)
	if err == nil {
		f32 := float32(f)
		return &f32, true
	}

	// "2.5 ]" とかをパースやってみる
	v := strings.TrimRight(value, " ]")
	f, err = strconv.ParseFloat(v, 32)
	if err == nil {
		f32 := float32(f)
		return &f32, false
	}

	return nil, false
}

func getFloatValueWithQuality(value string) FloatWithQuality {
	fwq := FloatWithQuality{}
	value = strings.TrimRight(value, " ")
	f, ok := parseFloatFromString(value)
	if !ok {
		return FloatWithQuality{
			Value:        f,
			IsBadQuality: true,
		}
	}
	fwq.Value = f
	fwq.IsBadQuality = false

	return fwq
}

const NilValue = "///"

func getStringValueWithQuality(value string) StringWithQuality {
	swq := StringWithQuality{}
	value = strings.TrimRight(value, " ")

	if strings.ContainsAny(value, ")]") {
		swq.IsBadQuality = true
	}

	value = strings.TrimRight(value, ")]")
	swq.Value = &value

	if value == NilValue || value == "×" || value == "#" {
		swq.Value = nil
		swq.IsBadQuality = true
	}

	return swq
}
<|endoftext|>


<commit_before>package main

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

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/engine"
	"github.com/getsentry/raven-go"
)

const (
	Version = "0.0.8"
)

var (
	sentryDsn  = ""
	exitResult = false
)

func main() {
	showVersion := flag.Bool("version", false, "Prints changes-client version")
	flag.Parse()

	if *showVersion {
		fmt.Println(Version)
		return
	}

	if sentryDsn != "" {
		sentryClient, err := raven.NewClient(sentryDsn, map[string]string{
			"version": Version,
		})
		if err != nil {
			log.Fatal(err)
		}

		defer func() {
			var packet *raven.Packet
			p := recover()
			switch rval := p.(type) {
			case nil:
				return
			case error:
				packet = raven.NewPacket(rval.Error(), raven.NewException(rval, raven.NewStacktrace(2, 3, nil)))
			default:
				rvalStr := fmt.Sprint(rval)
				packet = raven.NewPacket(rvalStr, raven.NewException(errors.New(rvalStr), raven.NewStacktrace(2, 3, nil)))
			}

			log.Printf("[client] Sending panic to Sentry")
			_, ch := sentryClient.Capture(packet, map[string]string{})
			<-ch
			panic(p)
		}()

		run()
	} else {
		run()
	}
}

func run() {
	config, err := client.GetConfig()
	if err != nil {
		panic(err)
	}

	result, err := engine.RunBuildPlan(config)
	log.Printf("[client] Finished: %s", result)
	if (err != nil || result != engine.RESULT_PASSED) && exitResult {
		log.Printf("[client] exit: 1")
		os.Exit(1)
	}
	log.Printf("[client] exit: 0")
}

func init() {
	flag.StringVar(&sentryDsn, "sentry-dsn", "", "Sentry DSN for reporting errors")
	flag.BoolVar(&exitResult, "exit-result", false, "Determine exit code from result")
}
<commit_msg>Print out errors.<commit_after>package main

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

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/engine"
	"github.com/getsentry/raven-go"
)

const (
	Version = "0.0.8"
)

var (
	sentryDsn  = ""
	exitResult = false
)

func main() {
	showVersion := flag.Bool("version", false, "Prints changes-client version")
	flag.Parse()

	if *showVersion {
		fmt.Println(Version)
		return
	}

	if sentryDsn != "" {
		sentryClient, err := raven.NewClient(sentryDsn, map[string]string{
			"version": Version,
		})
		if err != nil {
			log.Fatal(err)
		}

		defer func() {
			var packet *raven.Packet
			p := recover()
			switch rval := p.(type) {
			case nil:
				return
			case error:
				packet = raven.NewPacket(rval.Error(), raven.NewException(rval, raven.NewStacktrace(2, 3, nil)))
			default:
				rvalStr := fmt.Sprint(rval)
				packet = raven.NewPacket(rvalStr, raven.NewException(errors.New(rvalStr), raven.NewStacktrace(2, 3, nil)))
			}

			log.Printf("[client] Sending panic to Sentry")
			_, ch := sentryClient.Capture(packet, map[string]string{})
			<-ch
			panic(p)
		}()

		run()
	} else {
		run()
	}
}

func run() {
	config, err := client.GetConfig()
	if err != nil {
		panic(err)
	}

	result, err := engine.RunBuildPlan(config)
	log.Printf("[client] Finished: %s", result)
	if err != nil {
		log.Printf("[client] error: %s", err.Error())
	}
	if (err != nil || result != engine.RESULT_PASSED) && exitResult {
		log.Printf("[client] exit: 1")
		os.Exit(1)
	}
	log.Printf("[client] exit: 0")
}

func init() {
	flag.StringVar(&sentryDsn, "sentry-dsn", "", "Sentry DSN for reporting errors")
	flag.BoolVar(&exitResult, "exit-result", false, "Determine exit code from result")
}
<|endoftext|>


<commit_before>// Package webfinger provides a simple client implementation of the WebFinger
// protocol.
//
// (This is a work in progress, the API is not frozen)
//
// This implementation tries to follow the last spec:
// http://tools.ietf.org/html/draft-ietf-appsawg-webfinger-05
//
// And also tries to provide backwark compatibility with the original spec:
// https://code.google.com/p/webfinger/wiki/WebFingerProtocol
//
// Example:
//
//  package main
//
//  import (
//          "fmt"
//          "github.com/ant0ine/go-webfinger"
//          "os"
//  )
//
//  func main() {
//          email := os.Args[1]
//
//          client := webfinger.NewClient(nil)
//
//          resource, err := webfinger.MakeResource(email)
//          if err != nil {
//                  panic(err)
//          }
//
//          jrd, err := client.GetJRD(resource)
//          if err != nil {
//                  fmt.Println(err)
//                  return
//          }
//
//          fmt.Printf("JRD: %+v", jrd)
//  }
package webfinger

import (
	"errors"
	"fmt"
	"github.com/ant0ine/go-webfinger/jrd"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

// Resource represents a WebFinger resource.
type Resource struct {
	Local  string
	Domain string
}

// MakeResource constructs a WebFinger resource for the provided email string.
func MakeResource(email string) (*Resource, error) {
	// TODO validate address, see http://www.ietf.org/rfc/rfc2822.txt
	// TODO accept an email address URI
	// TODO support mailto: http:  <= rework that
	parts := strings.SplitN(email, "@", 2)
	if len(parts) < 2 {
		return nil, errors.New("not a valid email")
	}
	return &Resource{
		Local:  parts[0],
		Domain: parts[1],
	}, nil
}

// AsURIString returns the resource as an URI string (eg: acct:user@domain).
func (self *Resource) AsURIString() string {
	return fmt.Sprintf("acct:%s@%s", self.Local, self.Domain)
}

// JRDURL returns the WebFinger URL that points to the JRD data for this resource.
func (self *Resource) JRDURL(rels []string) *url.URL {
	return &url.URL{
		Scheme: "https",
		Host:   self.Domain,
		Path:   "/.well-known/webfinger",
		RawQuery: url.Values{
			"resource": []string{self.AsURIString()},
			"rel":      rels,
		}.Encode(),
	}
}

// A Client is a WebFinger client.
type Client struct {
	// HTTP client used to perform WebFinger lookups.
	client *http.Client
}

// NewClient returns a new WebFinger client.  If a nil http.Client is provied,
// http.DefaultClient will be used.
func NewClient(httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return &Client{client: httpClient}
}

// GetJRDPart returns the JRD for the specified resource, with the ability to
// specify which "rel" links to include.
func (self *Client) GetJRDPart(resource *Resource, rels []string) (*jrd.JRD, error) {

	log.Printf("Trying to get WebFinger JRD data for: %s", resource.AsURIString())

	resourceJRD, err := self.fetchJRD(resource.JRDURL(rels))
	if err != nil {
		return nil, err
	}

	return resourceJRD, nil
}

// GetJRD returns the JRD data for this resource.
// It follows redirect, and retries with http if https is not available.
func (self *Client) GetJRD(resource *Resource) (*jrd.JRD, error) {
	return self.GetJRDPart(resource, nil)
}

func (self *Client) fetchJRD(jrdURL *url.URL) (*jrd.JRD, error) {
	// TODO verify signature if not https
	// TODO extract http cache info

	// Get follows up to 10 redirects
	log.Printf("GET %s", jrdURL.String())
	res, err := self.client.Get(jrdURL.String())
	if err != nil {
		// retry with http instead of https
		if strings.Contains(err.Error(), "connection refused") {
			jrdURL.Scheme = "http"
			log.Printf("GET %s", jrdURL.String())
			res, err = self.client.Get(jrdURL.String())
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	if !(200 <= res.StatusCode && res.StatusCode < 300) {
		return nil, errors.New(res.Status)
	}

	content, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return nil, err
	}

	ct := strings.ToLower(res.Header.Get("content-type"))
	if strings.Contains(ct, "application/jrd+json") ||
		strings.Contains(ct, "application/json") {
		parsed, err := jrd.ParseJRD(content)
		if err != nil {
			return nil, err
		}
		return parsed, nil
	}

	return nil, errors.New(fmt.Sprintf("invalid content-type: %s", ct))
}
<commit_msg>Update the gdoc description<commit_after>// Package webfinger provides a simple client implementation of the WebFinger
// protocol.
//
// It is a work in progress, the API is not frozen.
// We're trying to catchup with the last draft of the protocol:
// http://tools.ietf.org/html/draft-ietf-appsawg-webfinger-14
// and to support the http://webfist.org
//
// Example:
//
//  package main
//
//  import (
//          "fmt"
//          "github.com/ant0ine/go-webfinger"
//          "os"
//  )
//
//  func main() {
//          email := os.Args[1]
//
//          client := webfinger.NewClient(nil)
//
//          resource, err := webfinger.MakeResource(email)
//          if err != nil {
//                  panic(err)
//          }
//
//          jrd, err := client.GetJRD(resource)
//          if err != nil {
//                  fmt.Println(err)
//                  return
//          }
//
//          fmt.Printf("JRD: %+v", jrd)
//  }
package webfinger

import (
	"errors"
	"fmt"
	"github.com/ant0ine/go-webfinger/jrd"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"strings"
)

// Resource represents a WebFinger resource.
type Resource struct {
	Local  string
	Domain string
}

// MakeResource constructs a WebFinger resource for the provided email string.
func MakeResource(email string) (*Resource, error) {
	// TODO validate address, see http://www.ietf.org/rfc/rfc2822.txt
	// TODO accept an email address URI
	// TODO support mailto: http:  <= rework that
	parts := strings.SplitN(email, "@", 2)
	if len(parts) < 2 {
		return nil, errors.New("not a valid email")
	}
	return &Resource{
		Local:  parts[0],
		Domain: parts[1],
	}, nil
}

// AsURIString returns the resource as an URI string (eg: acct:user@domain).
func (self *Resource) AsURIString() string {
	return fmt.Sprintf("acct:%s@%s", self.Local, self.Domain)
}

// JRDURL returns the WebFinger URL that points to the JRD data for this resource.
func (self *Resource) JRDURL(rels []string) *url.URL {
	return &url.URL{
		Scheme: "https",
		Host:   self.Domain,
		Path:   "/.well-known/webfinger",
		RawQuery: url.Values{
			"resource": []string{self.AsURIString()},
			"rel":      rels,
		}.Encode(),
	}
}

// A Client is a WebFinger client.
type Client struct {
	// HTTP client used to perform WebFinger lookups.
	client *http.Client
}

// NewClient returns a new WebFinger client.  If a nil http.Client is provied,
// http.DefaultClient will be used.
func NewClient(httpClient *http.Client) *Client {
	if httpClient == nil {
		httpClient = http.DefaultClient
	}
	return &Client{client: httpClient}
}

// GetJRDPart returns the JRD for the specified resource, with the ability to
// specify which "rel" links to include.
func (self *Client) GetJRDPart(resource *Resource, rels []string) (*jrd.JRD, error) {

	log.Printf("Trying to get WebFinger JRD data for: %s", resource.AsURIString())

	resourceJRD, err := self.fetchJRD(resource.JRDURL(rels))
	if err != nil {
		return nil, err
	}

	return resourceJRD, nil
}

// GetJRD returns the JRD data for this resource.
// It follows redirect, and retries with http if https is not available.
func (self *Client) GetJRD(resource *Resource) (*jrd.JRD, error) {
	return self.GetJRDPart(resource, nil)
}

func (self *Client) fetchJRD(jrdURL *url.URL) (*jrd.JRD, error) {
	// TODO verify signature if not https
	// TODO extract http cache info

	// Get follows up to 10 redirects
	log.Printf("GET %s", jrdURL.String())
	res, err := self.client.Get(jrdURL.String())
	if err != nil {
		// retry with http instead of https
		if strings.Contains(err.Error(), "connection refused") {
			jrdURL.Scheme = "http"
			log.Printf("GET %s", jrdURL.String())
			res, err = self.client.Get(jrdURL.String())
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}

	if !(200 <= res.StatusCode && res.StatusCode < 300) {
		return nil, errors.New(res.Status)
	}

	content, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		return nil, err
	}

	ct := strings.ToLower(res.Header.Get("content-type"))
	if strings.Contains(ct, "application/jrd+json") ||
		strings.Contains(ct, "application/json") {
		parsed, err := jrd.ParseJRD(content)
		if err != nil {
			return nil, err
		}
		return parsed, nil
	}

	return nil, errors.New(fmt.Sprintf("invalid content-type: %s", ct))
}
<|endoftext|>


<commit_before>package discoverd

import (
	"errors"
	"net"
	"os"
	"sort"
	"sync"
	"time"

	"github.com/flynn/discoverd/agent"
	"github.com/flynn/rpcplus"
)

type Service struct {
	Created uint
	Name    string
	Host    string
	Port    string
	Addr    string
	Attrs   map[string]string
}

type ServiceSet struct {
	sync.Mutex
	services map[string]*Service
	filters  map[string]string
	watches  map[chan *agent.ServiceUpdate]bool
	leaders  chan *Service
	call     *rpcplus.Call
	self     *Service
	SelfAddr string
}

func copyService(service *Service) *Service {
	s := *service
	s.Attrs = make(map[string]string, len(service.Attrs))
	for k, v := range service.Attrs {
		s.Attrs[k] = v
	}
	return &s
}

func makeServiceSet(call *rpcplus.Call) *ServiceSet {
	return &ServiceSet{
		services: make(map[string]*Service),
		filters:  make(map[string]string),
		watches:  make(map[chan *agent.ServiceUpdate]bool),
		call:     call,
	}
}

func (s *ServiceSet) bind(updates chan *agent.ServiceUpdate) chan struct{} {
	// current is an event when enough service updates have been
	// received to bring us to "current" state (when subscribed)
	current := make(chan struct{})
	go func() {
		isCurrent := false
		for update := range updates {
			if update.Addr == "" && update.Name == "" && !isCurrent {
				close(current)
				isCurrent = true
				continue
			}
			s.Lock()
			if s.filters != nil && !s.matchFilters(update.Attrs) {
				s.Unlock()
				continue
			}
			if s.SelfAddr != update.Addr && update.Online {
				if _, exists := s.services[update.Addr]; !exists {
					host, port, _ := net.SplitHostPort(update.Addr)
					s.services[update.Addr] = &Service{
						Name:    update.Name,
						Addr:    update.Addr,
						Host:    host,
						Port:    port,
						Created: update.Created,
					}
				}
				s.services[update.Addr].Attrs = update.Attrs
			} else {
				if _, exists := s.services[update.Addr]; exists {
					delete(s.services, update.Addr)
				} else {
					s.Unlock()
					if s.SelfAddr == update.Addr {
						s.updateWatches(update)
					}
					continue
				}
			}
			s.Unlock()
			s.updateWatches(update)
		}
		s.closeWatches()
	}()
	return current
}

func (s *ServiceSet) updateWatches(update *agent.ServiceUpdate) {
	s.Lock()
	defer s.Unlock()
	for ch, once := range s.watches {
		ch <- update
		if once {
			delete(s.watches, ch)
		}
	}
}

func (s *ServiceSet) closeWatches() {
	s.Lock()
	defer s.Unlock()
	for ch := range s.watches {
		close(ch)
	}
}

func (s *ServiceSet) matchFilters(attrs map[string]string) bool {
	for key, value := range s.filters {
		if attrs[key] != value {
			return false
		}
	}
	return true
}

func (s *ServiceSet) Leader() *Service {
	services := s.Services()
	if len(services) > 0 {
		if s.self != nil && services[0].Created > s.self.Created {
			return s.self
		}
		return services[0]
	}
	if s.self != nil {
		return s.self
	}
	return nil
}

func (s *ServiceSet) Leaders() chan *Service {
	if s.leaders != nil {
		return s.leaders
	}
	s.leaders = make(chan *Service)
	updates := make(chan *agent.ServiceUpdate)
	s.Watch(updates, false, false)
	go func() {
		leader := s.Leader()
		s.leaders <- leader
		for update := range updates {
			if !update.Online && update.Addr == leader.Addr {
				leader = s.Leader()
				s.leaders <- leader
			}
		}
	}()
	return s.leaders
}

type serviceByAge []*Service

func (a serviceByAge) Len() int           { return len(a) }
func (a serviceByAge) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a serviceByAge) Less(i, j int) bool { return a[i].Created < a[j].Created }

func (s *ServiceSet) Services() []*Service {
	s.Lock()
	defer s.Unlock()
	list := make([]*Service, 0, len(s.services))
	for _, service := range s.services {
		list = append(list, copyService(service))
	}
	if len(list) > 0 {
		sort.Sort(serviceByAge(list))
	}
	return list
}

func (s *ServiceSet) Addrs() []string {
	list := make([]string, 0, len(s.services))
	for _, service := range s.Services() {
		list = append(list, service.Addr)
	}
	return list
}

func (s *ServiceSet) Select(attrs map[string]string) []*Service {
	s.Lock()
	defer s.Unlock()
	list := make([]*Service, 0, len(s.services))
outer:
	for _, service := range s.services {
		for key, value := range attrs {
			if service.Attrs[key] != value {
				continue outer
			}
		}
		list = append(list, service)
	}
	return list
}

func (s *ServiceSet) Filter(attrs map[string]string) {
	s.Lock()
	defer s.Unlock()
	s.filters = attrs
	for key, service := range s.services {
		if !s.matchFilters(service.Attrs) {
			delete(s.services, key)
		}
	}
}

func (s *ServiceSet) Watch(ch chan *agent.ServiceUpdate, bringCurrent bool, fireOnce bool) {
	s.Lock()
	s.watches[ch] = fireOnce
	s.Unlock()
	if bringCurrent {
		go func() {
			s.Lock()
			defer s.Unlock()
			for _, service := range s.services {
				ch <- &agent.ServiceUpdate{
					Name:    service.Name,
					Addr:    service.Addr,
					Online:  true,
					Attrs:   service.Attrs,
					Created: service.Created,
				}
			}
		}()
	}
}

func (s *ServiceSet) Unwatch(ch chan *agent.ServiceUpdate) {
	s.Lock()
	defer s.Unlock()
	delete(s.watches, ch)
}

func (s *ServiceSet) Wait() chan *agent.ServiceUpdate {
	updateCh := make(chan *agent.ServiceUpdate, 1024) // buffer because of Watch bringCurrent race bug
	s.Watch(updateCh, true, true)
	return updateCh
}

func (s *ServiceSet) Close() error {
	return s.call.CloseStream()
}

type Client struct {
	sync.Mutex
	client        *rpcplus.Client
	heartbeats    map[string]chan struct{}
	expandedAddrs map[string]string
}

func NewClient() (*Client, error) {
	addr := os.Getenv("DISCOVERD")
	if addr == "" {
		addr = "127.0.0.1:1111"
	}
	return NewClientUsingAddress(addr)
}

func NewClientUsingAddress(addr string) (*Client, error) {
	client, err := rpcplus.DialHTTP("tcp", addr)
	return &Client{
		client:        client,
		heartbeats:    make(map[string]chan struct{}),
		expandedAddrs: make(map[string]string),
	}, err
}

func (c *Client) ServiceSet(name string) (*ServiceSet, error) {
	updates := make(chan *agent.ServiceUpdate)
	call := c.client.StreamGo("Agent.Subscribe", &agent.Args{
		Name: name,
	}, updates)
	set := makeServiceSet(call)
	<-set.bind(updates)
	return set, nil
}

func (c *Client) Services(name string, timeout time.Duration) ([]*Service, error) {
	set, err := c.ServiceSet(name)
	if err != nil {
		return nil, err
	}
	defer set.Close()
	select {
	case <-set.Wait():
		return set.Services(), nil
	case <-time.After(time.Duration(timeout) * time.Second):
		return nil, errors.New("discover: wait timeout exceeded")
	}
}

func (c *Client) Register(name, addr string) error {
	return c.RegisterWithAttributes(name, addr, nil)
}

func (c *Client) RegisterWithSet(name, addr string, attributes map[string]string) (*ServiceSet, error) {
	err := c.RegisterWithAttributes(name, addr, attributes)
	if err != nil {
		return nil, err
	}
	set, err := c.ServiceSet(name)
	if err != nil {
		c.Unregister(name, addr)
		return nil, err
	}
	set.SelfAddr = c.expandedAddrs[addr]
	_, exists := set.services[set.SelfAddr]
	if !exists {
		update := <-set.Wait()
		for update.Addr != set.SelfAddr {
			update = <-set.Wait()
		}
	}
	set.Lock()
	set.self = set.services[set.SelfAddr]
	delete(set.services, set.SelfAddr)
	set.Unlock()
	return set, nil
}

func (c *Client) RegisterAndStandby(name, addr string, attributes map[string]string) (chan *Service, error) {
	set, err := c.RegisterWithSet(name, addr, attributes)
	if err != nil {
		return nil, err
	}
	standbyCh := make(chan *Service)
	go func() {
		for leader := range set.Leaders() {
			if leader.Addr == set.SelfAddr {
				set.Close()
				standbyCh <- leader
				return
			}
		}
	}()
	return standbyCh, nil
}

func (c *Client) RegisterWithAttributes(name, addr string, attributes map[string]string) error {
	args := &agent.Args{
		Name:  name,
		Addr:  addr,
		Attrs: attributes,
	}
	var ret string
	err := c.client.Call("Agent.Register", args, &ret)
	if err != nil {
		return errors.New("discover: register failed: " + err.Error())
	}
	done := make(chan struct{})
	c.Lock()
	c.heartbeats[args.Addr] = done
	c.expandedAddrs[args.Addr] = ret
	c.Unlock()
	go func() {
		ticker := time.NewTicker(agent.HeartbeatIntervalSecs * time.Second) // TODO: add jitter
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				// TODO: log error here
				c.client.Call("Agent.Heartbeat", &agent.Args{
					Name: name,
					Addr: args.Addr,
				}, &struct{}{})
			case <-done:
				return
			}
		}
	}()
	return nil
}

func (c *Client) Unregister(name, addr string) error {
	args := &agent.Args{
		Name: name,
		Addr: addr,
	}
	c.Lock()
	close(c.heartbeats[args.Addr])
	delete(c.heartbeats, args.Addr)
	c.Unlock()
	err := c.client.Call("Agent.Unregister", args, &struct{}{})
	if err != nil {
		return errors.New("discover: unregister failed: " + err.Error())
	}
	return nil
}
<commit_msg>discoverd/client: put mutexes behind unexported member<commit_after>package discoverd

import (
	"errors"
	"net"
	"os"
	"sort"
	"sync"
	"time"

	"github.com/flynn/discoverd/agent"
	"github.com/flynn/rpcplus"
)

type Service struct {
	Created uint
	Name    string
	Host    string
	Port    string
	Addr    string
	Attrs   map[string]string
}

type ServiceSet struct {
	l        sync.Mutex
	services map[string]*Service
	filters  map[string]string
	watches  map[chan *agent.ServiceUpdate]bool
	leaders  chan *Service
	call     *rpcplus.Call
	self     *Service
	SelfAddr string
}

func copyService(service *Service) *Service {
	s := *service
	s.Attrs = make(map[string]string, len(service.Attrs))
	for k, v := range service.Attrs {
		s.Attrs[k] = v
	}
	return &s
}

func makeServiceSet(call *rpcplus.Call) *ServiceSet {
	return &ServiceSet{
		services: make(map[string]*Service),
		filters:  make(map[string]string),
		watches:  make(map[chan *agent.ServiceUpdate]bool),
		call:     call,
	}
}

func (s *ServiceSet) bind(updates chan *agent.ServiceUpdate) chan struct{} {
	// current is an event when enough service updates have been
	// received to bring us to "current" state (when subscribed)
	current := make(chan struct{})
	go func() {
		isCurrent := false
		for update := range updates {
			if update.Addr == "" && update.Name == "" && !isCurrent {
				close(current)
				isCurrent = true
				continue
			}
			s.l.Lock()
			if s.filters != nil && !s.matchFilters(update.Attrs) {
				s.l.Unlock()
				continue
			}
			if s.SelfAddr != update.Addr && update.Online {
				if _, exists := s.services[update.Addr]; !exists {
					host, port, _ := net.SplitHostPort(update.Addr)
					s.services[update.Addr] = &Service{
						Name:    update.Name,
						Addr:    update.Addr,
						Host:    host,
						Port:    port,
						Created: update.Created,
					}
				}
				s.services[update.Addr].Attrs = update.Attrs
			} else {
				if _, exists := s.services[update.Addr]; exists {
					delete(s.services, update.Addr)
				} else {
					s.l.Unlock()
					if s.SelfAddr == update.Addr {
						s.updateWatches(update)
					}
					continue
				}
			}
			s.l.Unlock()
			s.updateWatches(update)
		}
		s.closeWatches()
	}()
	return current
}

func (s *ServiceSet) updateWatches(update *agent.ServiceUpdate) {
	s.l.Lock()
	defer s.l.Unlock()
	for ch, once := range s.watches {
		ch <- update
		if once {
			delete(s.watches, ch)
		}
	}
}

func (s *ServiceSet) closeWatches() {
	s.l.Lock()
	defer s.l.Unlock()
	for ch := range s.watches {
		close(ch)
	}
}

func (s *ServiceSet) matchFilters(attrs map[string]string) bool {
	for key, value := range s.filters {
		if attrs[key] != value {
			return false
		}
	}
	return true
}

func (s *ServiceSet) Leader() *Service {
	services := s.Services()
	if len(services) > 0 {
		if s.self != nil && services[0].Created > s.self.Created {
			return s.self
		}
		return services[0]
	}
	if s.self != nil {
		return s.self
	}
	return nil
}

func (s *ServiceSet) Leaders() chan *Service {
	if s.leaders != nil {
		return s.leaders
	}
	s.leaders = make(chan *Service)
	updates := make(chan *agent.ServiceUpdate)
	s.Watch(updates, false, false)
	go func() {
		leader := s.Leader()
		s.leaders <- leader
		for update := range updates {
			if !update.Online && update.Addr == leader.Addr {
				leader = s.Leader()
				s.leaders <- leader
			}
		}
	}()
	return s.leaders
}

type serviceByAge []*Service

func (a serviceByAge) Len() int           { return len(a) }
func (a serviceByAge) Swap(i, j int)      { a[i], a[j] = a[j], a[i] }
func (a serviceByAge) Less(i, j int) bool { return a[i].Created < a[j].Created }

func (s *ServiceSet) Services() []*Service {
	s.l.Lock()
	defer s.l.Unlock()
	list := make([]*Service, 0, len(s.services))
	for _, service := range s.services {
		list = append(list, copyService(service))
	}
	if len(list) > 0 {
		sort.Sort(serviceByAge(list))
	}
	return list
}

func (s *ServiceSet) Addrs() []string {
	list := make([]string, 0, len(s.services))
	for _, service := range s.Services() {
		list = append(list, service.Addr)
	}
	return list
}

func (s *ServiceSet) Select(attrs map[string]string) []*Service {
	s.l.Lock()
	defer s.l.Unlock()
	list := make([]*Service, 0, len(s.services))
outer:
	for _, service := range s.services {
		for key, value := range attrs {
			if service.Attrs[key] != value {
				continue outer
			}
		}
		list = append(list, service)
	}
	return list
}

func (s *ServiceSet) Filter(attrs map[string]string) {
	s.l.Lock()
	defer s.l.Unlock()
	s.filters = attrs
	for key, service := range s.services {
		if !s.matchFilters(service.Attrs) {
			delete(s.services, key)
		}
	}
}

func (s *ServiceSet) Watch(ch chan *agent.ServiceUpdate, bringCurrent bool, fireOnce bool) {
	s.l.Lock()
	s.watches[ch] = fireOnce
	s.l.Unlock()
	if bringCurrent {
		go func() {
			s.l.Lock()
			defer s.l.Unlock()
			for _, service := range s.services {
				ch <- &agent.ServiceUpdate{
					Name:    service.Name,
					Addr:    service.Addr,
					Online:  true,
					Attrs:   service.Attrs,
					Created: service.Created,
				}
			}
		}()
	}
}

func (s *ServiceSet) Unwatch(ch chan *agent.ServiceUpdate) {
	s.l.Lock()
	defer s.l.Unlock()
	delete(s.watches, ch)
}

func (s *ServiceSet) Wait() chan *agent.ServiceUpdate {
	updateCh := make(chan *agent.ServiceUpdate, 1024) // buffer because of Watch bringCurrent race bug
	s.Watch(updateCh, true, true)
	return updateCh
}

func (s *ServiceSet) Close() error {
	return s.call.CloseStream()
}

type Client struct {
	l             sync.Mutex
	client        *rpcplus.Client
	heartbeats    map[string]chan struct{}
	expandedAddrs map[string]string
}

func NewClient() (*Client, error) {
	addr := os.Getenv("DISCOVERD")
	if addr == "" {
		addr = "127.0.0.1:1111"
	}
	return NewClientUsingAddress(addr)
}

func NewClientUsingAddress(addr string) (*Client, error) {
	client, err := rpcplus.DialHTTP("tcp", addr)
	return &Client{
		client:        client,
		heartbeats:    make(map[string]chan struct{}),
		expandedAddrs: make(map[string]string),
	}, err
}

func (c *Client) ServiceSet(name string) (*ServiceSet, error) {
	updates := make(chan *agent.ServiceUpdate)
	call := c.client.StreamGo("Agent.Subscribe", &agent.Args{
		Name: name,
	}, updates)
	set := makeServiceSet(call)
	<-set.bind(updates)
	return set, nil
}

func (c *Client) Services(name string, timeout time.Duration) ([]*Service, error) {
	set, err := c.ServiceSet(name)
	if err != nil {
		return nil, err
	}
	defer set.Close()
	select {
	case <-set.Wait():
		return set.Services(), nil
	case <-time.After(time.Duration(timeout) * time.Second):
		return nil, errors.New("discover: wait timeout exceeded")
	}
}

func (c *Client) Register(name, addr string) error {
	return c.RegisterWithAttributes(name, addr, nil)
}

func (c *Client) RegisterWithSet(name, addr string, attributes map[string]string) (*ServiceSet, error) {
	err := c.RegisterWithAttributes(name, addr, attributes)
	if err != nil {
		return nil, err
	}
	set, err := c.ServiceSet(name)
	if err != nil {
		c.Unregister(name, addr)
		return nil, err
	}
	set.SelfAddr = c.expandedAddrs[addr]
	_, exists := set.services[set.SelfAddr]
	if !exists {
		update := <-set.Wait()
		for update.Addr != set.SelfAddr {
			update = <-set.Wait()
		}
	}
	set.l.Lock()
	set.self = set.services[set.SelfAddr]
	delete(set.services, set.SelfAddr)
	set.l.Unlock()
	return set, nil
}

func (c *Client) RegisterAndStandby(name, addr string, attributes map[string]string) (chan *Service, error) {
	set, err := c.RegisterWithSet(name, addr, attributes)
	if err != nil {
		return nil, err
	}
	standbyCh := make(chan *Service)
	go func() {
		for leader := range set.Leaders() {
			if leader.Addr == set.SelfAddr {
				set.Close()
				standbyCh <- leader
				return
			}
		}
	}()
	return standbyCh, nil
}

func (c *Client) RegisterWithAttributes(name, addr string, attributes map[string]string) error {
	args := &agent.Args{
		Name:  name,
		Addr:  addr,
		Attrs: attributes,
	}
	var ret string
	err := c.client.Call("Agent.Register", args, &ret)
	if err != nil {
		return errors.New("discover: register failed: " + err.Error())
	}
	done := make(chan struct{})
	c.l.Lock()
	c.heartbeats[args.Addr] = done
	c.expandedAddrs[args.Addr] = ret
	c.l.Unlock()
	go func() {
		ticker := time.NewTicker(agent.HeartbeatIntervalSecs * time.Second) // TODO: add jitter
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				// TODO: log error here
				c.client.Call("Agent.Heartbeat", &agent.Args{
					Name: name,
					Addr: args.Addr,
				}, &struct{}{})
			case <-done:
				return
			}
		}
	}()
	return nil
}

func (c *Client) Unregister(name, addr string) error {
	args := &agent.Args{
		Name: name,
		Addr: addr,
	}
	c.l.Lock()
	close(c.heartbeats[args.Addr])
	delete(c.heartbeats, args.Addr)
	c.l.Unlock()
	err := c.client.Call("Agent.Unregister", args, &struct{}{})
	if err != nil {
		return errors.New("discover: unregister failed: " + err.Error())
	}
	return nil
}
<|endoftext|>


<commit_before>package cmd

import (
	"fmt"

	"github.com/phil-mansfield/shellfish/parse"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/cmd/catalog"

	"github.com/phil-mansfield/shellfish/cmd/halo"
	"github.com/phil-mansfield/shellfish/io"
)

const finderCells = 150

// IDConfig contains the configuration fileds for the 'id' mode of the shellfish
// tool.
type IDConfig struct {
	idType string
	ids []int64
	idStart, idEnd, snap, mult int64

	exclusionStrategy string
	exclusionRadiusMult float64
}

var _ Mode = &IDConfig{}

// ExampleConfig creates an example id.config file.
func (config *IDConfig) ExampleConfig() string {
	return `[id.config]
#####################
## Required Fields ##
#####################

# Index of the snapshot to be analyzed.
Snap = 100

IDs = 10, 11, 12, 13, 14

#####################
## Optional Fields ##
#####################

# IDType indicates what the input IDs correspond to. It can be set to the
# following modes:
# halo-id - The numeric IDs given in the halo catalog.
# m200m   - The rank of the halos when sorted by M200m.
#
# Defaults to m200m if not set.
# IDType = m200m

# An alternative way of specifying IDs is to select start and end (inclusive)
# ID values. If the IDs variable is not set, both of these values must be set.
#
# IDStart = 10
# IDEnd = 15

# ExclusionStrategy determines how to exclude IDs from the given set. This is
# useful because splashback shells are not particularly meaningful for
# subhalos. It can be set to the following modes:
# none    - No halos are removed
# subhalo - Halos flagged as subhalos in the catalog are removed
# overlap - Halos which have an R200m shell that overlaps with a larger halo's
#           R200m shell are removed
#
# ExclusionStrategy defaults to overlap if not set.
#
# ExclusionStrategy = overlap

# ExclusionRadiusMult is a multiplier of R200m applied for the sake of
# determining exclusions.
#
# ExclusionRadiusMult defaults to 1 if not set.
#
# ExclustionRadiusMult = 1

# Mult is the number of times a given ID should be repeated. This is most useful
# if you want to estimate the scatter in shell measurements for halos with a
# given set of shell parameters.
#
# Mult defaults to 1 if not set.
#
# Mult = 1`
}

// ReadConfig reads in an id.config file into config.
func (config *IDConfig) ReadConfig(fname string) error {

	vars := parse.NewConfigVars("id.config")
	vars.String(&config.idType, "IDType", "m200m")
	vars.Ints(&config.ids, "IDs", []int64{})
	vars.Int(&config.idStart, "IDStart", -1)
	vars.Int(&config.idEnd, "IDEnd", -1)
	vars.Int(&config.mult, "Mult", 1)
	vars.Int(&config.snap, "Snap", -1)
	vars.String(&config.exclusionStrategy, "ExclusionStrategy", "subhalo")
	vars.Float(&config.exclusionRadiusMult, "ExclusionRadiusMult", 1)

	if fname == "" { return nil }
	if err := parse.ReadConfig(fname, vars); err != nil { return err }
	return config.validate()
}

// validate checks whether all the fields of config are valid.
func (config *IDConfig) validate() error {
	switch config.idType {
	case "halo-id", "m200m":
	default:
		return fmt.Errorf("The 'IDType' variable is set to '%s', which I " +
			"don't recognize.", config.idType)
	}

	switch config.exclusionStrategy {
	case "none", "subhalo":
	case "overlap":
		if config.exclusionRadiusMult <= 0 {
			return fmt.Errorf("The 'ExclusionRadiusMult' varaible is set to " +
				"%g, but it needs to be positive.", config.exclusionRadiusMult)
		}
	default:
		return fmt.Errorf("The 'ExclusionStrategy' variable is set to '%s', " +
		"which I don't recognize.", config.exclusionStrategy)
	}

	// TODO: Check the ranges of the IDs as well as IDStart and IDEnd
	if len(config.ids) == 0 {
		switch {
		case config.idStart == -1 && config.idEnd == -1:
			return fmt.Errorf("'IDs' variable not set.")
		case config.idStart == -1:
			return fmt.Errorf("'IDStart variable not set.")
		case config.idEnd == -1:
			return fmt.Errorf("'IDEnd' variable not set.")
		case config.idEnd < config.idStart:
			return fmt.Errorf("'IDEnd' variable set to %d, but 'IDStart' " +
				"variable set to %d.", config.idEnd, config.idStart)
		}
	}

	switch {
	case config.snap == -1:
		return fmt.Errorf("'Snap' variable not set.")
	case config.snap < 0:
		return fmt.Errorf("'Snap' variable set to %d.", config.snap)
	}

	if config.mult <= 0 {
		return fmt.Errorf("'Mult' variable set to %d", config.mult)
	}

	return nil
}

// Run executes the ID mode of shellfish tool.
func (config *IDConfig) Run(
	flags []string, gConfig *GlobalConfig, e *env.Environment, stdin []string,
) ([]string, error) {

	if config.snap < gConfig.SnapMin || config.snap > gConfig.SnapMax {
		return nil, fmt.Errorf("'Snap' = %d, but 'SnapMin' = %d and " +
			"'SnapMax = %d'", config.snap, gConfig.SnapMin, gConfig.SnapMax)
	}

	// Get IDs and snapshots

	rawIds := getIDs(config.idStart, config.idEnd, config.ids)

	vars := &halo.VarColumns{
		ID: int(gConfig.HaloIDColumn),
		X: int(gConfig.HaloPositionColumns[0]),
		Y: int(gConfig.HaloPositionColumns[1]),
		Z: int(gConfig.HaloPositionColumns[2]),
		M200m: int(gConfig.HaloM200mColumn),
	}

	var ids, snaps []int
	switch config.idType {
	case "halo-id":
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = int(config.snap) }
		ids = rawIds
	case "m200m":
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = int(config.snap) }

		var err error
		ids, err = convertSortedIDs(rawIds, int(config.snap), vars, e)
		if err != nil { return nil, err }
	default:
		panic("Impossible")
	}

	// Tag subhalos, if neccessary.
	exclude := make([]bool, len(ids))
	switch config.exclusionStrategy {
	case "none":
	case "subhalo":
		panic("subhalo is not implemented")
	case "overlap":
		var err error
		exclude, err = findOverlapSubs(ids, snaps, vars, e, config)
		if err != nil { return nil, err }
	}

	// Generate lines
	intCols := [][]int{ids, snaps}
	floatCols := [][]float64{}
	colOrder := []int{0, 1}
	lines := catalog.FormatCols(intCols, floatCols, colOrder)

	// Filter
	fLines := []string{}
	for i := range lines {
		if !exclude[i] { fLines = append(fLines, lines[i]) }
	}

	// Multiply
	mLines := []string{}
	for i := range fLines {
		for j := 0; j < int(config.mult); j++ {
			mLines = append(mLines, fLines[i])
		}
	}

	cString := catalog.CommentString(
		[]string{"ID", "Snapshot"}, []string{}, []int{0, 1},
	)
	mLines = append([]string{cString}, mLines...)

	return mLines, nil
}

func getIDs(idStart, idEnd int64, ids []int64) []int {
	if idStart != -1 {
		out := make([]int, idEnd - idStart)
		for i := range out {
			out[i] = int(idStart) + i
		}
		return out
	} else {
		out := make([]int, len(ids))
		for i := range out {
			out[i] = int(ids[i])
		}
		return out
	}
}

func convertSortedIDs(
	rawIDs []int, snap int, vars *halo.VarColumns, e *env.Environment,
) ([]int, error) {
	maxID := 0
	for _, id := range rawIDs {
		if id > maxID { maxID = id }
	}

	rids, err := memo.ReadSortedRockstarIDs(snap, maxID, vars, e)
	if err != nil { return nil, err }

	ids := make([]int, len(rawIDs))
	for i := range ids { ids[i] = rids[rawIDs[i]] }
	return ids, nil
}

func findOverlapSubs(
	rawIDs, snaps []int, vars *halo.VarColumns,
	e *env.Environment, config *IDConfig,
) ([]bool, error) {
	isSub := make([]bool, len(rawIDs))

	// Group by snapshot.
	snapGroups := make(map[int][]int)
	groupIdxs := make(map[int][]int)
	for i, id := range rawIDs {
		snap := snaps[i]
		snapGroups[snap] = append(snapGroups[snap], id)
		groupIdxs[snap] = append(groupIdxs[snap], i)
	}

	// Load each snapshot.
	hd := &io.Header{}
	buf, err := io.NewGotetraBuffer(e.ParticleCatalog(snaps[0], 0))
	if err != nil { return nil, err }

	for snap, group := range snapGroups {
		err := buf.ReadHeader(e.ParticleCatalog(snap, 0), hd)
		if err != nil { return nil, err }

		rids, err := memo.ReadSortedRockstarIDs(snap, -1, vars, e)
		if err != nil { return nil, err }
		_, xs, ys, zs, _, rs, err := memo.ReadRockstar(snap, rids, vars, e)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, config.exclusionRadiusMult)

		for i, id := range group {
			origIdx := groupIdxs[snap][i]
			// TODO: Holy linear search, batman! Fix this.
			for j, checkID := range rids {
				if checkID == id {
					isSub[origIdx] = sf.HostCount(j) > 0
					break
				} else if j == len(rids) - 1 {
					return nil, fmt.Errorf("ID %d not in halo list.", id)
				}
			}
		}
	}
	return isSub, nil
}
<commit_msg>Removed unneccessary header read from shellfish id<commit_after>package cmd

import (
	"fmt"

	"github.com/phil-mansfield/shellfish/parse"
	"github.com/phil-mansfield/shellfish/cmd/env"
	"github.com/phil-mansfield/shellfish/cmd/memo"
	"github.com/phil-mansfield/shellfish/cmd/catalog"

	"github.com/phil-mansfield/shellfish/cmd/halo"
)

const finderCells = 150

// IDConfig contains the configuration fileds for the 'id' mode of the shellfish
// tool.
type IDConfig struct {
	idType string
	ids []int64
	idStart, idEnd, snap, mult int64

	exclusionStrategy string
	exclusionRadiusMult float64
}

var _ Mode = &IDConfig{}

// ExampleConfig creates an example id.config file.
func (config *IDConfig) ExampleConfig() string {
	return `[id.config]
#####################
## Required Fields ##
#####################

# Index of the snapshot to be analyzed.
Snap = 100

IDs = 10, 11, 12, 13, 14

#####################
## Optional Fields ##
#####################

# IDType indicates what the input IDs correspond to. It can be set to the
# following modes:
# halo-id - The numeric IDs given in the halo catalog.
# m200m   - The rank of the halos when sorted by M200m.
#
# Defaults to m200m if not set.
# IDType = m200m

# An alternative way of specifying IDs is to select start and end (inclusive)
# ID values. If the IDs variable is not set, both of these values must be set.
#
# IDStart = 10
# IDEnd = 15

# ExclusionStrategy determines how to exclude IDs from the given set. This is
# useful because splashback shells are not particularly meaningful for
# subhalos. It can be set to the following modes:
# none    - No halos are removed
# subhalo - Halos flagged as subhalos in the catalog are removed
# overlap - Halos which have an R200m shell that overlaps with a larger halo's
#           R200m shell are removed
#
# ExclusionStrategy defaults to overlap if not set.
#
# ExclusionStrategy = overlap

# ExclusionRadiusMult is a multiplier of R200m applied for the sake of
# determining exclusions.
#
# ExclusionRadiusMult defaults to 1 if not set.
#
# ExclustionRadiusMult = 1

# Mult is the number of times a given ID should be repeated. This is most useful
# if you want to estimate the scatter in shell measurements for halos with a
# given set of shell parameters.
#
# Mult defaults to 1 if not set.
#
# Mult = 1`
}

// ReadConfig reads in an id.config file into config.
func (config *IDConfig) ReadConfig(fname string) error {

	vars := parse.NewConfigVars("id.config")
	vars.String(&config.idType, "IDType", "m200m")
	vars.Ints(&config.ids, "IDs", []int64{})
	vars.Int(&config.idStart, "IDStart", -1)
	vars.Int(&config.idEnd, "IDEnd", -1)
	vars.Int(&config.mult, "Mult", 1)
	vars.Int(&config.snap, "Snap", -1)
	vars.String(&config.exclusionStrategy, "ExclusionStrategy", "subhalo")
	vars.Float(&config.exclusionRadiusMult, "ExclusionRadiusMult", 1)

	if fname == "" { return nil }
	if err := parse.ReadConfig(fname, vars); err != nil { return err }
	return config.validate()
}

// validate checks whether all the fields of config are valid.
func (config *IDConfig) validate() error {
	switch config.idType {
	case "halo-id", "m200m":
	default:
		return fmt.Errorf("The 'IDType' variable is set to '%s', which I " +
			"don't recognize.", config.idType)
	}

	switch config.exclusionStrategy {
	case "none", "subhalo":
	case "overlap":
		if config.exclusionRadiusMult <= 0 {
			return fmt.Errorf("The 'ExclusionRadiusMult' varaible is set to " +
				"%g, but it needs to be positive.", config.exclusionRadiusMult)
		}
	default:
		return fmt.Errorf("The 'ExclusionStrategy' variable is set to '%s', " +
		"which I don't recognize.", config.exclusionStrategy)
	}

	// TODO: Check the ranges of the IDs as well as IDStart and IDEnd
	if len(config.ids) == 0 {
		switch {
		case config.idStart == -1 && config.idEnd == -1:
			return fmt.Errorf("'IDs' variable not set.")
		case config.idStart == -1:
			return fmt.Errorf("'IDStart variable not set.")
		case config.idEnd == -1:
			return fmt.Errorf("'IDEnd' variable not set.")
		case config.idEnd < config.idStart:
			return fmt.Errorf("'IDEnd' variable set to %d, but 'IDStart' " +
				"variable set to %d.", config.idEnd, config.idStart)
		}
	}

	switch {
	case config.snap == -1:
		return fmt.Errorf("'Snap' variable not set.")
	case config.snap < 0:
		return fmt.Errorf("'Snap' variable set to %d.", config.snap)
	}

	if config.mult <= 0 {
		return fmt.Errorf("'Mult' variable set to %d", config.mult)
	}

	return nil
}

// Run executes the ID mode of shellfish tool.
func (config *IDConfig) Run(
	flags []string, gConfig *GlobalConfig, e *env.Environment, stdin []string,
) ([]string, error) {

	if config.snap < gConfig.SnapMin || config.snap > gConfig.SnapMax {
		return nil, fmt.Errorf("'Snap' = %d, but 'SnapMin' = %d and " +
			"'SnapMax = %d'", config.snap, gConfig.SnapMin, gConfig.SnapMax)
	}

	// Get IDs and snapshots

	rawIds := getIDs(config.idStart, config.idEnd, config.ids)

	vars := &halo.VarColumns{
		ID: int(gConfig.HaloIDColumn),
		X: int(gConfig.HaloPositionColumns[0]),
		Y: int(gConfig.HaloPositionColumns[1]),
		Z: int(gConfig.HaloPositionColumns[2]),
		M200m: int(gConfig.HaloM200mColumn),
	}

	var ids, snaps []int
	switch config.idType {
	case "halo-id":
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = int(config.snap) }
		ids = rawIds
	case "m200m":
		snaps = make([]int, len(rawIds))
		for i := range snaps { snaps[i] = int(config.snap) }

		var err error
		ids, err = convertSortedIDs(rawIds, int(config.snap), vars, e)
		if err != nil { return nil, err }
	default:
		panic("Impossible")
	}

	// Tag subhalos, if neccessary.
	exclude := make([]bool, len(ids))
	switch config.exclusionStrategy {
	case "none":
	case "subhalo":
		panic("subhalo is not implemented")
	case "overlap":
		var err error
		exclude, err = findOverlapSubs(ids, snaps, vars, e, config)
		if err != nil { return nil, err }
	}

	// Generate lines
	intCols := [][]int{ids, snaps}
	floatCols := [][]float64{}
	colOrder := []int{0, 1}
	lines := catalog.FormatCols(intCols, floatCols, colOrder)

	// Filter
	fLines := []string{}
	for i := range lines {
		if !exclude[i] { fLines = append(fLines, lines[i]) }
	}

	// Multiply
	mLines := []string{}
	for i := range fLines {
		for j := 0; j < int(config.mult); j++ {
			mLines = append(mLines, fLines[i])
		}
	}

	cString := catalog.CommentString(
		[]string{"ID", "Snapshot"}, []string{}, []int{0, 1},
	)
	mLines = append([]string{cString}, mLines...)

	return mLines, nil
}

func getIDs(idStart, idEnd int64, ids []int64) []int {
	if idStart != -1 {
		out := make([]int, idEnd - idStart)
		for i := range out {
			out[i] = int(idStart) + i
		}
		return out
	} else {
		out := make([]int, len(ids))
		for i := range out {
			out[i] = int(ids[i])
		}
		return out
	}
}

func convertSortedIDs(
	rawIDs []int, snap int, vars *halo.VarColumns, e *env.Environment,
) ([]int, error) {
	maxID := 0
	for _, id := range rawIDs {
		if id > maxID { maxID = id }
	}

	rids, err := memo.ReadSortedRockstarIDs(snap, maxID, vars, e)
	if err != nil { return nil, err }

	ids := make([]int, len(rawIDs))
	for i := range ids { ids[i] = rids[rawIDs[i]] }
	return ids, nil
}

func findOverlapSubs(
	rawIDs, snaps []int, vars *halo.VarColumns,
	e *env.Environment, config *IDConfig,
) ([]bool, error) {
	isSub := make([]bool, len(rawIDs))

	// Group by snapshot.
	snapGroups := make(map[int][]int)
	groupIdxs := make(map[int][]int)
	for i, id := range rawIDs {
		snap := snaps[i]
		snapGroups[snap] = append(snapGroups[snap], id)
		groupIdxs[snap] = append(groupIdxs[snap], i)
	}

	// Load each snapshot.
	hds, _, err := memo.ReadHeaders(snaps[0], e)
	if err != nil { return nil, err }
	hd := hds[0]

	for snap, group := range snapGroups {
		rids, err := memo.ReadSortedRockstarIDs(snap, -1, vars, e)
		if err != nil { return nil, err }
		_, xs, ys, zs, _, rs, err := memo.ReadRockstar(snap, rids, vars, e)

		g := halo.NewGrid(finderCells, hd.TotalWidth, len(xs))
		g.Insert(xs, ys, zs)
		sf := halo.NewSubhaloFinder(g)
		sf.FindSubhalos(xs, ys, zs, rs, config.exclusionRadiusMult)

		for i, id := range group {
			origIdx := groupIdxs[snap][i]
			// TODO: Holy linear search, batman! Fix this.
			for j, checkID := range rids {
				if checkID == id {
					isSub[origIdx] = sf.HostCount(j) > 0
					break
				} else if j == len(rids) - 1 {
					return nil, fmt.Errorf("ID %d not in halo list.", id)
				}
			}
		}
	}
	return isSub, nil
}
<|endoftext|>


<commit_before>package azureSdkForGo

import (
	"fmt"
	"io"
	"bytes"
	"time"
	"strings"
	"errors"
	"os/exec"
	"encoding/xml"
	"github.com/MSOpenTech/azure-sdk-for-go/core/tls"
	"github.com/MSOpenTech/azure-sdk-for-go/core/http"
)

const (
	azureManagementDnsName = "https://management.core.windows.net"
	msVersionHeader = "x-ms-version"
	msVersionHeaderValue = "2014-05-01"
	contentHeader = "Content-Type"
	contentHeaderValue = "application/xml"
	requestIdHeader = "X-Ms-Request-Id"
)

func SendAzureGetRequest(url string) ([]byte, error){
	response, err := SendAzureRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}

	responseContent := getResponseBody(response)
	return responseContent, nil
}

func SendAzurePostRequest(url string, data []byte) (string, error){
	response, err := SendAzureRequest(url, "POST", data)
	if err != nil {
		return "", err
	}

	requestId := response.Header[requestIdHeader]
	return requestId[0], nil
}

func SendAzureDeleteRequest(url string) ([]byte, error){
	response, err := SendAzureRequest(url, "DELETE", nil)
	if err != nil {
		return "", err
	}

	requestId := response.Header[requestIdHeader]
	return requestId[0], nil
}

func SendAzureRequest(url string, requestType string,  data []byte) (*http.Response, error){
	client := createHttpClient()

	response, err := sendRequest(client, url, requestType, data, 5)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func sendRequest(client *http.Client, url string, requestType string, data []byte, numberOfRetries int) (*http.Response, error){
	request, reqErr := createAzureRequest(url, requestType, data)
	if reqErr != nil {
		return nil, reqErr
	}

	response, err := client.Do(request)
	if err != nil {
		if numberOfRetries == 0 {
			return nil, err
		}

		return sendRequest(client, url, requestType, data, numberOfRetries-1)
	}

	if response.StatusCode > 299 {
		responseContent := getResponseBody(response)
		azureErr := getAzureError(responseContent)
		if azureErr != nil {
			if numberOfRetries == 0 {
				return nil, azureErr
			}

			return sendRequest(client, url, requestType, data, numberOfRetries-1)
		}
	}

	return response, nil
}

func ExecuteCommand(command string) ([]byte, error) {
	parts := strings.Fields(command)
	head := parts[0]
	parts = parts[1:len(parts)]

	cmd := exec.Command(head, parts...)

	out, err := cmd.Output()

	if err != nil {
		return nil, err
	}

	return out, nil
}

func GetOperationStatus(operationId string) (*Operation, error){
	operation := new(Operation)
	url := "operations/" + operationId
	response, azureErr := SendAzureGetRequest(url)
	if azureErr != nil {
		return nil, azureErr
	}

	err := xml.Unmarshal(response, operation)
	if err != nil {
		return nil, err
	}

	return operation, nil
}

func WaitAsyncOperation(operationId string) (error) {
	status := "InProgress"
	operation := new(Operation)
	err := errors.New("")
	for status == "InProgress" {
		time.Sleep(2000 * time.Millisecond)
		operation, err = GetOperationStatus(operationId)
		if err != nil {
			return err
		}

		status = operation.Status
	}

	if status == "Failed" {
		return errors.New(operation.Error.Message)
	}

	return nil
}

func getAzureError(responseBody []byte) (error){
	error := new(AzureError)
	err := xml.Unmarshal(responseBody, error)
	if err != nil {
		return err
	}

	return error
}

func createAzureRequest(url string, requestType string,  data []byte) (*http.Request, error){
	var request *http.Request
	var err error

	url = fmt.Sprintf("%s/%s/%s", azureManagementDnsName, GetPublishSettings().SubscriptionID, url)
	if data != nil {
		body := bytes.NewBuffer(data)
		request, err = http.NewRequest(requestType, url, body)
	} else {
		request, err = http.NewRequest(requestType, url, nil)
	}

	if err != nil {
		return nil, err
	}

	request.Header.Add(msVersionHeader, msVersionHeaderValue)
	request.Header.Add(contentHeader, contentHeaderValue)

	return request, nil
}

func createHttpClient() (*http.Client){
	cert, _ := tls.X509KeyPair(GetPublishSettings().SubscriptionCert, GetPublishSettings().SubscriptionKey)

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

	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: ssl,
		},
	}

	return client
}

func getResponseBody(response *http.Response) ([]byte){

	responseBody := make([]byte, response.ContentLength)
	io.ReadFull(response.Body, responseBody)
	return responseBody
}

type AzureError struct {
	XMLName   			xml.Name `xml:"Error"`
	Code				string
	Message				string
}

func (e *AzureError) Error() string {
	return fmt.Sprintf("Code: %s, Message: %s", e.Code, e.Message)
}

type Operation struct {
	XMLName   			xml.Name `xml:"Operation"`
	ID					string
	Status				string
	HttpStatusCode		string
	Error 				AzureError
}
<commit_msg>Fixed SendAzureDeleteRequest method<commit_after>package azureSdkForGo

import (
	"fmt"
	"io"
	"bytes"
	"time"
	"strings"
	"errors"
	"os/exec"
	"encoding/xml"
	"github.com/MSOpenTech/azure-sdk-for-go/core/tls"
	"github.com/MSOpenTech/azure-sdk-for-go/core/http"
)

const (
	azureManagementDnsName = "https://management.core.windows.net"
	msVersionHeader = "x-ms-version"
	msVersionHeaderValue = "2014-05-01"
	contentHeader = "Content-Type"
	contentHeaderValue = "application/xml"
	requestIdHeader = "X-Ms-Request-Id"
)

func SendAzureGetRequest(url string) ([]byte, error){
	response, err := SendAzureRequest(url, "GET", nil)
	if err != nil {
		return nil, err
	}

	responseContent := getResponseBody(response)
	return responseContent, nil
}

func SendAzurePostRequest(url string, data []byte) (string, error){
	response, err := SendAzureRequest(url, "POST", data)
	if err != nil {
		return "", err
	}

	requestId := response.Header[requestIdHeader]
	return requestId[0], nil
}

func SendAzureDeleteRequest(url string) (string, error){
	response, err := SendAzureRequest(url, "DELETE", nil)
	if err != nil {
		return "", err
	}

	requestId := response.Header[requestIdHeader]
	return requestId[0], nil
}

func SendAzureRequest(url string, requestType string,  data []byte) (*http.Response, error){
	client := createHttpClient()

	response, err := sendRequest(client, url, requestType, data, 5)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func sendRequest(client *http.Client, url string, requestType string, data []byte, numberOfRetries int) (*http.Response, error){
	request, reqErr := createAzureRequest(url, requestType, data)
	if reqErr != nil {
		return nil, reqErr
	}

	response, err := client.Do(request)
	if err != nil {
		if numberOfRetries == 0 {
			return nil, err
		}

		return sendRequest(client, url, requestType, data, numberOfRetries-1)
	}

	if response.StatusCode > 299 {
		responseContent := getResponseBody(response)
		azureErr := getAzureError(responseContent)
		if azureErr != nil {
			if numberOfRetries == 0 {
				return nil, azureErr
			}

			return sendRequest(client, url, requestType, data, numberOfRetries-1)
		}
	}

	return response, nil
}

func ExecuteCommand(command string) ([]byte, error) {
	parts := strings.Fields(command)
	head := parts[0]
	parts = parts[1:len(parts)]

	cmd := exec.Command(head, parts...)

	out, err := cmd.Output()

	if err != nil {
		return nil, err
	}

	return out, nil
}

func GetOperationStatus(operationId string) (*Operation, error){
	operation := new(Operation)
	url := "operations/" + operationId
	response, azureErr := SendAzureGetRequest(url)
	if azureErr != nil {
		return nil, azureErr
	}

	err := xml.Unmarshal(response, operation)
	if err != nil {
		return nil, err
	}

	return operation, nil
}

func WaitAsyncOperation(operationId string) (error) {
	status := "InProgress"
	operation := new(Operation)
	err := errors.New("")
	for status == "InProgress" {
		time.Sleep(2000 * time.Millisecond)
		operation, err = GetOperationStatus(operationId)
		if err != nil {
			return err
		}

		status = operation.Status
	}

	if status == "Failed" {
		return errors.New(operation.Error.Message)
	}

	return nil
}

func getAzureError(responseBody []byte) (error){
	error := new(AzureError)
	err := xml.Unmarshal(responseBody, error)
	if err != nil {
		return err
	}

	return error
}

func createAzureRequest(url string, requestType string,  data []byte) (*http.Request, error){
	var request *http.Request
	var err error

	url = fmt.Sprintf("%s/%s/%s", azureManagementDnsName, GetPublishSettings().SubscriptionID, url)
	if data != nil {
		body := bytes.NewBuffer(data)
		request, err = http.NewRequest(requestType, url, body)
	} else {
		request, err = http.NewRequest(requestType, url, nil)
	}

	if err != nil {
		return nil, err
	}

	request.Header.Add(msVersionHeader, msVersionHeaderValue)
	request.Header.Add(contentHeader, contentHeaderValue)

	return request, nil
}

func createHttpClient() (*http.Client){
	cert, _ := tls.X509KeyPair(GetPublishSettings().SubscriptionCert, GetPublishSettings().SubscriptionKey)

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

	client := &http.Client{
		Transport: &http.Transport{
			TLSClientConfig: ssl,
		},
	}

	return client
}

func getResponseBody(response *http.Response) ([]byte){

	responseBody := make([]byte, response.ContentLength)
	io.ReadFull(response.Body, responseBody)
	return responseBody
}

type AzureError struct {
	XMLName   			xml.Name `xml:"Error"`
	Code				string
	Message				string
}

func (e *AzureError) Error() string {
	return fmt.Sprintf("Code: %s, Message: %s", e.Code, e.Message)
}

type Operation struct {
	XMLName   			xml.Name `xml:"Operation"`
	ID					string
	Status				string
	HttpStatusCode		string
	Error 				AzureError
}
<|endoftext|>