<commit_before>package viewLib

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/boltdb/bolt"
)

//Counter is an instance of a [pageName]pageView hash map. This is implemented with a
//mutex RW lock to stop goroutine data races
var Counter = struct {
	sync.RWMutex
	M map[string]int
}{M: make(map[string]int)}

//IPs is an instance of a [ipAdress]bool hash map. We don't care about the bool,
//using a hash map in this case just for the IP Key, as it offers a
//easy implementation on a set with quick insertion. This struct has a
//mutex RW lock to stop goroutine data races
var IPs = struct {
	sync.RWMutex
	M map[string]bool
}{M: make(map[string]bool)}

//IPList struct is used to marshal/unmarshal IP visitor data into JSON
//to be sent to current storage
type IPList struct {
	IPs map[string]bool
}

//SavePoint struct is used to marshal/unmarshal pageview data into JSON
//to be sent to current and historic storage
type SavePoint struct {
	PageCounts  map[string]int
	UniqueViews int
}

//init checks checks for previos data, sets up multithreading and then
//initiates the HTTP server. init() does not need to be called, it runs
//automatically when the package is called.
func init() {
	//checks for present DB storage and loads it into memory
	checkForRecords()

	//start goroutine to periodicly write IP and page view sets to disk
	go periodicMemoryWriter()
}

//ViewInc locks the Counter and ip set mutexes, writes to both then unlocks
func ViewInc(ip string, page string) {
	log.Println(ip + " requests " + page)

	Counter.Lock()
	Counter.M[page]++
	Counter.Unlock()
	Counter.RLock()
	Counter.RUnlock()

	IPs.Lock()
	IPs.M[ip] = true
	IPs.Unlock()
}

//periodicMemoryWriter initiates a BoltDB client, sets up a ticker and
//then wrties the IP and pageView maps to on persistant memory via BoltDB.
//This means that in the highly unlikely ;) case that the program crashes,
//a restart will reload the data and your view count won't vanish.
func periodicMemoryWriter() {
	//start the bolt client
	boltClient, err := bolt.Open("viewCounter.db", 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer boltClient.Close()

	//check and create a bucket in bolt to store the data
	boltClient.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists([]byte("historicData"))
		return nil
	})

	//start a ticker for auto uploading the IPs and view count to bolt
	//that triggers every ten minutes
	ticker := time.NewTicker(time.Minute * 10)

	for {

		<-ticker.C
		log.Println("Tick")
		fmt.Println("start:", time.Now())

		date := strconv.Itoa((time.Now().YearDay() * 10000) + time.Now().Year())
		fmt.Println(date)

		Counter.RLock()
		IPs.RLock()

		m1 := SavePoint{
			PageCounts:  Counter.M,
			UniqueViews: len(IPs.M),
		}

		m2 := IPList{
			IPs: IPs.M,
		}

		Counter.RUnlock()
		IPs.RUnlock()

		m1json, err := json.Marshal(m1)
		errLog(err)
		m2json, err := json.Marshal(m2)
		errLog(err)
		boltClient.Update(func(tx *bolt.Tx) error {

			err = tx.Bucket([]byte("historicData")).Put([]byte(date), []byte(m1json))
			errLog(err)

			err = tx.Bucket([]byte("historicData")).Put([]byte("current"), []byte(m1json))
			errLog(err)

			err = tx.Bucket([]byte("historicData")).Put([]byte("IPs"), []byte(m2json))
			errLog(err)
			return nil
		})

		fmt.Println("end:", time.Now())

	}
}

//checkForRecords is used to see if a BoltDB database is present in the file system,
//and if it is then to load the IP and pageview sets into program memory.
func checkForRecords() {
	if _, err := os.Stat("viewCounter.db"); err == nil {
		log.Println("viewCount.db database already exists; processing old entries")

		boltClient, err := bolt.Open("viewCounter.db", 0600, nil) //maybe change the 600 to a read only value
		if err != nil {
			log.Fatal(err)
		}
		defer boltClient.Close()

		var b1, b2 []byte
		boltClient.View(func(tx *bolt.Tx) error {
			// Set the value "bar" for the key "foo".
			b1 = tx.Bucket([]byte("historicData")).Get([]byte("current"))
			errLog(err)

			b2 = tx.Bucket([]byte("historicData")).Get([]byte("IPs"))
			errLog(err)

			return nil
		})

		var mjson1 SavePoint
		err = json.Unmarshal(b1, &mjson1)
		errLog(err)

		for k, v := range mjson1.PageCounts {
			Counter.M[k] = v
		}

		var mjson2 IPList
		err = json.Unmarshal(b2, &mjson2)
		errLog(err)

		for k := range mjson2.IPs {
			IPs.M[k] = true
		}

	} else {
		log.Println("viewCount.db not present; creating database")

	}
}

func errLog(err error) {
	if err != nil {
		log.Print(err)
	}
}
<commit_msg>added the SaveDuration global Var<commit_after>package viewLib

import (
	"encoding/json"
	"log"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/boltdb/bolt"
)

//Counter is an instance of a [pageName]pageView hash map. This is implemented with a
//mutex RW lock to stop goroutine data races
var Counter = struct {
	sync.RWMutex
	M map[string]int
}{M: make(map[string]int)}

//IPs is an instance of a [ipAdress]bool hash map. We don't care about the bool,
//using a hash map in this case just for the IP Key, as it offers a
//easy implementation on a set with quick insertion. This struct has a
//mutex RW lock to stop goroutine data races
var IPs = struct {
	sync.RWMutex
	M map[string]bool
}{M: make(map[string]bool)}

//RefreshTime is how often the page views and IPs are saved to disk, set at a
//default of 1 minute
var SaveDuration = time.Minute * 1

//IPList struct is used to marshal/unmarshal IP visitor data into JSON
//to be sent to current storage
type IPList struct {
	IPs map[string]bool
}

//SavePoint struct is used to marshal/unmarshal pageview data into JSON
//to be sent to current and historic storage
type SavePoint struct {
	PageCounts  map[string]int
	UniqueViews int
}

//init checks checks for previos data, sets up multithreading and then
//initiates the HTTP server. init() does not need to be called, it runs
//automatically when the package is called.
func init() {
	//checks for present DB storage and loads it into memory
	checkForRecords()

	//start goroutine to periodicly write IP and page view sets to disk
	go periodicMemoryWriter()
}

//ViewInc locks the Counter and ip set mutexes, writes to both then unlocks
func ViewInc(ip string, page string) {
	log.Println(ip + " requests " + page)

	Counter.Lock()
	Counter.M[page]++
	Counter.Unlock()
	Counter.RLock()
	Counter.RUnlock()

	IPs.Lock()
	IPs.M[ip] = true
	IPs.Unlock()
}

//periodicMemoryWriter initiates a BoltDB client, sets up a ticker and
//then wrties the IP and pageView maps to on persistant memory via BoltDB.
//This means that in the highly unlikely ;) case that the program crashes,
//a restart will reload the data and your view count won't vanish.
func periodicMemoryWriter() {
	//start the bolt client
	boltClient, err := bolt.Open("viewCounter.db", 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer boltClient.Close()

	//check and create a bucket in bolt to store the data
	boltClient.Update(func(tx *bolt.Tx) error {
		tx.CreateBucketIfNotExists([]byte("historicData"))
		return nil
	})

	//start a ticker for auto uploading the IPs and view count to bolt
	//that triggers every ten minutes
	ticker := time.NewTicker(SaveDuration)

	for {

		<-ticker.C
		//Debug fmt.Println("Save start time: ", time.Now())

		//The date is made of the day number concatinated with the year, e.g. the
		//05/01/2015 would be 52015
		date := strconv.Itoa((time.Now().YearDay() * 10000) + time.Now().Year())

		Counter.RLock()
		IPs.RLock()

		m1 := SavePoint{
			PageCounts:  Counter.M,
			UniqueViews: len(IPs.M),
		}

		m2 := IPList{
			IPs: IPs.M,
		}

		Counter.RUnlock()
		IPs.RUnlock()

		m1json, err := json.Marshal(m1)
		errLog(err)
		m2json, err := json.Marshal(m2)
		errLog(err)
		boltClient.Update(func(tx *bolt.Tx) error {

			err = tx.Bucket([]byte("historicData")).Put([]byte(date), []byte(m1json))
			errLog(err)

			err = tx.Bucket([]byte("historicData")).Put([]byte("current"), []byte(m1json))
			errLog(err)

			err = tx.Bucket([]byte("historicData")).Put([]byte("IPs"), []byte(m2json))
			errLog(err)
			return nil
		})

		//Debug fmt.Println("Save finish time: ", time.Now())

	}
}

//checkForRecords is used to see if a BoltDB database is present in the file system,
//and if it is then to load the IP and pageview sets into program memory.
func checkForRecords() {
	if _, err := os.Stat("viewCounter.db"); err == nil {
		log.Println("viewCount.db database already exists; processing old entries")

		boltClient, err := bolt.Open("viewCounter.db", 0600, nil) //maybe change the 600 to a read only value
		if err != nil {
			log.Fatal(err)
		}
		defer boltClient.Close()

		var b1, b2 []byte
		boltClient.View(func(tx *bolt.Tx) error {
			// Set the value "bar" for the key "foo".
			b1 = tx.Bucket([]byte("historicData")).Get([]byte("current"))
			errLog(err)

			b2 = tx.Bucket([]byte("historicData")).Get([]byte("IPs"))
			errLog(err)

			return nil
		})

		var mjson1 SavePoint
		err = json.Unmarshal(b1, &mjson1)
		errLog(err)

		for k, v := range mjson1.PageCounts {
			Counter.M[k] = v
		}

		var mjson2 IPList
		err = json.Unmarshal(b2, &mjson2)
		errLog(err)

		for k := range mjson2.IPs {
			IPs.M[k] = true
		}

	} else {
		log.Println("viewCount.db not present; creating database")

	}
}

func errLog(err error) {
	if err != nil {
		log.Print(err)
	}
}
<|endoftext|>


<commit_before>// Go Riemann client
package raidman

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"reflect"
	"sync"
	"time"

	"github.com/amir/raidman/proto"
	pb "github.com/golang/protobuf/proto"
)

type network interface {
	Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error)
}

type tcp struct{}

type udp struct{}

// Client represents a connection to a Riemann server
type Client struct {
	sync.Mutex
	net        network
	connection net.Conn
	timeout    time.Duration
}

// An Event represents a single Riemann event
type Event struct {
	Ttl         float32
	Time        int64
	Tags        []string
	Host        string // Defaults to os.Hostname()
	State       string
	Service     string
	Metric      interface{} // Could be Int, Float32, Float64
	Description string
	Attributes  map[string]string
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk, with a timeout of timeout
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func DialWithTimeout(netwrk, addr string, timeout time.Duration) (c *Client, err error) {
	c = new(Client)

	var cnet network
	switch netwrk {
	case "tcp", "tcp4", "tcp6":
		cnet = new(tcp)
	case "udp", "udp4", "udp6":
		cnet = new(udp)
	default:
		return nil, fmt.Errorf("dial %q: unsupported network %q", netwrk, netwrk)
	}

	c.net = cnet
	c.timeout = timeout
	c.connection, err = net.Dial(netwrk, addr)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk.
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func Dial(netwrk, addr string) (c *Client, err error) {
	return DialWithTimeout(netwrk, addr, 0)
}

func (network *tcp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	msg := &proto.Msg{}
	data, err := pb.Marshal(message)
	if err != nil {
		return msg, err
	}
	b := new(bytes.Buffer)
	if err = binary.Write(b, binary.BigEndian, uint32(len(data))); err != nil {
		return msg, err
	}
	if _, err = conn.Write(b.Bytes()); err != nil {
		return msg, err
	}
	if _, err = conn.Write(data); err != nil {
		return msg, err
	}
	var header uint32
	if err = binary.Read(conn, binary.BigEndian, &header); err != nil {
		return msg, err
	}
	response := make([]byte, header)
	if err = readFully(conn, response); err != nil {
		return msg, err
	}
	if err = pb.Unmarshal(response, msg); err != nil {
		return msg, err
	}
	if msg.GetOk() != true {
		return msg, errors.New(msg.GetError())
	}
	return msg, nil
}

func readFully(r io.Reader, p []byte) error {
	for len(p) > 0 {
		n, err := r.Read(p)
		p = p[n:]
		if err != nil {
			return err
		}
	}
	return nil
}

func (network *udp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	data, err := pb.Marshal(message)
	if err != nil {
		return nil, err
	}
	if _, err = conn.Write(data); err != nil {
		return nil, err
	}

	return nil, nil
}

func isZero(v reflect.Value) bool {
	switch v.Kind() {
	case reflect.Map:
		return v.IsNil()
	case reflect.Slice:
		zero := true
		for i := 0; i < v.Len(); i++ {
			zero = zero && isZero(v.Index(i))
		}
		return zero
	}
	zero := reflect.Zero(v.Type())
	return v.Interface() == zero.Interface()
}

func eventToPbEvent(event *Event) (*proto.Event, error) {
	var e proto.Event

	if event.Host == "" {
		event.Host, _ = os.Hostname()
	}
	t := reflect.ValueOf(&e).Elem()
	s := reflect.ValueOf(event).Elem()
	typeOfEvent := s.Type()
	for i := 0; i < s.NumField(); i++ {
		f := s.Field(i)
		value := reflect.ValueOf(f.Interface())
		if !isZero(f) {
			name := typeOfEvent.Field(i).Name
			switch name {
			case "State", "Service", "Host", "Description":
				tmp := reflect.ValueOf(pb.String(value.String()))
				t.FieldByName(name).Set(tmp)
			case "Ttl":
				tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
				t.FieldByName(name).Set(tmp)
			case "Time":
				tmp := reflect.ValueOf(pb.Int64(value.Int()))
				t.FieldByName(name).Set(tmp)
			case "Tags":
				tmp := reflect.ValueOf(value.Interface().([]string))
				t.FieldByName(name).Set(tmp)
			case "Metric":
				switch reflect.TypeOf(f.Interface()).Kind() {
				case reflect.Int:
					tmp := reflect.ValueOf(pb.Int64(int64(value.Int())))
					t.FieldByName("MetricSint64").Set(tmp)
				case reflect.Int64:
					tmp := reflect.ValueOf(pb.Int64(int64(value.Int())))
					t.FieldByName("MetricSint64").Set(tmp)
				case reflect.Float32:
					tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
					t.FieldByName("MetricF").Set(tmp)
				case reflect.Float64:
					tmp := reflect.ValueOf(pb.Float64(value.Float()))
					t.FieldByName("MetricD").Set(tmp)
				default:
					return nil, fmt.Errorf("Metric of invalid type (type %v)",
						reflect.TypeOf(f.Interface()).Kind())
				}
			case "Attributes":
				var attrs []*proto.Attribute
				for k, v := range value.Interface().(map[string]string) {
					// Copy k,v so we can take
					// pointers to the new
					// temporaries
					k_, v_ := k, v
					attrs = append(attrs, &proto.Attribute{
						Key:   &k_,
						Value: &v_,
					})
				}
				t.FieldByName(name).Set(reflect.ValueOf(attrs))
			}
		}
	}

	return &e, nil
}

func pbEventsToEvents(pbEvents []*proto.Event) []Event {
	var events []Event

	for _, event := range pbEvents {
		e := Event{
			State:       event.GetState(),
			Service:     event.GetService(),
			Host:        event.GetHost(),
			Description: event.GetDescription(),
			Ttl:         event.GetTtl(),
			Time:        event.GetTime(),
			Tags:        event.GetTags(),
		}
		if event.MetricF != nil {
			e.Metric = event.GetMetricF()
		} else if event.MetricD != nil {
			e.Metric = event.GetMetricD()
		} else {
			e.Metric = event.GetMetricSint64()
		}
		if event.Attributes != nil {
			e.Attributes = make(map[string]string, len(event.GetAttributes()))
			for _, attr := range event.GetAttributes() {
				e.Attributes[attr.GetKey()] = attr.GetValue()
			}
		}

		events = append(events, e)
	}

	return events
}

// Send sends an event to Riemann
func (c *Client) Send(event *Event) error {
	return c.SendMulti([]*Event{event})
}

// SendMulti sends multiple events to Riemann
func (c *Client) SendMulti(events []*Event) error {
	message := &proto.Msg{}

	for _, event := range events {
		e, err := eventToPbEvent(event)
		if err != nil {
			return err
		}

		message.Events = append(message.Events, e)
	}

	c.Lock()
	defer c.Unlock()

	if c.timeout > 0 {
		err := c.connection.SetDeadline(time.Now().Add(c.timeout))
		if err != nil {
			return err
		}
	}

	_, err := c.net.Send(message, c.connection)
	if err != nil {
		return err
	}

	return nil
}

// Query returns a list of events matched by query
func (c *Client) Query(q string) ([]Event, error) {
	switch c.net.(type) {
	case *udp:
		return nil, errors.New("Querying over UDP is not supported")
	}
	query := &proto.Query{}
	query.String_ = pb.String(q)
	message := &proto.Msg{}
	message.Query = query
	c.Lock()
	defer c.Unlock()
	response, err := c.net.Send(message, c.connection)
	if err != nil {
		return nil, err
	}
	return pbEventsToEvents(response.GetEvents()), nil
}

// Close closes the connection to Riemann
func (c *Client) Close() {
	c.Lock()
	c.connection.Close()
	c.Unlock()
}
<commit_msg>Add uint64 metrics support<commit_after>// Go Riemann client
package raidman

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"reflect"
	"sync"
	"time"

	"github.com/amir/raidman/proto"
	pb "github.com/golang/protobuf/proto"
)

type network interface {
	Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error)
}

type tcp struct{}

type udp struct{}

// Client represents a connection to a Riemann server
type Client struct {
	sync.Mutex
	net        network
	connection net.Conn
	timeout    time.Duration
}

// An Event represents a single Riemann event
type Event struct {
	Ttl         float32
	Time        int64
	Tags        []string
	Host        string // Defaults to os.Hostname()
	State       string
	Service     string
	Metric      interface{} // Could be Int, Float32, Float64
	Description string
	Attributes  map[string]string
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk, with a timeout of timeout
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func DialWithTimeout(netwrk, addr string, timeout time.Duration) (c *Client, err error) {
	c = new(Client)

	var cnet network
	switch netwrk {
	case "tcp", "tcp4", "tcp6":
		cnet = new(tcp)
	case "udp", "udp4", "udp6":
		cnet = new(udp)
	default:
		return nil, fmt.Errorf("dial %q: unsupported network %q", netwrk, netwrk)
	}

	c.net = cnet
	c.timeout = timeout
	c.connection, err = net.Dial(netwrk, addr)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk.
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func Dial(netwrk, addr string) (c *Client, err error) {
	return DialWithTimeout(netwrk, addr, 0)
}

func (network *tcp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	msg := &proto.Msg{}
	data, err := pb.Marshal(message)
	if err != nil {
		return msg, err
	}
	b := new(bytes.Buffer)
	if err = binary.Write(b, binary.BigEndian, uint32(len(data))); err != nil {
		return msg, err
	}
	if _, err = conn.Write(b.Bytes()); err != nil {
		return msg, err
	}
	if _, err = conn.Write(data); err != nil {
		return msg, err
	}
	var header uint32
	if err = binary.Read(conn, binary.BigEndian, &header); err != nil {
		return msg, err
	}
	response := make([]byte, header)
	if err = readFully(conn, response); err != nil {
		return msg, err
	}
	if err = pb.Unmarshal(response, msg); err != nil {
		return msg, err
	}
	if msg.GetOk() != true {
		return msg, errors.New(msg.GetError())
	}
	return msg, nil
}

func readFully(r io.Reader, p []byte) error {
	for len(p) > 0 {
		n, err := r.Read(p)
		p = p[n:]
		if err != nil {
			return err
		}
	}
	return nil
}

func (network *udp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	data, err := pb.Marshal(message)
	if err != nil {
		return nil, err
	}
	if _, err = conn.Write(data); err != nil {
		return nil, err
	}

	return nil, nil
}

func isZero(v reflect.Value) bool {
	switch v.Kind() {
	case reflect.Map:
		return v.IsNil()
	case reflect.Slice:
		zero := true
		for i := 0; i < v.Len(); i++ {
			zero = zero && isZero(v.Index(i))
		}
		return zero
	}
	zero := reflect.Zero(v.Type())
	return v.Interface() == zero.Interface()
}

func eventToPbEvent(event *Event) (*proto.Event, error) {
	var e proto.Event

	if event.Host == "" {
		event.Host, _ = os.Hostname()
	}
	t := reflect.ValueOf(&e).Elem()
	s := reflect.ValueOf(event).Elem()
	typeOfEvent := s.Type()
	for i := 0; i < s.NumField(); i++ {
		f := s.Field(i)
		value := reflect.ValueOf(f.Interface())
		if !isZero(f) {
			name := typeOfEvent.Field(i).Name
			switch name {
			case "State", "Service", "Host", "Description":
				tmp := reflect.ValueOf(pb.String(value.String()))
				t.FieldByName(name).Set(tmp)
			case "Ttl":
				tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
				t.FieldByName(name).Set(tmp)
			case "Time":
				tmp := reflect.ValueOf(pb.Int64(value.Int()))
				t.FieldByName(name).Set(tmp)
			case "Tags":
				tmp := reflect.ValueOf(value.Interface().([]string))
				t.FieldByName(name).Set(tmp)
			case "Metric":
				switch reflect.TypeOf(f.Interface()).Kind() {
				case reflect.Int, reflect.Int64:
					tmp := reflect.ValueOf(pb.Int64(int64(value.Int())))
					t.FieldByName("MetricSint64").Set(tmp)
				case reflect.Uint64:
					tmp := reflect.ValueOf(pb.Int64(int64(value.Uint())))
					t.FieldByName("MetricSint64").Set(tmp)
				case reflect.Float32:
					tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
					t.FieldByName("MetricF").Set(tmp)
				case reflect.Float64:
					tmp := reflect.ValueOf(pb.Float64(value.Float()))
					t.FieldByName("MetricD").Set(tmp)
				default:
					return nil, fmt.Errorf("Metric of invalid type (type %v)",
						reflect.TypeOf(f.Interface()).Kind())
				}
			case "Attributes":
				var attrs []*proto.Attribute
				for k, v := range value.Interface().(map[string]string) {
					// Copy k,v so we can take
					// pointers to the new
					// temporaries
					k_, v_ := k, v
					attrs = append(attrs, &proto.Attribute{
						Key:   &k_,
						Value: &v_,
					})
				}
				t.FieldByName(name).Set(reflect.ValueOf(attrs))
			}
		}
	}

	return &e, nil
}

func pbEventsToEvents(pbEvents []*proto.Event) []Event {
	var events []Event

	for _, event := range pbEvents {
		e := Event{
			State:       event.GetState(),
			Service:     event.GetService(),
			Host:        event.GetHost(),
			Description: event.GetDescription(),
			Ttl:         event.GetTtl(),
			Time:        event.GetTime(),
			Tags:        event.GetTags(),
		}
		if event.MetricF != nil {
			e.Metric = event.GetMetricF()
		} else if event.MetricD != nil {
			e.Metric = event.GetMetricD()
		} else {
			e.Metric = event.GetMetricSint64()
		}
		if event.Attributes != nil {
			e.Attributes = make(map[string]string, len(event.GetAttributes()))
			for _, attr := range event.GetAttributes() {
				e.Attributes[attr.GetKey()] = attr.GetValue()
			}
		}

		events = append(events, e)
	}

	return events
}

// Send sends an event to Riemann
func (c *Client) Send(event *Event) error {
	return c.SendMulti([]*Event{event})
}

// SendMulti sends multiple events to Riemann
func (c *Client) SendMulti(events []*Event) error {
	message := &proto.Msg{}

	for _, event := range events {
		e, err := eventToPbEvent(event)
		if err != nil {
			return err
		}

		message.Events = append(message.Events, e)
	}

	c.Lock()
	defer c.Unlock()

	if c.timeout > 0 {
		err := c.connection.SetDeadline(time.Now().Add(c.timeout))
		if err != nil {
			return err
		}
	}

	_, err := c.net.Send(message, c.connection)
	if err != nil {
		return err
	}

	return nil
}

// Query returns a list of events matched by query
func (c *Client) Query(q string) ([]Event, error) {
	switch c.net.(type) {
	case *udp:
		return nil, errors.New("Querying over UDP is not supported")
	}
	query := &proto.Query{}
	query.String_ = pb.String(q)
	message := &proto.Msg{}
	message.Query = query
	c.Lock()
	defer c.Unlock()
	response, err := c.net.Send(message, c.connection)
	if err != nil {
		return nil, err
	}
	return pbEventsToEvents(response.GetEvents()), nil
}

// Close closes the connection to Riemann
func (c *Client) Close() {
	c.Lock()
	c.connection.Close()
	c.Unlock()
}
<|endoftext|>


<commit_before>// Go Riemann client
package raidman

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"reflect"
	"sync"
	"time"

	pb "github.com/golang/protobuf/proto"
	"github.com/amir/raidman/proto"
)

type network interface {
	Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error)
}

type tcp struct{}

type udp struct{}

// Client represents a connection to a Riemann server
type Client struct {
	sync.Mutex
	net        network
	connection net.Conn
	timeout	   time.Duration
}

// An Event represents a single Riemann event
type Event struct {
	Ttl         float32
	Time        int64
	Tags        []string
	Host        string // Defaults to os.Hostname()
	State       string
	Service     string
	Metric      interface{} // Could be Int, Float32, Float64
	Description string
	Attributes  map[string]string
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk, with a timeout of timeout
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func DialWithTimeout(netwrk, addr string, timeout time.Duration) (c *Client, err error) {
	c = new(Client)

	var cnet network
	switch netwrk {
	case "tcp", "tcp4", "tcp6":
		cnet = new(tcp)
	case "udp", "udp4", "udp6":
		cnet = new(udp)
	default:
		return nil, fmt.Errorf("dial %q: unsupported network %q", netwrk, netwrk)
	}

	c.net = cnet
	c.timeout = timeout
	c.connection, err = net.Dial(netwrk, addr)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk.
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func Dial(netwrk, addr string) (c *Client, err error) {
	return DialWithTimeout(netwrk, addr, 0)
}

func (network *tcp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	msg := &proto.Msg{}
	data, err := pb.Marshal(message)
	if err != nil {
		return msg, err
	}
	b := new(bytes.Buffer)
	if err = binary.Write(b, binary.BigEndian, uint32(len(data))); err != nil {
		return msg, err
	}
	if _, err = conn.Write(b.Bytes()); err != nil {
		return msg, err
	}
	if _, err = conn.Write(data); err != nil {
		return msg, err
	}
	var header uint32
	if err = binary.Read(conn, binary.BigEndian, &header); err != nil {
		return msg, err
	}
	response := make([]byte, header)
	if err = readFully(conn, response); err != nil {
		return msg, err
	}
	if err = pb.Unmarshal(response, msg); err != nil {
		return msg, err
	}
	if msg.GetOk() != true {
		return msg, errors.New(msg.GetError())
	}
	return msg, nil
}

func readFully(r io.Reader, p []byte) error {
	for len(p) > 0 {
		n, err := r.Read(p)
		p = p[n:]
		if err != nil {
			return err
		}
	}
	return nil
}

func (network *udp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	data, err := pb.Marshal(message)
	if err != nil {
		return nil, err
	}
	if _, err = conn.Write(data); err != nil {
		return nil, err
	}

	return nil, nil
}

func eventToPbEvent(event *Event) (*proto.Event, error) {
	var e proto.Event

	if event.Host == "" {
		event.Host, _ = os.Hostname()
	}
	t := reflect.ValueOf(&e).Elem()
	s := reflect.ValueOf(event).Elem()
	typeOfEvent := s.Type()
	for i := 0; i < s.NumField(); i++ {
		f := s.Field(i)
		value := reflect.ValueOf(f.Interface())
		if reflect.Zero(f.Type()) != value && f.Interface() != nil {
			name := typeOfEvent.Field(i).Name
			switch name {
			case "State", "Service", "Host", "Description":
				tmp := reflect.ValueOf(pb.String(value.String()))
				t.FieldByName(name).Set(tmp)
			case "Ttl":
				tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
				t.FieldByName(name).Set(tmp)
			case "Time":
				tmp := reflect.ValueOf(pb.Int64(value.Int()))
				t.FieldByName(name).Set(tmp)
			case "Tags":
				tmp := reflect.ValueOf(value.Interface().([]string))
				t.FieldByName(name).Set(tmp)
			case "Metric":
				switch reflect.TypeOf(f.Interface()).Kind() {
				case reflect.Int:
					tmp := reflect.ValueOf(pb.Int64(int64(value.Int())))
					t.FieldByName("MetricSint64").Set(tmp)
				case reflect.Float32:
					tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
					t.FieldByName("MetricF").Set(tmp)
				case reflect.Float64:
					tmp := reflect.ValueOf(pb.Float64(value.Float()))
					t.FieldByName("MetricD").Set(tmp)
				default:
					return nil, fmt.Errorf("Metric of invalid type (type %v)",
						reflect.TypeOf(f.Interface()).Kind())
				}
			case "Attributes":
				var attrs []*proto.Attribute
				for k, v := range value.Interface().(map[string]string) {
					// Copy k,v so we can take
					// pointers to the new
					// temporaries
					k_, v_ := k, v
					attrs = append(attrs, &proto.Attribute{
						Key:   &k_,
						Value: &v_,
					})
				}
				t.FieldByName(name).Set(reflect.ValueOf(attrs))
			}
		}
	}

	return &e, nil
}

func pbEventsToEvents(pbEvents []*proto.Event) []Event {
	var events []Event

	for _, event := range pbEvents {
		e := Event{
			State:       event.GetState(),
			Service:     event.GetService(),
			Host:        event.GetHost(),
			Description: event.GetDescription(),
			Ttl:         event.GetTtl(),
			Time:        event.GetTime(),
			Tags:        event.GetTags(),
		}
		if event.MetricF != nil {
			e.Metric = event.GetMetricF()
		} else if event.MetricD != nil {
			e.Metric = event.GetMetricD()
		} else {
			e.Metric = event.GetMetricSint64()
		}
		if event.Attributes != nil {
			e.Attributes = make(map[string]string, len(event.GetAttributes()))
			for _, attr := range event.GetAttributes() {
				e.Attributes[attr.GetKey()] = attr.GetValue()
			}
		}

		events = append(events, e)
	}

	return events
}

// Send sends an event to Riemann
func (c *Client) Send(event *Event) error {
	e, err := eventToPbEvent(event)
	if err != nil {
		return err
	}
	message := &proto.Msg{}
	message.Events = append(message.Events, e)
	c.Lock()
	defer c.Unlock()

	if c.timeout > 0 {
		err = c.connection.SetDeadline(time.Now().Add(c.timeout))
		if err != nil {
			return err
		}
	}

	_, err = c.net.Send(message, c.connection)
	if err != nil {
		return err
	}

	return nil
}

// Query returns a list of events matched by query
func (c *Client) Query(q string) ([]Event, error) {
	switch c.net.(type) {
	case *udp:
		return nil, errors.New("Querying over UDP is not supported")
	}
	query := &proto.Query{}
	query.String_ = pb.String(q)
	message := &proto.Msg{}
	message.Query = query
	c.Lock()
	defer c.Unlock()
	response, err := c.net.Send(message, c.connection)
	if err != nil {
		return nil, err
	}
	return pbEventsToEvents(response.GetEvents()), nil
}

// Close closes the connection to Riemann
func (c *Client) Close() {
	c.Lock()
	c.connection.Close()
	c.Unlock()
}
<commit_msg>Cosmetics<commit_after>// Go Riemann client
package raidman

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"reflect"
	"sync"
	"time"

	"github.com/amir/raidman/proto"
	pb "github.com/golang/protobuf/proto"
)

type network interface {
	Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error)
}

type tcp struct{}

type udp struct{}

// Client represents a connection to a Riemann server
type Client struct {
	sync.Mutex
	net        network
	connection net.Conn
	timeout    time.Duration
}

// An Event represents a single Riemann event
type Event struct {
	Ttl         float32
	Time        int64
	Tags        []string
	Host        string // Defaults to os.Hostname()
	State       string
	Service     string
	Metric      interface{} // Could be Int, Float32, Float64
	Description string
	Attributes  map[string]string
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk, with a timeout of timeout
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func DialWithTimeout(netwrk, addr string, timeout time.Duration) (c *Client, err error) {
	c = new(Client)

	var cnet network
	switch netwrk {
	case "tcp", "tcp4", "tcp6":
		cnet = new(tcp)
	case "udp", "udp4", "udp6":
		cnet = new(udp)
	default:
		return nil, fmt.Errorf("dial %q: unsupported network %q", netwrk, netwrk)
	}

	c.net = cnet
	c.timeout = timeout
	c.connection, err = net.Dial(netwrk, addr)
	if err != nil {
		return nil, err
	}

	return c, nil
}

// Dial establishes a connection to a Riemann server at addr, on the network
// netwrk.
//
// Known networks are "tcp", "tcp4", "tcp6", "udp", "udp4", and "udp6".
func Dial(netwrk, addr string) (c *Client, err error) {
	return DialWithTimeout(netwrk, addr, 0)
}

func (network *tcp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	msg := &proto.Msg{}
	data, err := pb.Marshal(message)
	if err != nil {
		return msg, err
	}
	b := new(bytes.Buffer)
	if err = binary.Write(b, binary.BigEndian, uint32(len(data))); err != nil {
		return msg, err
	}
	if _, err = conn.Write(b.Bytes()); err != nil {
		return msg, err
	}
	if _, err = conn.Write(data); err != nil {
		return msg, err
	}
	var header uint32
	if err = binary.Read(conn, binary.BigEndian, &header); err != nil {
		return msg, err
	}
	response := make([]byte, header)
	if err = readFully(conn, response); err != nil {
		return msg, err
	}
	if err = pb.Unmarshal(response, msg); err != nil {
		return msg, err
	}
	if msg.GetOk() != true {
		return msg, errors.New(msg.GetError())
	}
	return msg, nil
}

func readFully(r io.Reader, p []byte) error {
	for len(p) > 0 {
		n, err := r.Read(p)
		p = p[n:]
		if err != nil {
			return err
		}
	}
	return nil
}

func (network *udp) Send(message *proto.Msg, conn net.Conn) (*proto.Msg, error) {
	data, err := pb.Marshal(message)
	if err != nil {
		return nil, err
	}
	if _, err = conn.Write(data); err != nil {
		return nil, err
	}

	return nil, nil
}

func eventToPbEvent(event *Event) (*proto.Event, error) {
	var e proto.Event

	if event.Host == "" {
		event.Host, _ = os.Hostname()
	}
	t := reflect.ValueOf(&e).Elem()
	s := reflect.ValueOf(event).Elem()
	typeOfEvent := s.Type()
	for i := 0; i < s.NumField(); i++ {
		f := s.Field(i)
		value := reflect.ValueOf(f.Interface())
		if reflect.Zero(f.Type()) != value && f.Interface() != nil {
			name := typeOfEvent.Field(i).Name
			switch name {
			case "State", "Service", "Host", "Description":
				tmp := reflect.ValueOf(pb.String(value.String()))
				t.FieldByName(name).Set(tmp)
			case "Ttl":
				tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
				t.FieldByName(name).Set(tmp)
			case "Time":
				tmp := reflect.ValueOf(pb.Int64(value.Int()))
				t.FieldByName(name).Set(tmp)
			case "Tags":
				tmp := reflect.ValueOf(value.Interface().([]string))
				t.FieldByName(name).Set(tmp)
			case "Metric":
				switch reflect.TypeOf(f.Interface()).Kind() {
				case reflect.Int:
					tmp := reflect.ValueOf(pb.Int64(int64(value.Int())))
					t.FieldByName("MetricSint64").Set(tmp)
				case reflect.Float32:
					tmp := reflect.ValueOf(pb.Float32(float32(value.Float())))
					t.FieldByName("MetricF").Set(tmp)
				case reflect.Float64:
					tmp := reflect.ValueOf(pb.Float64(value.Float()))
					t.FieldByName("MetricD").Set(tmp)
				default:
					return nil, fmt.Errorf("Metric of invalid type (type %v)",
						reflect.TypeOf(f.Interface()).Kind())
				}
			case "Attributes":
				var attrs []*proto.Attribute
				for k, v := range value.Interface().(map[string]string) {
					// Copy k,v so we can take
					// pointers to the new
					// temporaries
					k_, v_ := k, v
					attrs = append(attrs, &proto.Attribute{
						Key:   &k_,
						Value: &v_,
					})
				}
				t.FieldByName(name).Set(reflect.ValueOf(attrs))
			}
		}
	}

	return &e, nil
}

func pbEventsToEvents(pbEvents []*proto.Event) []Event {
	var events []Event

	for _, event := range pbEvents {
		e := Event{
			State:       event.GetState(),
			Service:     event.GetService(),
			Host:        event.GetHost(),
			Description: event.GetDescription(),
			Ttl:         event.GetTtl(),
			Time:        event.GetTime(),
			Tags:        event.GetTags(),
		}
		if event.MetricF != nil {
			e.Metric = event.GetMetricF()
		} else if event.MetricD != nil {
			e.Metric = event.GetMetricD()
		} else {
			e.Metric = event.GetMetricSint64()
		}
		if event.Attributes != nil {
			e.Attributes = make(map[string]string, len(event.GetAttributes()))
			for _, attr := range event.GetAttributes() {
				e.Attributes[attr.GetKey()] = attr.GetValue()
			}
		}

		events = append(events, e)
	}

	return events
}

// Send sends an event to Riemann
func (c *Client) Send(event *Event) error {
	e, err := eventToPbEvent(event)
	if err != nil {
		return err
	}
	message := &proto.Msg{}
	message.Events = append(message.Events, e)
	c.Lock()
	defer c.Unlock()

	if c.timeout > 0 {
		err = c.connection.SetDeadline(time.Now().Add(c.timeout))
		if err != nil {
			return err
		}
	}

	_, err = c.net.Send(message, c.connection)
	if err != nil {
		return err
	}

	return nil
}

// Query returns a list of events matched by query
func (c *Client) Query(q string) ([]Event, error) {
	switch c.net.(type) {
	case *udp:
		return nil, errors.New("Querying over UDP is not supported")
	}
	query := &proto.Query{}
	query.String_ = pb.String(q)
	message := &proto.Msg{}
	message.Query = query
	c.Lock()
	defer c.Unlock()
	response, err := c.net.Send(message, c.connection)
	if err != nil {
		return nil, err
	}
	return pbEventsToEvents(response.GetEvents()), nil
}

// Close closes the connection to Riemann
func (c *Client) Close() {
	c.Lock()
	c.connection.Close()
	c.Unlock()
}
<|endoftext|>


<commit_before>package stats

import (
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

var _ = log.Println

type ServerStats struct {
	mutex       sync.Mutex
	phaseCounts map[string]int
	*CpuStatSampler
	*DiskStatSampler
	PhaseOrder []string
}

func NewServerStats() *ServerStats {
	return &ServerStats{
		phaseCounts:     map[string]int{},
		CpuStatSampler:  NewCpuStatSampler(),
		DiskStatSampler: NewDiskStatSampler(),
	}
}

func (me *ServerStats) Enter(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]++
}

func (me *ServerStats) Exit(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]--
}

func CpuStatsWriteHttp(w http.ResponseWriter, stats []CpuStat, disk []DiskStat) {
	if len(stats) == 0 {
		return
	}

	statm := CpuStat{}
	stat5 := CpuStat{}

	diskm := DiskStat{}
	disk5 := DiskStat{}

	count5 := int64(0)
	for i, v := range stats {
		statm = statm.Add(v)
		diskm.Add(&disk[i])
		if i >= len(stats)-5 {
			count5++
			stat5 = stat5.Add(v)

			disk5.Add(&disk[i])
		}
	}

	printChild := statm.ChildCpu+statm.ChildSys > 0
	chHeader := ""
	if printChild {
		chHeader = "child cpu (ms)	child sys (ms)	"
	}
	fmt.Fprintf(w, "

self cpu (ms)	self sys (ms)	%stotal (ms)	read ops	write ops
",
		chHeader)
	for i, v := range stats {
		if i < len(stats)-5 {
			continue
		}

		chRow := ""
		if printChild {
			chRow = fmt.Sprintf("%d	%d	",
				v.ChildCpu/time.Millisecond, v.ChildSys/time.Millisecond)
		}

		fmt.Fprintf(w, "
%d	%d	%s%d	%d	%d
",
			v.SelfCpu/time.Millisecond, v.SelfSys/time.Millisecond, chRow,
			(v.Total())/time.Millisecond,
			disk[i].ReadsCompleted,
			disk[i].WritesCompleted)
	}
	fmt.Fprintf(w, "
")

	fmt.Fprintf(w, "
CPU (last min): %d s self %d s sys", statm.SelfCpu/time.Second, statm.SelfSys/time.Second)
	if printChild {
		fmt.Fprintf(w, " %d s child %d s sys", statm.ChildCpu/time.Second, statm.ChildSys/time.Second)
	}

	fmt.Fprintf(w, " %.2f CPU", float64(statm.Total())/1e9/float64(len(stats)))
	fmt.Fprintf(w, "
CPU (last %ds): %.2f self %.2f sys",
		count5, float64(stat5.SelfCpu)/float64(time.Second), float64(stat5.SelfSys)/float64(time.Second))
	if printChild {
		fmt.Fprintf(w, "%.2f s child %.2f s sys",
			float64(stat5.ChildCpu)/float64(time.Second),
			float64(stat5.ChildSys)/float64(time.Second))
	}
	fmt.Fprintf(w, " %.2f CPU", float64(stat5.Total())/1e9/float64(count5))
	fmt.Fprintf(w, "
Disk (last %ds): Read %.2f/s Write %.2f/s",
		count5,
		float64(disk5.MergedReadsCompleted)/(float64(count5)*float64(time.Second)),
		float64(disk5.WritesCompleted)/(float64(5)*float64(time.Second)))
}

func CountStatsWriteHttp(w http.ResponseWriter, names []string, counts []int) {
	fmt.Fprintf(w, "
")
	for i, c := range counts {
		fmt.Fprintf(w, "
Jobs in phase %s: %d ", names[i], c)
	}
	fmt.Fprintf(w, "
")
}

func (me *ServerStats) WriteHttp(w http.ResponseWriter) {
	cpu := me.CpuStats()
	disk := me.DiskStats()

	l := len(cpu)
	if len(disk) < l {
		l = len(disk)
	}
	CpuStatsWriteHttp(w, cpu[:l], disk[:l])
	CountStatsWriteHttp(w, me.PhaseOrder, me.PhaseCounts())
}

func (me *ServerStats) PhaseCounts() (r []int) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for _, n := range me.PhaseOrder {
		r = append(r, me.phaseCounts[n])
	}
	return r
}
<commit_msg>stats: don't provide stats if disk/memory stat sizes are out of sync<commit_after>package stats

import (
	"fmt"
	"log"
	"net/http"
	"sync"
	"time"
)

var _ = log.Println

type ServerStats struct {
	mutex       sync.Mutex
	phaseCounts map[string]int
	*CpuStatSampler
	*DiskStatSampler
	PhaseOrder []string
}

func NewServerStats() *ServerStats {
	return &ServerStats{
		phaseCounts:     map[string]int{},
		CpuStatSampler:  NewCpuStatSampler(),
		DiskStatSampler: NewDiskStatSampler(),
	}
}

func (me *ServerStats) Enter(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]++
}

func (me *ServerStats) Exit(phase string) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	me.phaseCounts[phase]--
}

func CpuStatsWriteHttp(w http.ResponseWriter, stats []CpuStat, disk []DiskStat) {
	if len(stats) == 0 {
		return
	}
	if len(stats) != len(disk) {
		return
	}
	statm := CpuStat{}
	stat5 := CpuStat{}

	diskm := DiskStat{}
	disk5 := DiskStat{}

	count5 := int64(0)
	for i, v := range stats {
		statm = statm.Add(v)
		diskm.Add(&disk[i])
		if i >= len(stats)-5 {
			count5++
			stat5 = stat5.Add(v)

			disk5.Add(&disk[i])
		}
	}

	printChild := statm.ChildCpu+statm.ChildSys > 0
	chHeader := ""
	if printChild {
		chHeader = "child cpu (ms)	child sys (ms)	"
	}
	fmt.Fprintf(w, "

self cpu (ms)	self sys (ms)	%stotal (ms)	read ops	write ops
",
		chHeader)
	for i, v := range stats {
		if i < len(stats)-5 {
			continue
		}

		chRow := ""
		if printChild {
			chRow = fmt.Sprintf("%d	%d	",
				v.ChildCpu/time.Millisecond, v.ChildSys/time.Millisecond)
		}

		fmt.Fprintf(w, "
%d	%d	%s%d	%d	%d
",
			v.SelfCpu/time.Millisecond, v.SelfSys/time.Millisecond, chRow,
			(v.Total())/time.Millisecond,
			disk[i].ReadsCompleted,
			disk[i].WritesCompleted)
	}
	fmt.Fprintf(w, "
")

	fmt.Fprintf(w, "
CPU (last min): %d s self %d s sys", statm.SelfCpu/time.Second, statm.SelfSys/time.Second)
	if printChild {
		fmt.Fprintf(w, " %d s child %d s sys", statm.ChildCpu/time.Second, statm.ChildSys/time.Second)
	}

	fmt.Fprintf(w, " %.2f CPU", float64(statm.Total())/1e9/float64(len(stats)))
	fmt.Fprintf(w, "
CPU (last %ds): %.2f self %.2f sys",
		count5, float64(stat5.SelfCpu)/float64(time.Second), float64(stat5.SelfSys)/float64(time.Second))
	if printChild {
		fmt.Fprintf(w, "%.2f s child %.2f s sys",
			float64(stat5.ChildCpu)/float64(time.Second),
			float64(stat5.ChildSys)/float64(time.Second))
	}
	fmt.Fprintf(w, " %.2f CPU", float64(stat5.Total())/1e9/float64(count5))
	fmt.Fprintf(w, "
Disk (last %ds): Read %.2f/s Write %.2f/s",
		count5,
		float64(disk5.MergedReadsCompleted)/(float64(count5)*float64(time.Second)),
		float64(disk5.WritesCompleted)/(float64(5)*float64(time.Second)))
}

func CountStatsWriteHttp(w http.ResponseWriter, names []string, counts []int) {
	fmt.Fprintf(w, "
")
	for i, c := range counts {
		fmt.Fprintf(w, "
Jobs in phase %s: %d ", names[i], c)
	}
	fmt.Fprintf(w, "
")
}

func (me *ServerStats) WriteHttp(w http.ResponseWriter) {
	cpu := me.CpuStats()
	disk := me.DiskStats()

	l := len(cpu)
	if len(disk) < l {
		l = len(disk)
	}
	CpuStatsWriteHttp(w, cpu[:l], disk[:l])
	CountStatsWriteHttp(w, me.PhaseOrder, me.PhaseCounts())
}

func (me *ServerStats) PhaseCounts() (r []int) {
	me.mutex.Lock()
	defer me.mutex.Unlock()
	for _, n := range me.PhaseOrder {
		r = append(r, me.phaseCounts[n])
	}
	return r
}
<|endoftext|>


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

package view

import (
	"fmt"
	"time"

	"go.opencensus.io/stats"
	"go.opencensus.io/stats/internal"
	"go.opencensus.io/tag"
)

func init() {
	defaultWorker = newWorker()
	go defaultWorker.start()
	internal.DefaultRecorder = record
}

type measureRef struct {
	measure string
	views   map[*viewInternal]struct{}
}

type worker struct {
	measures   map[string]*measureRef
	views      map[string]*viewInternal
	startTimes map[*viewInternal]time.Time

	timer      *time.Ticker
	c          chan command
	quit, done chan bool
}

var defaultWorker *worker

var defaultReportingDuration = 10 * time.Second

// Find returns a subscribed view associated with this name.
// If no subscribed view is found, nil is returned.
func Find(name string) (v *View) {
	req := &getViewByNameReq{
		name: name,
		c:    make(chan *getViewByNameResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.v
}

// Deprecated: Use the Register function
func Subscribe(v ...*View) error {
	return Register(v...)
}

// Deprecated: Use the Unregister function
func Unsubscribe(v ...*View) error {
	Unregister(v...)
	return nil
}

// Deprecated: Use the Register function.
func (v *View) Register() error {
	return Register(v)
}

// Register begins collecting data for the given views.
// Once a view is subscribed, it reports data to the registered exporters.
func Register(views ...*View) error {
	for _, v := range views {
		if err := v.canonicalize(); err != nil {
			return err
		}
	}
	req := &registerViewReq{
		views: views,
		err:   make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// Unregister the given views. Data will not longer be exported for these views
// after Unregister returns.
// It is not necessary to unregister from views you expect to collect for the
// duration of your program execution.
func Unregister(views ...*View) {
	names := make([]string, len(views))
	for i := range views {
		names[i] = views[i].Name
	}
	req := &unsubscribeFromViewReq{
		views: names,
		done:  make(chan struct{}),
	}
	defaultWorker.c <- req
	<-req.done
}

// Deprecated: Use the Unregister function instead.
func (v *View) Unregister() error {
	if v == nil {
		return nil
	}
	Unregister(v)
	return nil
}

func RetrieveData(viewName string) ([]*Row, error) {
	req := &retrieveDataReq{
		now: time.Now(),
		v:   viewName,
		c:   make(chan *retrieveDataResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.rows, resp.err
}

func record(tags *tag.Map, ms interface{}) {
	req := &recordReq{
		tm: tags,
		ms: ms.([]stats.Measurement),
	}
	defaultWorker.c <- req
}

// SetReportingPeriod sets the interval between reporting aggregated views in
// the program. If duration is less than or
// equal to zero, it enables the default behavior.
func SetReportingPeriod(d time.Duration) {
	// TODO(acetechnologist): ensure that the duration d is more than a certain
	// value. e.g. 1s
	req := &setReportingPeriodReq{
		d: d,
		c: make(chan bool),
	}
	defaultWorker.c <- req
	<-req.c // don't return until the timer is set to the new duration.
}

func newWorker() *worker {
	return &worker{
		measures:   make(map[string]*measureRef),
		views:      make(map[string]*viewInternal),
		startTimes: make(map[*viewInternal]time.Time),
		timer:      time.NewTicker(defaultReportingDuration),
		c:          make(chan command, 1024),
		quit:       make(chan bool),
		done:       make(chan bool),
	}
}

func (w *worker) start() {
	for {
		select {
		case cmd := <-w.c:
			if cmd != nil {
				cmd.handleCommand(w)
			}
		case <-w.timer.C:
			w.reportUsage(time.Now())
		case <-w.quit:
			w.timer.Stop()
			close(w.c)
			w.done <- true
			return
		}
	}
}

func (w *worker) stop() {
	w.quit <- true
	<-w.done
}

func (w *worker) getMeasureRef(name string) *measureRef {
	if mr, ok := w.measures[name]; ok {
		return mr
	}
	mr := &measureRef{
		measure: name,
		views:   make(map[*viewInternal]struct{}),
	}
	w.measures[name] = mr
	return mr
}

func (w *worker) tryRegisterView(v *View) (*viewInternal, error) {
	vi, err := newViewInternal(v)
	if err != nil {
		return nil, err
	}
	if x, ok := w.views[vi.view.Name]; ok {
		if !x.view.same(vi.view) {
			return nil, fmt.Errorf("cannot subscribe view %q; a different view with the same name is already subscribed", v.Name)
		}

		// the view is already registered so there is nothing to do and the
		// command is considered successful.
		return x, nil
	}
	w.views[vi.view.Name] = vi
	ref := w.getMeasureRef(vi.view.Measure.Name())
	ref.views[vi] = struct{}{}
	return vi, nil
}

func (w *worker) reportUsage(now time.Time) {
	for _, v := range w.views {
		if !v.isSubscribed() {
			continue
		}
		rows := v.collectedRows()
		_, ok := w.startTimes[v]
		if !ok {
			w.startTimes[v] = now
		}
		// Make sure collector is never going
		// to mutate the exported data.
		rows = deepCopyRowData(rows)
		viewData := &Data{
			View:  v.view,
			Start: w.startTimes[v],
			End:   time.Now(),
			Rows:  rows,
		}
		exportersMu.Lock()
		for e := range exporters {
			e.ExportView(viewData)
		}
		exportersMu.Unlock()
	}
}

func deepCopyRowData(rows []*Row) []*Row {
	newRows := make([]*Row, 0, len(rows))
	for _, r := range rows {
		newRows = append(newRows, &Row{
			Data: r.Data.clone(),
			Tags: r.Tags,
		})
	}
	return newRows
}
<commit_msg>Revert and keep deprecated *View.*ubscribe<commit_after>// Copyright 2017, OpenCensus Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

package view

import (
	"fmt"
	"time"

	"go.opencensus.io/stats"
	"go.opencensus.io/stats/internal"
	"go.opencensus.io/tag"
)

func init() {
	defaultWorker = newWorker()
	go defaultWorker.start()
	internal.DefaultRecorder = record
}

type measureRef struct {
	measure string
	views   map[*viewInternal]struct{}
}

type worker struct {
	measures   map[string]*measureRef
	views      map[string]*viewInternal
	startTimes map[*viewInternal]time.Time

	timer      *time.Ticker
	c          chan command
	quit, done chan bool
}

var defaultWorker *worker

var defaultReportingDuration = 10 * time.Second

// Find returns a subscribed view associated with this name.
// If no subscribed view is found, nil is returned.
func Find(name string) (v *View) {
	req := &getViewByNameReq{
		name: name,
		c:    make(chan *getViewByNameResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.v
}

// Deprecated: Use the Register function
func Subscribe(v ...*View) error {
	return Register(v...)
}

// Deprecated: Use the Unregister function
func Unsubscribe(v ...*View) error {
	Unregister(v...)
	return nil
}

// Deprecated: Use the Register function.
func (v *View) Subscribe() error {
	return Register(v)
}

// Register begins collecting data for the given views.
// Once a view is subscribed, it reports data to the registered exporters.
func Register(views ...*View) error {
	for _, v := range views {
		if err := v.canonicalize(); err != nil {
			return err
		}
	}
	req := &registerViewReq{
		views: views,
		err:   make(chan error),
	}
	defaultWorker.c <- req
	return <-req.err
}

// Unregister the given views. Data will not longer be exported for these views
// after Unregister returns.
// It is not necessary to unregister from views you expect to collect for the
// duration of your program execution.
func Unregister(views ...*View) {
	names := make([]string, len(views))
	for i := range views {
		names[i] = views[i].Name
	}
	req := &unsubscribeFromViewReq{
		views: names,
		done:  make(chan struct{}),
	}
	defaultWorker.c <- req
	<-req.done
}

// Deprecated: Use the Unregister function instead.
func (v *View) Unsubscribe() error {
	if v == nil {
		return nil
	}
	Unregister(v)
	return nil
}

func RetrieveData(viewName string) ([]*Row, error) {
	req := &retrieveDataReq{
		now: time.Now(),
		v:   viewName,
		c:   make(chan *retrieveDataResp),
	}
	defaultWorker.c <- req
	resp := <-req.c
	return resp.rows, resp.err
}

func record(tags *tag.Map, ms interface{}) {
	req := &recordReq{
		tm: tags,
		ms: ms.([]stats.Measurement),
	}
	defaultWorker.c <- req
}

// SetReportingPeriod sets the interval between reporting aggregated views in
// the program. If duration is less than or
// equal to zero, it enables the default behavior.
func SetReportingPeriod(d time.Duration) {
	// TODO(acetechnologist): ensure that the duration d is more than a certain
	// value. e.g. 1s
	req := &setReportingPeriodReq{
		d: d,
		c: make(chan bool),
	}
	defaultWorker.c <- req
	<-req.c // don't return until the timer is set to the new duration.
}

func newWorker() *worker {
	return &worker{
		measures:   make(map[string]*measureRef),
		views:      make(map[string]*viewInternal),
		startTimes: make(map[*viewInternal]time.Time),
		timer:      time.NewTicker(defaultReportingDuration),
		c:          make(chan command, 1024),
		quit:       make(chan bool),
		done:       make(chan bool),
	}
}

func (w *worker) start() {
	for {
		select {
		case cmd := <-w.c:
			if cmd != nil {
				cmd.handleCommand(w)
			}
		case <-w.timer.C:
			w.reportUsage(time.Now())
		case <-w.quit:
			w.timer.Stop()
			close(w.c)
			w.done <- true
			return
		}
	}
}

func (w *worker) stop() {
	w.quit <- true
	<-w.done
}

func (w *worker) getMeasureRef(name string) *measureRef {
	if mr, ok := w.measures[name]; ok {
		return mr
	}
	mr := &measureRef{
		measure: name,
		views:   make(map[*viewInternal]struct{}),
	}
	w.measures[name] = mr
	return mr
}

func (w *worker) tryRegisterView(v *View) (*viewInternal, error) {
	vi, err := newViewInternal(v)
	if err != nil {
		return nil, err
	}
	if x, ok := w.views[vi.view.Name]; ok {
		if !x.view.same(vi.view) {
			return nil, fmt.Errorf("cannot subscribe view %q; a different view with the same name is already subscribed", v.Name)
		}

		// the view is already registered so there is nothing to do and the
		// command is considered successful.
		return x, nil
	}
	w.views[vi.view.Name] = vi
	ref := w.getMeasureRef(vi.view.Measure.Name())
	ref.views[vi] = struct{}{}
	return vi, nil
}

func (w *worker) reportUsage(now time.Time) {
	for _, v := range w.views {
		if !v.isSubscribed() {
			continue
		}
		rows := v.collectedRows()
		_, ok := w.startTimes[v]
		if !ok {
			w.startTimes[v] = now
		}
		// Make sure collector is never going
		// to mutate the exported data.
		rows = deepCopyRowData(rows)
		viewData := &Data{
			View:  v.view,
			Start: w.startTimes[v],
			End:   time.Now(),
			Rows:  rows,
		}
		exportersMu.Lock()
		for e := range exporters {
			e.ExportView(viewData)
		}
		exportersMu.Unlock()
	}
}

func deepCopyRowData(rows []*Row) []*Row {
	newRows := make([]*Row, 0, len(rows))
	for _, r := range rows {
		newRows = append(newRows, &Row{
			Data: r.Data.clone(),
			Tags: r.Tags,
		})
	}
	return newRows
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/coreos/go-etcd/etcd"
	"log"
	"strings"
)

type watcher struct {
	etcdLeader   string
	etcdMachines []string
	client       *etcd.Client
}

func NewWatcher(etcdURLS string) *watcher {
	w := &watcher{}
	w.Init(etcdURLS)
	return w
}

func (w *watcher) Init(etcdURLS string) {
	w.client = etcd.NewClient(strings.Split(etcdURLS, ","))

	if len(w.etcdMachines) > 0 {
		w.client.SetCluster(w.etcdMachines)
	}
}

func (w *watcher) StartApplications(p *proxy) {
	w.loadApplications(p)

	go func() {
		appsChannel := make(chan *etcd.Response, 10)
		w.watchApplications(p, appsChannel)

		w.client.Watch("applications", 0, appsChannel, nil)
	}()
}

func (w *watcher) loadApplications(p *proxy) {
	values, err := w.client.Get("applications")

	if err == nil {
		for _, entry := range values {
			app := strings.Split(entry.Key, "/")[2]
			w.registerApp(app, p)
		}
	}
}

func (w *watcher) watchApplications(p *proxy, appsChannel chan *etcd.Response) {
	for {
		entry := <-appsChannel
		app := strings.Split(entry.Key, "/")[2]

		w.registerApp(app, p)
	}
}

func (w *watcher) registerApp(app string, p *proxy) {
	values, err := w.client.Get("applications/" + app)

	if err != nil {
		log.Printf("Error getting settings for: %s\nReason: %s", app, err.Error())
	} else {
		a := &application{Name: app}

		for _, value := range values {
			switch value.Key {
			case "/applications/" + app + "/port":
				a.Port = value.Value
			case "/applications/" + app + "/test":
				a.Test = value.Value
			}
		}

		p.Route(a)
	}
}
<commit_msg>More idiomatic code, fix watching<commit_after>package main

import (
	"github.com/coreos/go-etcd/etcd"
	"log"
	"strings"
)

type watcher struct {
	etcdLeader   string
	etcdMachines []string
	client       *etcd.Client
}

func NewWatcher(etcdURLS string) *watcher {
	w := &watcher{}
	w.Init(etcdURLS)
	return w
}

func (w *watcher) Init(etcdURLS string) {
	w.client = etcd.NewClient(strings.Split(etcdURLS, ","))

	if len(w.etcdMachines) > 0 {
		w.client.SetCluster(w.etcdMachines)
	}
}

func (w *watcher) StartApplications(p *proxy) {
	go w.loadApplications(p)
	go w.watchApplications(p)
}

func (w *watcher) loadApplications(p *proxy) {
	values, err := w.client.Get("applications", true, true)

	if err == nil {
		for _, entry := range values.Node.Nodes {
			app := strings.Split(entry.Key, "/")[2]
			w.registerApp(app, p)
		}
	}
}

func (w *watcher) watchApplications(p *proxy) {
	appsChannel := make(chan *etcd.Response, 10)

	go w.client.Watch("applications", 0, true, appsChannel, nil)

	for entry := range appsChannel {
		app := strings.Split(entry.Node.Key, "/")[2]

		w.registerApp(app, p)
	}
}

func (w *watcher) registerApp(app string, p *proxy) {
	values, err := w.client.Get("applications/"+app, true, true)

	if err != nil {
		log.Printf("Error getting settings for: %s\nReason: %s", app, err.Error())
	} else {
		a := &application{Name: app}

		for _, value := range values.Node.Nodes {
			switch value.Key {
			case "/applications/" + app + "/port":
				a.Port = value.Value
			case "/applications/" + app + "/test":
				a.Test = value.Value
			}
		}

		p.Route(a)
	}
}
<|endoftext|>


<commit_before>// watches the current directory for changes and runs the specificed program on change

package main

import (
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"time"
)

var help = `watcher [command to execute]`

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [flags] [command to execute and args]\n", os.Args[0])
	flag.PrintDefaults()
}

var verbose = flag.Bool("v", false, "verbose")
var recurse = flag.Bool("r", true, "recurse")
var quiet = flag.Int("quiet", 800, "quiet period after command execution in milliseconds")

func main() {
	flag.Usage = usage
	flag.Parse()

	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Fatal(err)
	}
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	cmd, args := flag.Args()[0], flag.Args()[1:]

	fileEvents := make(chan *fsnotify.FileEvent, 100)

	// start watchAndExecute goroutine
	go watchAndExecute(fileEvents, cmd, args)

	// pipe all events to fileEvents (for buffering and draining)
	go func() {
		for {
			select {
			case ev := <-watcher.Event:
				fileEvents <- ev
				// @todo handle created/renamed/deleted dirs
			case err := <-watcher.Error:
				log.Println("fsnotify error:", err)
			}
		}
	}()

	cwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	if *recurse {
		err = watchDirAndChildren(watcher, cwd);
	} else {
		err = watcher.Watch(cwd)
	}
	if err != nil {
		log.Fatal(err)
	}
	<-make(chan struct{})
	watcher.Close()
}

// Execute cmd with args when a file event occurs
func watchAndExecute(fileEvents chan *fsnotify.FileEvent, cmd string, args []string) {
	for {
		// execute command
		c := exec.Command(cmd, args...)
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr
		c.Stdin = os.Stdin

		fmt.Fprintln(os.Stderr, "running", cmd, args)
		if err := c.Run(); err != nil {
			fmt.Fprintln(os.Stderr, "error running:", err)
		}
		if *verbose {
			fmt.Fprintln(os.Stderr, "done.")
		}
		// drain until quiet period is over
		drainUntil(time.After(time.Duration(*quiet)*time.Millisecond), fileEvents)
		ev := <-fileEvents
		if *verbose {
			fmt.Fprintln(os.Stderr, "File changed:", ev)
		}
	}
}

// Add dir and children (recursively) to watcher
func watchDirAndChildren(watcher *fsnotify.Watcher, dir string) error {
	return filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			if *verbose {
				fmt.Fprintln(os.Stderr, "Watching", path)
			}
			if err := watcher.Watch(path); err != nil {
                            return err
                        }
		}
                return nil
	})
}

// Drain events from channel until a particular time
func drainUntil(until <-chan time.Time, c chan *fsnotify.FileEvent) {
	for {
		select {
		case <-c:
		case <-until:
			return
		}
	}
}
<commit_msg>Simplify types, change to recursion depth.<commit_after>// watches the current directory for changes and runs the specificed program on change

package main

import (
	"flag"
	"fmt"
	"github.com/howeyc/fsnotify"
	"log"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"
)

var help = `watcher [command to execute]`

func usage() {
	fmt.Fprintf(os.Stderr, "usage: %s [flags] [command to execute and args]\n", os.Args[0])
	flag.PrintDefaults()
}

var verbose = flag.Bool("v", false, "verbose")
var depth = flag.Int("d", 1, "recursion depth")
var quiet = flag.Int("quiet", 800, "quiet period after command execution in milliseconds")

func main() {
	flag.Usage = usage
	flag.Parse()

	watcher, err := newWatcher()
	if err != nil {
		log.Fatal(err)
	}
	if len(flag.Args()) == 0 {
		flag.Usage()
		os.Exit(1)
	}
	cmd, args := flag.Args()[0], flag.Args()[1:]

	fileEvents := make(chan interface{}, 100)

	// start watchAndExecute goroutine
	go watchAndExecute(fileEvents, cmd, args)

	// pipe all events to fileEvents (for buffering and draining)
	go func() {
		for {
			select {
			case ev := <-watcher.Event:
				fileEvents <- ev
				// @todo handle created/renamed/deleted dirs
			case err := <-watcher.Error:
				log.Println("fsnotify error:", err)
			}
		}
	}()

	cwd, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	err = watcher.watchDirAndChildren(cwd, *depth)
	if err != nil {
		log.Fatal(err)
	}
	<-make(chan struct{})
	watcher.Close()
}

type watcher struct {
	*fsnotify.Watcher
}

func newWatcher() (watcher, error) {
	fsnw, err := fsnotify.NewWatcher()
	return watcher{fsnw}, err
}

// Execute cmd with args when a file event occurs
func watchAndExecute(fileEvents chan interface{}, cmd string, args []string) {
	for {
		// execute command
		c := exec.Command(cmd, args...)
		c.Stdout = os.Stdout
		c.Stderr = os.Stderr
		c.Stdin = os.Stdin

		fmt.Fprintln(os.Stderr, "running", cmd, args)
		if err := c.Run(); err != nil {
			fmt.Fprintln(os.Stderr, "error running:", err)
		}
		if *verbose {
			fmt.Fprintln(os.Stderr, "done.")
		}
		// drain until quiet period is over
		drainFor(*quiet, fileEvents)
		ev := <-fileEvents
		if *verbose {
			fmt.Fprintln(os.Stderr, "File changed:", ev)
		}
	}
}

// Add dir and children (recursively) to watcher
func (w watcher) watchDirAndChildren(path string, depth int) error {
	if err := w.Watch(path); err != nil {
		return err
	}
	baseNumSeps := strings.Count(path, string(os.PathSeparator))
	return filepath.Walk(path, func(path string, info os.FileInfo, err error) error {
		if info.IsDir() {
			pathDepth := strings.Count(path, string(os.PathSeparator)) - baseNumSeps
			if pathDepth > depth {
				return filepath.SkipDir
			}
			if *verbose {
				fmt.Fprintln(os.Stderr, "Watching", path)
			}
			if err := w.Watch(path); err != nil {
				return err
			}
		}
		return nil
	})
}

// Drain events from channel until a particular time
func drainFor(drainTimeMs int, c chan interface{}) {
	for {
		select {
		case <-c:
		case <-time.After(time.Duration(drainTimeMs) * time.Millisecond):
			return
		}
	}
}
<|endoftext|>


<commit_before>package client

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"time"

	mesosjson "github.com/Dataman-Cloud/swan/mesosproto/json"
)

type Client struct {
	StreamID string
	url      string
	client   *http.Client
}

func New(addr, path string) *Client {
	return &Client{
		url: "http://" + addr + path,
		client: &http.Client{
			Transport: &http.Transport{
				Dial: (&net.Dialer{
					Timeout:   10 * time.Second,
					KeepAlive: 30 * time.Second,
				}).Dial,
			},
		},
	}
}

func (c *Client) Send(payload []byte) (*http.Response, error) {

	httpReq, err := http.NewRequest("POST", c.url, bytes.NewReader(payload))
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/x-protobuf")
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("User-Agent", "mesos-demo/0.1")
	if c.StreamID != "" {
		httpReq.Header.Set("Mesos-Stream-Id", c.StreamID)
	}
	//log.Printf("SENDING:%v", httpReq)

	httpResp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("Unable to do request: %s", err)
	}
	if httpResp.Header.Get("Mesos-Stream-Id") != "" {
		c.StreamID = httpResp.Header.Get("Mesos-Stream-Id")
	}
	return httpResp, nil
}

func (c *Client) SendAsJson(call *mesosjson.Call) (*http.Response, error) {
	payload := new(bytes.Buffer)
	if err := json.NewEncoder(payload).Encode(call); err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequest("POST", c.url, payload)
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("User-Agent", "mesos-demo/0.1")

	httpResp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("Unable to do request: %s", err)
	}
	c.StreamID = httpResp.Header.Get("Mesos-Stream-Id")
	log.Println("Stream-ID: ", c.StreamID)
	return httpResp, nil
}
<commit_msg>fixed UA to `swan`<commit_after>package client

import (
	"bytes"
	"encoding/json"
	"fmt"
	"log"
	"net"
	"net/http"
	"time"

	mesosjson "github.com/Dataman-Cloud/swan/mesosproto/json"
)

type Client struct {
	StreamID string
	url      string
	client   *http.Client
}

func New(addr, path string) *Client {
	return &Client{
		url: "http://" + addr + path,
		client: &http.Client{
			Transport: &http.Transport{
				Dial: (&net.Dialer{
					Timeout:   10 * time.Second,
					KeepAlive: 30 * time.Second,
				}).Dial,
			},
		},
	}
}

func (c *Client) Send(payload []byte) (*http.Response, error) {

	httpReq, err := http.NewRequest("POST", c.url, bytes.NewReader(payload))
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/x-protobuf")
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("User-Agent", "swan/0.1")
	if c.StreamID != "" {
		httpReq.Header.Set("Mesos-Stream-Id", c.StreamID)
	}
	//log.Printf("SENDING:%v", httpReq)

	httpResp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("Unable to do request: %s", err)
	}
	if httpResp.Header.Get("Mesos-Stream-Id") != "" {
		c.StreamID = httpResp.Header.Get("Mesos-Stream-Id")
	}
	return httpResp, nil
}

func (c *Client) SendAsJson(call *mesosjson.Call) (*http.Response, error) {
	payload := new(bytes.Buffer)
	if err := json.NewEncoder(payload).Encode(call); err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequest("POST", c.url, payload)
	if err != nil {
		return nil, err
	}

	httpReq.Header.Set("Content-Type", "application/json")
	httpReq.Header.Set("Accept", "application/json")
	httpReq.Header.Set("User-Agent", "swan/0.1")

	httpResp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("Unable to do request: %s", err)
	}
	c.StreamID = httpResp.Header.Get("Mesos-Stream-Id")
	log.Println("Stream-ID: ", c.StreamID)
	return httpResp, nil
}
<|endoftext|>


<commit_before>package godirwalk

import (
	"os"
	"path/filepath"

	"github.com/pkg/errors"
)

// Dirent stores the name and file system mode type of discovered file system
// entries.
type Dirent struct {
	name     string
	modeType os.FileMode
}

// NewDirent returns a newly initialized Dirent structure, or an error. This
// function does not follow symbolic links.
//
// This function is rarely used, as Dirent structures are provided by other
// functions in this library that read and walk directories.
func NewDirent(osPathname string) (*Dirent, error) {
	fi, err = os.Lstat(osPathname)
	if err != nil {
		return nil, errors.Wrap(err, "cannot lstat")
	}
	return &Dirent{
		name:     filepath.Base(osPathname),
		modeType: fi.Mode() & os.ModeType,
	}
}

// Name returns the basename of the file system entry.
func (de Dirent) Name() string { return de.name }

// ModeType returns the mode bits that specify the file system node type. We
// could make our own enum-like data type for encoding the file type, but Go's
// runtime already gives us architecture independent file modes, as discussed in
// `os/types.go`:
//
//    Go's runtime FileMode type has same definition on all systems, so that
//    information about files can be moved from one system to another portably.
func (de Dirent) ModeType() os.FileMode { return de.modeType }

// IsDir returns true if and only if the Dirent represents a file system
// directory. Note that on some operating systems, more than one file mode bit
// may be set for a node. For instance, on Windows, a symbolic link that points
// to a directory will have both the directory and the symbolic link bits set.
func (de Dirent) IsDir() bool { return de.modeType&os.ModeDir != 0 }

// IsRegular returns true if and only if the Dirent represents a regular
// file. That is, it ensures that no mode type bits are set.
func (de Dirent) IsRegular() bool { return de.modeType&os.ModeType == 0 }

// IsSymlink returns true if and only if the Dirent represents a file system
// symbolic link. Note that on some operating systems, more than one file mode
// bit may be set for a node. For instance, on Windows, a symbolic link that
// points to a directory will have both the directory and the symbolic link bits
// set.
func (de Dirent) IsSymlink() bool { return de.modeType&os.ModeSymlink != 0 }

// Dirents represents a slice of Dirent pointers, which are sortable by
// name. This type satisfies the `sort.Interface` interface.
type Dirents []*Dirent

// Len returns the count of Dirent structures in the slice.
func (l Dirents) Len() int { return len(l) }

// Less returns true if and only if the Name of the element specified by the
// first index is lexicographically less than that of the second index.
func (l Dirents) Less(i, j int) bool { return l[i].name < l[j].name }

// Swap exchanges the two Dirent entries specified by the two provided indexes.
func (l Dirents) Swap(i, j int) { l[i], l[j] = l[j], l[i] }

// ReadDirents returns a sortable slice of pointers to Dirent structures, each
// representing the file system name and mode type for one of the immediate
// descendant of the specified directory. If the specified directory is a
// symbolic link, it will be resolved.
//
// If an optional scratch buffer is provided that is at least one page of
// memory, it will be used when reading directory entries from the file system.
//
//    children, err := godirwalk.ReadDirents(osDirname, nil)
//    if err != nil {
//        return nil, errors.Wrap(err, "cannot get list of directory children")
//    }
//    sort.Sort(children)
//    for _, child := range children {
//        fmt.Printf("%s %s\n", child.ModeType, child.Name)
//    }
func ReadDirents(osDirname string, scratchBuffer []byte) (Dirents, error) {
	return readdirents(osDirname, scratchBuffer)
}

// ReadDirnames returns a slice of strings, representing the immediate
// descendants of the specified directory. If the specified directory is a
// symbolic link, it will be resolved.
//
// If an optional scratch buffer is provided that is at least one page of
// memory, it will be used when reading directory entries from the file system.
//
// Note that this function, depending on operating system, may or may not invoke
// the ReadDirents function, in order to prepare the list of immediate
// descendants. Therefore, if your program needs both the names and the file
// system mode types of descendants, it will always be faster to invoke
// ReadDirents directly, rather than calling this function, then looping over
// the results and calling os.Stat for each child.
//
//    children, err := godirwalk.ReadDirnames(osDirname, nil)
//    if err != nil {
//        return nil, errors.Wrap(err, "cannot get list of directory children")
//    }
//    sort.Strings(children)
//    for _, child := range children {
//        fmt.Printf("%s\n", child)
//    }
func ReadDirnames(osDirname string, scratchBuffer []byte) ([]string, error) {
	return readdirnames(osDirname, scratchBuffer)
}
<commit_msg>megacheck && golint && go test<commit_after>package godirwalk

import (
	"os"
	"path/filepath"

	"github.com/pkg/errors"
)

// Dirent stores the name and file system mode type of discovered file system
// entries.
type Dirent struct {
	name     string
	modeType os.FileMode
}

// NewDirent returns a newly initialized Dirent structure, or an error. This
// function does not follow symbolic links.
//
// This function is rarely used, as Dirent structures are provided by other
// functions in this library that read and walk directories.
func NewDirent(osPathname string) (*Dirent, error) {
	fi, err := os.Lstat(osPathname)
	if err != nil {
		return nil, errors.Wrap(err, "cannot lstat")
	}
	return &Dirent{
		name:     filepath.Base(osPathname),
		modeType: fi.Mode() & os.ModeType,
	}, nil
}

// Name returns the basename of the file system entry.
func (de Dirent) Name() string { return de.name }

// ModeType returns the mode bits that specify the file system node type. We
// could make our own enum-like data type for encoding the file type, but Go's
// runtime already gives us architecture independent file modes, as discussed in
// `os/types.go`:
//
//    Go's runtime FileMode type has same definition on all systems, so that
//    information about files can be moved from one system to another portably.
func (de Dirent) ModeType() os.FileMode { return de.modeType }

// IsDir returns true if and only if the Dirent represents a file system
// directory. Note that on some operating systems, more than one file mode bit
// may be set for a node. For instance, on Windows, a symbolic link that points
// to a directory will have both the directory and the symbolic link bits set.
func (de Dirent) IsDir() bool { return de.modeType&os.ModeDir != 0 }

// IsRegular returns true if and only if the Dirent represents a regular
// file. That is, it ensures that no mode type bits are set.
func (de Dirent) IsRegular() bool { return de.modeType&os.ModeType == 0 }

// IsSymlink returns true if and only if the Dirent represents a file system
// symbolic link. Note that on some operating systems, more than one file mode
// bit may be set for a node. For instance, on Windows, a symbolic link that
// points to a directory will have both the directory and the symbolic link bits
// set.
func (de Dirent) IsSymlink() bool { return de.modeType&os.ModeSymlink != 0 }

// Dirents represents a slice of Dirent pointers, which are sortable by
// name. This type satisfies the `sort.Interface` interface.
type Dirents []*Dirent

// Len returns the count of Dirent structures in the slice.
func (l Dirents) Len() int { return len(l) }

// Less returns true if and only if the Name of the element specified by the
// first index is lexicographically less than that of the second index.
func (l Dirents) Less(i, j int) bool { return l[i].name < l[j].name }

// Swap exchanges the two Dirent entries specified by the two provided indexes.
func (l Dirents) Swap(i, j int) { l[i], l[j] = l[j], l[i] }

// ReadDirents returns a sortable slice of pointers to Dirent structures, each
// representing the file system name and mode type for one of the immediate
// descendant of the specified directory. If the specified directory is a
// symbolic link, it will be resolved.
//
// If an optional scratch buffer is provided that is at least one page of
// memory, it will be used when reading directory entries from the file system.
//
//    children, err := godirwalk.ReadDirents(osDirname, nil)
//    if err != nil {
//        return nil, errors.Wrap(err, "cannot get list of directory children")
//    }
//    sort.Sort(children)
//    for _, child := range children {
//        fmt.Printf("%s %s\n", child.ModeType, child.Name)
//    }
func ReadDirents(osDirname string, scratchBuffer []byte) (Dirents, error) {
	return readdirents(osDirname, scratchBuffer)
}

// ReadDirnames returns a slice of strings, representing the immediate
// descendants of the specified directory. If the specified directory is a
// symbolic link, it will be resolved.
//
// If an optional scratch buffer is provided that is at least one page of
// memory, it will be used when reading directory entries from the file system.
//
// Note that this function, depending on operating system, may or may not invoke
// the ReadDirents function, in order to prepare the list of immediate
// descendants. Therefore, if your program needs both the names and the file
// system mode types of descendants, it will always be faster to invoke
// ReadDirents directly, rather than calling this function, then looping over
// the results and calling os.Stat for each child.
//
//    children, err := godirwalk.ReadDirnames(osDirname, nil)
//    if err != nil {
//        return nil, errors.Wrap(err, "cannot get list of directory children")
//    }
//    sort.Strings(children)
//    for _, child := range children {
//        fmt.Printf("%s\n", child)
//    }
func ReadDirnames(osDirname string, scratchBuffer []byte) ([]string, error) {
	return readdirnames(osDirname, scratchBuffer)
}
<|endoftext|>


<commit_before>package session

import (
	"fmt"
	"github.com/hashicorp/consul/api"
	"github.com/myENA/consultant/log"
	"github.com/myENA/consultant/util"
	"github.com/pkg/errors"
	"sync"
	"time"
)

type State uint8

const (
	StateStopped State = iota
	StateRunning
)

const (
	DefaultTTL = "30s"
)

type (
	Update struct {
		ID          string
		Name        string
		LastRenewed time.Time
		Error       error
		State       State
	}

	UpdateFunc func(Update)

	Config struct {
		// Key [suggested]
		//
		// Implementation-specific Key to be placed in session key.
		Key string

		// TTL [optional]
		//
		// Session TTL, defaults to value of DefaultTTL
		TTL string

		// Behavior [optional]
		//
		// Session timeout behavior, defaults to "release"
		Behavior string

		// Log [optional]
		//
		// Logger for this session.  One will be created if value is empty
		Log log.DebugLogger

		// Client [optional]
		//
		// Consul API client, default will be created if not provided
		Client *api.Client

		// UpdateFunc [optional]
		//
		// Callback to be executed after session state change
		UpdateFunc UpdateFunc
	}

	Session struct {
		mu     sync.Mutex
		log    log.DebugLogger
		client *api.Client

		node string
		key  string

		id       string
		name     string
		ttl      time.Duration
		behavior string

		interval    time.Duration
		lastRenewed time.Time

		stop  chan chan struct{}
		state State

		updateFunc UpdateFunc
	}
)

func New(conf *Config) (*Session, error) {
	var key, ttl, behavior string
	var client *api.Client
	var l log.DebugLogger
	var err error

	if conf != nil {
		key = conf.Key
		ttl = conf.TTL
		behavior = conf.Behavior
		l = conf.Log
		client = conf.Client
	}

	if behavior == "" {
		behavior = api.SessionBehaviorRelease
	} else {
		switch behavior {
		case api.SessionBehaviorDelete, api.SessionBehaviorRelease:
		default:
			return nil, fmt.Errorf("\"%s\" is not a valid session behavior", behavior)
		}
	}

	if client == nil {
		client, err = api.NewClient(api.DefaultConfig())
		if err != nil {
			return nil, fmt.Errorf("no Consul api client provided and unable to create: %s", err)
		}
	}

	if l == nil {
		if key == "" {
			l = log.New(fmt.Sprintf("session-%s", util.RandStr(8)))
		} else {
			l = log.New(fmt.Sprintf("session-%s", key))
		}
	}

	if ttl == "" {
		ttl = DefaultTTL
	}

	ttlTD, err := time.ParseDuration(ttl)
	if err != nil {
		return nil, fmt.Errorf("\"%s\" is not valid: %s", ttl, err)
	}

	ttlSeconds := ttlTD.Seconds()
	if ttlSeconds < 10 {
		ttlTD = 10 * time.Second
	} else if ttlSeconds > 86400 {
		ttlTD = 86400 * time.Second
	}

	cs := &Session{
		log:        l,
		client:     client,
		key:        key,
		ttl:        ttlTD,
		behavior:   behavior,
		interval:   time.Duration(ttlTD.Nanoseconds() / 2),
		stop:       make(chan chan struct{}, 1),
		updateFunc: conf.UpdateFunc,
	}

	if cs.node, err = client.Agent().NodeName(); err != nil {
		return nil, fmt.Errorf("unable to determine node: %s", err)
	}

	l.Debugf("Lock interval: %d seconds", int64(ttlTD.Seconds()))
	l.Debugf("Session renew interval: %d seconds", int64(cs.interval.Seconds()))

	return cs, nil
}

func (cs *Session) ID() string {
	cs.mu.Lock()
	sid := cs.id
	cs.mu.Unlock()
	return sid
}

func (cs *Session) Name() string {
	cs.mu.Lock()
	name := cs.name
	cs.mu.Unlock()
	return name
}

func (cs *Session) TTL() time.Duration {
	return cs.ttl
}

func (cs *Session) Key() string {
	return cs.key
}

func (cs *Session) Behavior() string {
	return cs.behavior
}

func (cs *Session) RenewInterval() time.Duration {
	return cs.interval
}

func (cs *Session) LastRenewed() time.Time {
	cs.mu.Lock()
	t := cs.lastRenewed
	cs.mu.Unlock()
	return t
}

func (cs *Session) Running() bool {
	cs.mu.Lock()
	b := cs.state == StateRunning
	cs.mu.Unlock()
	return b
}

func (cs *Session) Run() {
	cs.mu.Lock()
	if cs.state == StateRunning {
		cs.mu.Unlock()
		return
	}

	cs.state = StateRunning

	// try to create session immediately
	if err := cs.create(); err != nil {
		cs.log.Printf(
			"Unable to perform initial session creation, will try again in \"%d\" seconds: %s",
			int64(cs.interval.Seconds()),
			err)
	}

	cs.mu.Unlock()

	go cs.maintain()
}

func (cs *Session) Stop() {
	cs.mu.Lock()
	if cs.state == StateStopped {
		cs.mu.Unlock()
		return
	}
	cs.state = StateStopped
	cs.mu.Unlock()

	stopped := make(chan struct{}, 1)
	cs.stop <- stopped
	<-stopped
	close(stopped)
}

func (cs *Session) State() State {
	cs.mu.Lock()
	s := cs.state
	cs.mu.Unlock()
	return s
}

// sendUpdate will attempt to notify whoever cares that the session has been created / refreshed / errored
// Caller MUST hold lock!
func (cs *Session) sendUpdate(up Update) {
	if cs.updateFunc == nil {
		return
	}
	cs.updateFunc(up)
}

// create will attempt to do just that. Caller MUST hold lock!
func (cs *Session) create() error {
	var name string

	if cs.key == "" {
		name = fmt.Sprintf("%s_%s", cs.node, util.RandStr(12))
	} else {
		name = fmt.Sprintf("%s_%s_%s", cs.key, cs.node, util.RandStr(12))
	}

	cs.log.Debugf("Attempting to create Consul Session \"%s\"...", name)

	se := &api.SessionEntry{
		Name:     name,
		Behavior: cs.behavior,
		TTL:      cs.ttl.String(),
	}

	sid, _, err := cs.client.Session().Create(se, nil)
	if err != nil {
		cs.id = ""
		cs.name = ""
	} else if sid != "" {
		cs.id = sid
		cs.name = name
		cs.lastRenewed = time.Now()
	} else {
		cs.id = ""
		cs.name = ""
		err = errors.New("internal error creating session")
	}

	return err
}

// renew will attempt to do just that.  Caller MUST hold lock!
func (cs *Session) renew() error {
	if cs.id == "" {
		cs.log.Print("Session cannot be renewed as it doesn't exist yet")
		return errors.New("session does not exist yet")
	}

	se, _, err := cs.client.Session().Renew(cs.id, nil)
	if err != nil {
		cs.id = ""
		cs.name = ""
	} else if se != nil {
		cs.id = se.ID
		cs.lastRenewed = time.Now()
	} else {
		cs.id = ""
		cs.name = ""
		err = errors.New("internal error renewing session")
	}

	return err
}

func (cs *Session) maintain() {
	var stopped chan struct{}
	var up Update
	var err error

	intervalTicker := time.NewTicker(cs.interval)

maintaining:
	for {
		select {
		case <-intervalTicker.C:
			cs.mu.Lock()
			if cs.id == "" {
				if err = cs.create(); err != nil {
					cs.log.Printf("Unable to create Consul Session: %s", err)
				}
			} else if err = cs.renew(); err != nil {
				cs.log.Printf("Unable to renew Consul Session: %s", err)
			}

			// create update behind lock
			up = Update{
				ID:          cs.id,
				Name:        cs.name,
				LastRenewed: cs.lastRenewed,
				Error:       err,
				State:       cs.state,
			}

			cs.mu.Unlock()

			//send update after unlock
			cs.sendUpdate(up)

		case stopped = <-cs.stop:
			break maintaining
		}
	}

	cs.log.Debug("Stopping session...")

	intervalTicker.Stop()

	cs.mu.Lock()

	// store most recent session id
	sid := cs.id

	// modify state
	cs.id = ""
	cs.name = ""
	cs.state = StateStopped

	// prepare final update
	up = Update{LastRenewed: cs.lastRenewed, State: StateStopped}

	cs.mu.Unlock()

	// if there was a session id, attempt to destroy it.
	if sid != "" {
		if _, err = cs.client.Session().Destroy(sid, nil); err != nil {
			cs.log.Printf("Unable to destroy session \"%s\": %s", sid, err)
		}
	}

	// just in case...
	if stopped != nil {
		stopped <- struct{}{}
	}

	// send final update asynchronously, don't much care if they're around to receive it.
	go cs.sendUpdate(up)

	cs.log.Print("Session stopped")
}
<commit_msg>updateFunc if conf is nil.<commit_after>package session

import (
	"fmt"
	"github.com/hashicorp/consul/api"
	"github.com/myENA/consultant/log"
	"github.com/myENA/consultant/util"
	"github.com/pkg/errors"
	"sync"
	"time"
)

type State uint8

const (
	StateStopped State = iota
	StateRunning
)

const (
	DefaultTTL = "30s"
)

type (
	Update struct {
		ID          string
		Name        string
		LastRenewed time.Time
		Error       error
		State       State
	}

	UpdateFunc func(Update)

	Config struct {
		// Key [suggested]
		//
		// Implementation-specific Key to be placed in session key.
		Key string

		// TTL [optional]
		//
		// Session TTL, defaults to value of DefaultTTL
		TTL string

		// Behavior [optional]
		//
		// Session timeout behavior, defaults to "release"
		Behavior string

		// Log [optional]
		//
		// Logger for this session.  One will be created if value is empty
		Log log.DebugLogger

		// Client [optional]
		//
		// Consul API client, default will be created if not provided
		Client *api.Client

		// UpdateFunc [optional]
		//
		// Callback to be executed after session state change
		UpdateFunc UpdateFunc
	}

	Session struct {
		mu     sync.Mutex
		log    log.DebugLogger
		client *api.Client

		node string
		key  string

		id       string
		name     string
		ttl      time.Duration
		behavior string

		interval    time.Duration
		lastRenewed time.Time

		stop  chan chan struct{}
		state State

		updateFunc UpdateFunc
	}
)

func New(conf *Config) (*Session, error) {
	var key, ttl, behavior string
	var client *api.Client
	var l log.DebugLogger
	var err error
	var updateFunc UpdateFunc

	if conf != nil {
		key = conf.Key
		ttl = conf.TTL
		behavior = conf.Behavior
		l = conf.Log
		client = conf.Client
		updateFunc = conf.UpdateFunc
	}

	if behavior == "" {
		behavior = api.SessionBehaviorRelease
	} else {
		switch behavior {
		case api.SessionBehaviorDelete, api.SessionBehaviorRelease:
		default:
			return nil, fmt.Errorf("\"%s\" is not a valid session behavior", behavior)
		}
	}

	if client == nil {
		client, err = api.NewClient(api.DefaultConfig())
		if err != nil {
			return nil, fmt.Errorf("no Consul api client provided and unable to create: %s", err)
		}
	}

	if l == nil {
		if key == "" {
			l = log.New(fmt.Sprintf("session-%s", util.RandStr(8)))
		} else {
			l = log.New(fmt.Sprintf("session-%s", key))
		}
	}

	if ttl == "" {
		ttl = DefaultTTL
	}

	ttlTD, err := time.ParseDuration(ttl)
	if err != nil {
		return nil, fmt.Errorf("\"%s\" is not valid: %s", ttl, err)
	}

	ttlSeconds := ttlTD.Seconds()
	if ttlSeconds < 10 {
		ttlTD = 10 * time.Second
	} else if ttlSeconds > 86400 {
		ttlTD = 86400 * time.Second
	}

	cs := &Session{
		log:        l,
		client:     client,
		key:        key,
		ttl:        ttlTD,
		behavior:   behavior,
		interval:   time.Duration(ttlTD.Nanoseconds() / 2),
		stop:       make(chan chan struct{}, 1),
		updateFunc: updateFunc,
	}

	if cs.node, err = client.Agent().NodeName(); err != nil {
		return nil, fmt.Errorf("unable to determine node: %s", err)
	}

	l.Debugf("Lock interval: %d seconds", int64(ttlTD.Seconds()))
	l.Debugf("Session renew interval: %d seconds", int64(cs.interval.Seconds()))

	return cs, nil
}

func (cs *Session) ID() string {
	cs.mu.Lock()
	sid := cs.id
	cs.mu.Unlock()
	return sid
}

func (cs *Session) Name() string {
	cs.mu.Lock()
	name := cs.name
	cs.mu.Unlock()
	return name
}

func (cs *Session) TTL() time.Duration {
	return cs.ttl
}

func (cs *Session) Key() string {
	return cs.key
}

func (cs *Session) Behavior() string {
	return cs.behavior
}

func (cs *Session) RenewInterval() time.Duration {
	return cs.interval
}

func (cs *Session) LastRenewed() time.Time {
	cs.mu.Lock()
	t := cs.lastRenewed
	cs.mu.Unlock()
	return t
}

func (cs *Session) Running() bool {
	cs.mu.Lock()
	b := cs.state == StateRunning
	cs.mu.Unlock()
	return b
}

func (cs *Session) Run() {
	cs.mu.Lock()
	if cs.state == StateRunning {
		cs.mu.Unlock()
		return
	}

	cs.state = StateRunning

	// try to create session immediately
	if err := cs.create(); err != nil {
		cs.log.Printf(
			"Unable to perform initial session creation, will try again in \"%d\" seconds: %s",
			int64(cs.interval.Seconds()),
			err)
	}

	cs.mu.Unlock()

	go cs.maintain()
}

func (cs *Session) Stop() {
	cs.mu.Lock()
	if cs.state == StateStopped {
		cs.mu.Unlock()
		return
	}
	cs.state = StateStopped
	cs.mu.Unlock()

	stopped := make(chan struct{}, 1)
	cs.stop <- stopped
	<-stopped
	close(stopped)
}

func (cs *Session) State() State {
	cs.mu.Lock()
	s := cs.state
	cs.mu.Unlock()
	return s
}

// sendUpdate will attempt to notify whoever cares that the session has been created / refreshed / errored
// Caller MUST hold lock!
func (cs *Session) sendUpdate(up Update) {
	if cs.updateFunc == nil {
		return
	}
	cs.updateFunc(up)
}

// create will attempt to do just that. Caller MUST hold lock!
func (cs *Session) create() error {
	var name string

	if cs.key == "" {
		name = fmt.Sprintf("%s_%s", cs.node, util.RandStr(12))
	} else {
		name = fmt.Sprintf("%s_%s_%s", cs.key, cs.node, util.RandStr(12))
	}

	cs.log.Debugf("Attempting to create Consul Session \"%s\"...", name)

	se := &api.SessionEntry{
		Name:     name,
		Behavior: cs.behavior,
		TTL:      cs.ttl.String(),
	}

	sid, _, err := cs.client.Session().Create(se, nil)
	if err != nil {
		cs.id = ""
		cs.name = ""
	} else if sid != "" {
		cs.id = sid
		cs.name = name
		cs.lastRenewed = time.Now()
	} else {
		cs.id = ""
		cs.name = ""
		err = errors.New("internal error creating session")
	}

	return err
}

// renew will attempt to do just that.  Caller MUST hold lock!
func (cs *Session) renew() error {
	if cs.id == "" {
		cs.log.Print("Session cannot be renewed as it doesn't exist yet")
		return errors.New("session does not exist yet")
	}

	se, _, err := cs.client.Session().Renew(cs.id, nil)
	if err != nil {
		cs.id = ""
		cs.name = ""
	} else if se != nil {
		cs.id = se.ID
		cs.lastRenewed = time.Now()
	} else {
		cs.id = ""
		cs.name = ""
		err = errors.New("internal error renewing session")
	}

	return err
}

func (cs *Session) maintain() {
	var stopped chan struct{}
	var up Update
	var err error

	intervalTicker := time.NewTicker(cs.interval)

maintaining:
	for {
		select {
		case <-intervalTicker.C:
			cs.mu.Lock()
			if cs.id == "" {
				if err = cs.create(); err != nil {
					cs.log.Printf("Unable to create Consul Session: %s", err)
				}
			} else if err = cs.renew(); err != nil {
				cs.log.Printf("Unable to renew Consul Session: %s", err)
			}

			// create update behind lock
			up = Update{
				ID:          cs.id,
				Name:        cs.name,
				LastRenewed: cs.lastRenewed,
				Error:       err,
				State:       cs.state,
			}

			cs.mu.Unlock()

			//send update after unlock
			cs.sendUpdate(up)

		case stopped = <-cs.stop:
			break maintaining
		}
	}

	cs.log.Debug("Stopping session...")

	intervalTicker.Stop()

	cs.mu.Lock()

	// store most recent session id
	sid := cs.id

	// modify state
	cs.id = ""
	cs.name = ""
	cs.state = StateStopped

	// prepare final update
	up = Update{LastRenewed: cs.lastRenewed, State: StateStopped}

	cs.mu.Unlock()

	// if there was a session id, attempt to destroy it.
	if sid != "" {
		if _, err = cs.client.Session().Destroy(sid, nil); err != nil {
			cs.log.Printf("Unable to destroy session \"%s\": %s", sid, err)
		}
	}

	// just in case...
	if stopped != nil {
		stopped <- struct{}{}
	}

	// send final update asynchronously, don't much care if they're around to receive it.
	go cs.sendUpdate(up)

	cs.log.Print("Session stopped")
}
<|endoftext|>