udpconn.go

package statsd

import (
	"fmt"
	"net"
	"sync/atomic"
	"time"
	"unsafe"
)

// UDPConn is an implementation of an io.WriteCLoser, encapsulating a
// net.PacketConn, that supports DNS changes during operation.
// See also the WriteCloser, TrimTrailingNewline, and UDPCheck options.
type UDPConn struct {
	conn net.PacketConn
	addr unsafe.Pointer
	stop chan struct{}
	done chan struct{}
}

func NewUDPConn(network, address string, rate time.Duration) (*UDPConn, error) {
	if rate <= 0 {
		return nil, fmt.Errorf(`invalid rate: %s`, rate)
	}

	addr, err := resolveUDPAddress(network, address)
	if err != nil {
		return nil, err
	}

	conn, err := listenUDP(network, nil)
	if err != nil {
		return nil, err
	}

	r := UDPConn{
		conn: conn,
		addr: unsafe.Pointer(addr),
		stop: make(chan struct{}, 1),
		done: make(chan struct{}),
	}

	go r.worker(network, address, rate)

	return &r, nil
}

func (x *UDPConn) Write(b []byte) (int, error) {
	return x.conn.WriteTo(b, (*net.UDPAddr)(atomic.LoadPointer(&x.addr)))
}

func (x *UDPConn) Close() (err error) {
	select {
	case x.stop <- struct{}{}:
	default:
	}
	err = x.conn.Close()
	<-x.done
	return
}

func (x *UDPConn) worker(network, address string, rate time.Duration) {
	defer close(x.done)
	ticker := time.NewTicker(rate)
	defer ticker.Stop()
	for {
		select {
		case <-x.stop:
			return
		case <-ticker.C:
			// TODO cancellation would be nice
			newAddr, err := resolveUDPAddress(network, address)
			if err != nil {
				continue
			}
			if oldAddr := (*net.UDPAddr)(x.addr); newAddr.Port == oldAddr.Port &&
				newAddr.Zone == oldAddr.Zone && newAddr.IP.Equal(oldAddr.IP) {
				continue
			}
			atomic.StorePointer(&x.addr, unsafe.Pointer(newAddr))
		}
	}
}