<commit_before>package doozer

import (
	"flag"
	"doozer/client"
	"doozer/lock"
	"doozer/gc"
	"doozer/member"
	"doozer/mon"
	"doozer/paxos"
	"doozer/server"
	"doozer/session"
	"doozer/store"
	"doozer/util"
	"doozer/web"
	"net"
	"os"
	"time"
)

const (
	alpha           = 50
	checkinInterval = 1e9 // ns == 1s
	pulseInterval   = 1e9
)

func Main(clusterName, listenAddr, attachAddr, webAddr string) {
	util.LogWriter = os.Stderr
	logger := util.NewLogger("main")

	prefix := "/d/" + clusterName

	if listenAddr == "" {
		logger.Println("require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	var webListener net.Listener
	if webAddr != "" {
		wl, err := net.Listen("tcp", webAddr)
		if err != nil {
			panic(err)
		}
		webListener = wl
	}

	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		panic(err)
	}

	outs := make(paxos.ChanPutCloserTo)

	cal := make(chan int, 2)

	var cl *client.Client
	self := util.RandId()
	st := store.New()
	seqn := uint64(0)
	if attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn+1, self, listenAddr)
		seqn = addHostname(st, seqn+1, self, os.Getenv("HOSTNAME"))
		seqn = addMember(st, seqn+1, self, listenAddr)
		seqn = claimSlot(st, seqn+1, "1", self)
		seqn = claimLeader(st, seqn+1, self)
		//seqn = claimSlot(st, seqn+1, "2", "")
		//seqn = claimSlot(st, seqn+1, "3", "")
		//seqn = claimSlot(st, seqn+1, "4", "")
		//seqn = claimSlot(st, seqn+1, "5", "")
		seqn = addPing(st, seqn+1, "pong")

		cal <- 1
		cal <- 1

		cl, err = client.Dial(listenAddr)
		if err != nil {
			panic(err)
		}
	} else {
		cl, err = client.Dial(attachAddr)
		if err != nil {
			panic(err)
		}

		path := prefix + "/doozer/info/" + self + "/public-addr"
		_, err = cl.Set(path, listenAddr, store.Clobber)
		if err != nil {
			panic(err)
		}

		path = prefix + "/doozer/info/" + self + "/hostname"
		_, err = cl.Set(path, os.Getenv("HOSTNAME"), store.Clobber)
		if err != nil {
			panic(err)
		}

		var snap string
		seqn, snap, err = cl.Join(self, listenAddr)
		if err != nil {
			panic(err)
		}

		done := make(chan int)
		st.Ops <- store.Op{1, snap}

		go advanceUntil(cl, done)

		go func() {
			st.Sync(seqn + alpha)
			close(done)
			activate(st, self, prefix, cl, cal)
		}()

		// TODO sink needs a way to pick up missing values if there are any
		// gaps in its sequence
	}
	mg := paxos.NewManager(self, seqn, alpha, st, st.Ops, outs)

	if attachAddr == "" {
		// Skip ahead alpha steps so that the registrar can provide a
		// meaningful cluster.
		for i := seqn + 1; i < seqn+alpha; i++ {
			st.Ops <- store.Op{i, store.Nop}
		}
	}

	go func() {
		<-cal
		go lock.Clean(st, mg)
		go session.Clean(st, mg)
		go member.Clean(st, mg)
		go gc.Pulse(self, st.Seqns, cl, pulseInterval)
	}()

	sv := &server.Server{listenAddr, st, mg, self, prefix}

	go func() {
		cas := store.Missing
		for _ = range time.Tick(checkinInterval) {
			_, cas, err = cl.Checkin(self, cas)
			if err != nil {
				logger.Println(err)
			}
		}
	}()

	go func() {
		panic(mon.Monitor(self, prefix, st, cl))
	}()

	go func() {
		panic(sv.Serve(listener, cal))
	}()

	if webListener != nil {
		web.Store = st
		web.ClusterName = clusterName
		go web.Serve(webListener)
	}

	panic(sv.ListenAndServeUdp(outs))
}

func activate(st *store.Store, self, prefix string, c *client.Client, cal chan int) {
	logger := util.NewLogger("activate")
	ch := make(chan store.Event)
	st.GetDirAndWatch("/doozer/slot", ch)
	for ev := range ch {
		// TODO ev.IsEmpty()
		if ev.IsSet() && ev.Body == "" {
			_, err := c.Set(prefix+ev.Path, self, ev.Cas)
			if err != nil {
				logger.Println(err)
				continue
			}
			cal <- 1
			cal <- 1
			close(ch)
		}
	}
}

func advanceUntil(cl *client.Client, done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		cl.Noop()
	}
}

func addPublicAddr(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/public-addr"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addHostname(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/hostname"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addMember(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/members/"+self, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimSlot(st *store.Store, seqn uint64, slot, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/slot/"+slot, self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addPing(st *store.Store, seqn uint64, v string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/ping", v, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}
<commit_msg>sort imports<commit_after>package doozer

import (
	"doozer/client"
	"doozer/gc"
	"doozer/lock"
	"doozer/member"
	"doozer/mon"
	"doozer/paxos"
	"doozer/server"
	"doozer/session"
	"doozer/store"
	"doozer/util"
	"doozer/web"
	"flag"
	"net"
	"os"
	"time"
)

const (
	alpha           = 50
	checkinInterval = 1e9 // ns == 1s
	pulseInterval   = 1e9
)

func Main(clusterName, listenAddr, attachAddr, webAddr string) {
	util.LogWriter = os.Stderr
	logger := util.NewLogger("main")

	prefix := "/d/" + clusterName

	if listenAddr == "" {
		logger.Println("require a listen address")
		flag.Usage()
		os.Exit(1)
	}

	var webListener net.Listener
	if webAddr != "" {
		wl, err := net.Listen("tcp", webAddr)
		if err != nil {
			panic(err)
		}
		webListener = wl
	}

	listener, err := net.Listen("tcp", listenAddr)
	if err != nil {
		panic(err)
	}

	outs := make(paxos.ChanPutCloserTo)

	cal := make(chan int, 2)

	var cl *client.Client
	self := util.RandId()
	st := store.New()
	seqn := uint64(0)
	if attachAddr == "" { // we are the only node in a new cluster
		seqn = addPublicAddr(st, seqn+1, self, listenAddr)
		seqn = addHostname(st, seqn+1, self, os.Getenv("HOSTNAME"))
		seqn = addMember(st, seqn+1, self, listenAddr)
		seqn = claimSlot(st, seqn+1, "1", self)
		seqn = claimLeader(st, seqn+1, self)
		//seqn = claimSlot(st, seqn+1, "2", "")
		//seqn = claimSlot(st, seqn+1, "3", "")
		//seqn = claimSlot(st, seqn+1, "4", "")
		//seqn = claimSlot(st, seqn+1, "5", "")
		seqn = addPing(st, seqn+1, "pong")

		cal <- 1
		cal <- 1

		cl, err = client.Dial(listenAddr)
		if err != nil {
			panic(err)
		}
	} else {
		cl, err = client.Dial(attachAddr)
		if err != nil {
			panic(err)
		}

		path := prefix + "/doozer/info/" + self + "/public-addr"
		_, err = cl.Set(path, listenAddr, store.Clobber)
		if err != nil {
			panic(err)
		}

		path = prefix + "/doozer/info/" + self + "/hostname"
		_, err = cl.Set(path, os.Getenv("HOSTNAME"), store.Clobber)
		if err != nil {
			panic(err)
		}

		var snap string
		seqn, snap, err = cl.Join(self, listenAddr)
		if err != nil {
			panic(err)
		}

		done := make(chan int)
		st.Ops <- store.Op{1, snap}

		go advanceUntil(cl, done)

		go func() {
			st.Sync(seqn + alpha)
			close(done)
			activate(st, self, prefix, cl, cal)
		}()

		// TODO sink needs a way to pick up missing values if there are any
		// gaps in its sequence
	}
	mg := paxos.NewManager(self, seqn, alpha, st, st.Ops, outs)

	if attachAddr == "" {
		// Skip ahead alpha steps so that the registrar can provide a
		// meaningful cluster.
		for i := seqn + 1; i < seqn+alpha; i++ {
			st.Ops <- store.Op{i, store.Nop}
		}
	}

	go func() {
		<-cal
		go lock.Clean(st, mg)
		go session.Clean(st, mg)
		go member.Clean(st, mg)
		go gc.Pulse(self, st.Seqns, cl, pulseInterval)
	}()

	sv := &server.Server{listenAddr, st, mg, self, prefix}

	go func() {
		cas := store.Missing
		for _ = range time.Tick(checkinInterval) {
			_, cas, err = cl.Checkin(self, cas)
			if err != nil {
				logger.Println(err)
			}
		}
	}()

	go func() {
		panic(mon.Monitor(self, prefix, st, cl))
	}()

	go func() {
		panic(sv.Serve(listener, cal))
	}()

	if webListener != nil {
		web.Store = st
		web.ClusterName = clusterName
		go web.Serve(webListener)
	}

	panic(sv.ListenAndServeUdp(outs))
}

func activate(st *store.Store, self, prefix string, c *client.Client, cal chan int) {
	logger := util.NewLogger("activate")
	ch := make(chan store.Event)
	st.GetDirAndWatch("/doozer/slot", ch)
	for ev := range ch {
		// TODO ev.IsEmpty()
		if ev.IsSet() && ev.Body == "" {
			_, err := c.Set(prefix+ev.Path, self, ev.Cas)
			if err != nil {
				logger.Println(err)
				continue
			}
			cal <- 1
			cal <- 1
			close(ch)
		}
	}
}

func advanceUntil(cl *client.Client, done chan int) {
	for _, ok := <-done; !ok; _, ok = <-done {
		cl.Noop()
	}
}

func addPublicAddr(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/public-addr"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addHostname(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	path := "/doozer/info/" + self + "/hostname"
	mx, err := store.EncodeSet(path, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addMember(st *store.Store, seqn uint64, self, addr string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/members/"+self, addr, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimSlot(st *store.Store, seqn uint64, slot, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/slot/"+slot, self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func claimLeader(st *store.Store, seqn uint64, self string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/doozer/leader", self, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}

func addPing(st *store.Store, seqn uint64, v string) uint64 {
	// TODO pull out path as a const
	mx, err := store.EncodeSet("/ping", v, store.Missing)
	if err != nil {
		panic(err)
	}
	st.Ops <- store.Op{seqn, mx}
	return seqn
}
<|endoftext|>


<commit_before>package mockirc

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"
)

type MockIRC struct {
	Port       int
	ServerName string
	Socket     net.Listener
	ctl        chan bool
	running    bool
	events     map[string]*WhenEvent
	io         *bufio.ReadWriter
}

// New will create a new instance of mockirc.
// Returns a pointer to a MockIRC struct.
func New(server string, port int) *MockIRC {
	irc := &MockIRC{}
	if port == 0 {
		irc.Port = 6661
	} else {
		irc.Port = port
	}
	irc.ServerName = server
	irc.events = make(map[string]*WhenEvent)
	return irc
}

// Start will start the "irc" server and listen on the port passed to New.
// Returns a channel of type bool or an error.
func (irc *MockIRC) Start() (chan bool, error) {
	wait := make(chan bool, 1)
	sock, err := net.Listen("tcp", fmt.Sprintf(":%d", irc.Port))
	if err != nil {
		return nil, err
	}
	irc.Socket = sock
	irc.ctl = make(chan bool, 1)
	go func() {
		defer close(wait)

		accept := func() <-chan net.Conn {
			a := make(chan net.Conn)
			go func() {
				for {
					conn, err := irc.Socket.Accept()
					if err != nil {
						log.Println(err)
					}
					if conn != nil {
						a <- conn
					}
				}
			}()
			return a
		}()

		wait <- true

		for {
			select {
			case conn := <-accept:
				go irc.connectionWorker(conn)
			case <-irc.ctl:
				irc.running = false
				return
			}
		}
	}()
	irc.running = true
	return wait, nil
}

// Stop will send the shutdown message on the control channel and stop the server.
// It could return an error.
func (irc *MockIRC) Stop() error {
	irc.ctl <- true
	err := irc.Socket.Close()
	if err != nil {
		return err
	}
	<-time.After(time.Second)
	return nil
}

// connectionWorker will handle incoming connections from Accept.
// Runs in it's own goroutine.
func (irc *MockIRC) connectionWorker(conn net.Conn) {
	irc.io = bufio.NewReadWriter(
		bufio.NewReader(conn),
		bufio.NewWriter(conn))
	defer conn.Close()
	for {
		msg, err := irc.io.ReadString('\n')
		if err != nil && err == io.EOF {
			return
		} else {
			log.Println(err)
		}
		irc.handleMessage(msg)
	}
}

// handleMessage will figure out how to handle messages coming in. It looks at the
// events map to see if it matched anything to send a response.
func (irc *MockIRC) handleMessage(msg string) {
	msg = strings.TrimSuffix(msg, "\r\n")
	var err error
	if val, ok := irc.events[msg]; ok {
		for _, response := range val.responses {
			_, err = irc.io.WriteString(response + "\r\n")
			if err != nil {
				log.Println(err)
				return
			}
			err = irc.io.Flush()
			if err != nil {
				log.Println(err)
				return
			}
		}
	} else {
		log.Printf("Nothing to do for %s\n", msg)
	}
	log.Println(msg)
}

// Send will write the string to the connection.
func (irc *MockIRC) Send(thing string) {
	irc.io.WriteString(thing + "\r\n")
}

type WhenEvent struct {
	event     string
	responses []string
}

// When will take a string that represents an event. This stores the event in a map
// that is checked later when a message comes in over a connection.
// Example use: `mockircserver.When("PING mockirc").Respond(":PONG mockirc")
// Returns the new WhenEvent instance for method chaining.
func (irc *MockIRC) When(event string) *WhenEvent {
	when := &WhenEvent{event: event}
	irc.events[event] = when
	return when
}

// This will add to a list of reponses to send back when an event is matched.
// Returns the new WhenEvent instance for method chaining.
func (when *WhenEvent) Respond(response string) *WhenEvent {
	when.responses = append(when.responses, response)
	return when
}
<commit_msg>stop mockirc gracefully using irc.stop<commit_after>package mockirc

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"strings"
	"time"
)

type MockIRC struct {
	Port       int
	ServerName string
	Socket     net.Listener
	events     map[string]*WhenEvent
	io         *bufio.ReadWriter
	stop       bool
}

// New will create a new instance of mockirc.
// Returns a pointer to a MockIRC struct.
func New(server string, port int) *MockIRC {
	irc := &MockIRC{}
	if port == 0 {
		irc.Port = 6661
	} else {
		irc.Port = port
	}
	irc.ServerName = server
	irc.events = make(map[string]*WhenEvent)
	return irc
}

// Start will start the "irc" server and listen on the port passed to New.
// Returns a channel of type bool or an error.
func (irc *MockIRC) Start() (chan bool, error) {
	wait := make(chan bool, 1)
	sock, err := net.Listen("tcp", fmt.Sprintf(":%d", irc.Port))
	if err != nil {
		return nil, err
	}
	irc.Socket = sock
	go func() {
		defer close(wait)

		accept := func() <-chan net.Conn {
			a := make(chan net.Conn)
			go func() {
				for {
					conn, err := irc.Socket.Accept()
					if err != nil {
						if irc.stop {
							return
						}
						log.Println(err)
						continue
					}
					if conn != nil {
						a <- conn
					}
				}
			}()
			return a
		}()

		wait <- true

		for {
			conn := <-accept
			if irc.stop {
				return
			}
			go irc.connectionWorker(conn)
		}
	}()
	return wait, nil
}

// Stop will send the shutdown message on the control channel and stop the server.
// It could return an error.
func (irc *MockIRC) Stop() error {
	if irc.stop {
		return nil
	}

	irc.stop = true
	err := irc.Socket.Close()
	if err != nil {
		return err
	}
	<-time.After(time.Second)
	return nil
}

// connectionWorker will handle incoming connections from Accept.
// Runs in it's own goroutine.
func (irc *MockIRC) connectionWorker(conn net.Conn) {
	irc.io = bufio.NewReadWriter(
		bufio.NewReader(conn),
		bufio.NewWriter(conn))
	defer conn.Close()
	for {
		msg, err := irc.io.ReadString('\n')
		if err != nil {
			if !irc.stop {
				if err == io.EOF {
					log.Println(err)
				}
			}
			return
		}
		irc.handleMessage(msg)
	}
}

// handleMessage will figure out how to handle messages coming in. It looks at the
// events map to see if it matched anything to send a response.
func (irc *MockIRC) handleMessage(msg string) {
	if !irc.stop {
		msg = strings.TrimSuffix(msg, "\r\n")
		var err error
		if val, ok := irc.events[msg]; ok {
			for _, response := range val.responses {
				_, err = irc.io.WriteString(response + "\r\n")
				if err != nil {
					log.Println(err)
					return
				}
				err = irc.io.Flush()
				if err != nil {
					log.Println(err)
					return
				}
			}
		} else {
			log.Printf("Nothing to do for %s\n", msg)
		}
		log.Println(msg)
	}
}

// Send will write the string to the connection.
func (irc *MockIRC) Send(thing string) {
	if !irc.stop {
		irc.io.WriteString(thing + "\r\n")
	}
}

type WhenEvent struct {
	event     string
	responses []string
}

// When will take a string that represents an event. This stores the event in a map
// that is checked later when a message comes in over a connection.
// Example use: `mockircserver.When("PING mockirc").Respond(":PONG mockirc")
// Returns the new WhenEvent instance for method chaining.
func (irc *MockIRC) When(event string) *WhenEvent {
	when := &WhenEvent{event: event}
	irc.events[event] = when
	return when
}

// This will add to a list of reponses to send back when an event is matched.
// Returns the new WhenEvent instance for method chaining.
func (when *WhenEvent) Respond(response string) *WhenEvent {
	when.responses = append(when.responses, response)
	return when
}
<|endoftext|>


<commit_before>// Copyright 2016 Google LLC
//
// 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.

// TODO(jba): document in CONTRIBUTING.md that service account must be given "Logs Configuration Writer" IAM role for sink tests to pass.
// TODO(jba): [cont] (1) From top left menu, go to IAM & Admin. (2) In Roles dropdown for acct, select Logging > Logs Configuration Writer. (3) Save.
// TODO(jba): Also, cloud-logs@google.com must have Owner permission on the GCS bucket named for the test project.

package logadmin

import (
	"context"
	"log"
	"testing"
	"time"

	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/internal/uid"
	ltest "cloud.google.com/go/logging/internal/testing"
	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

var sinkIDs = uid.NewSpace("GO-CLIENT-TEST-SINK", nil)

const testFilter = ""

var testSinkDestination string

// Called just before TestMain calls m.Run.
// Returns a cleanup function to be called after the tests finish.
func initSinks(ctx context.Context) func() {
	// Create a unique GCS bucket so concurrent tests don't interfere with each other.
	bucketIDs := uid.NewSpace(testProjectID+"-log-sink", nil)
	testBucket := bucketIDs.New()
	testSinkDestination = "storage.googleapis.com/" + testBucket
	var storageClient *storage.Client
	if integrationTest {
		// Create a unique bucket as a sink destination, and give the cloud logging account
		// owner right.
		ts := testutil.TokenSource(ctx, storage.ScopeFullControl)
		var err error
		storageClient, err = storage.NewClient(ctx, option.WithTokenSource(ts))
		if err != nil {
			log.Fatalf("new storage client: %v", err)
		}
		bucket := storageClient.Bucket(testBucket)
		if err := bucket.Create(ctx, testProjectID, nil); err != nil {
			log.Fatalf("creating storage bucket %q: %v", testBucket, err)
		}
		log.Printf("successfully created bucket %s", testBucket)
		if err := bucket.ACL().Set(ctx, "group-cloud-logs@google.com", storage.RoleOwner); err != nil {
			log.Fatalf("setting owner role: %v", err)
		}
	}
	// Clean up from aborted tests.
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Printf("listing sinks: %v", err)
			break
		}
		if sinkIDs.Older(s.ID, 24*time.Hour) {
			client.DeleteSink(ctx, s.ID) // ignore error
		}
	}
	if integrationTest {
		for _, bn := range bucketNames(ctx, storageClient) {
			if bucketIDs.Older(bn, 24*time.Hour) {
				storageClient.Bucket(bn).Delete(ctx) // ignore error
			}
		}
		return func() {
			if err := storageClient.Bucket(testBucket).Delete(ctx); err != nil {
				log.Printf("deleting %q: %v", testBucket, err)
			}
			storageClient.Close()
		}
	}
	return func() {}
}

// Collect the name of all buckets for the test project.
func bucketNames(ctx context.Context, client *storage.Client) []string {
	var names []string
	it := client.Buckets(ctx, testProjectID)
loop:
	for {
		b, err := it.Next()
		switch err {
		case nil:
			names = append(names, b.Name)
		case iterator.Done:
			break loop
		default:
			log.Printf("listing buckets: %v", err)
			break loop
		}
	}
	return names
}

func TestCreateDeleteSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
	}
	got, err := client.CreateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, sink.ID)
	sink.WriterIdentity = ltest.SharedServiceAccount
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}

	if err := client.DeleteSink(ctx, sink.ID); err != nil {
		t.Fatal(err)
	}

	if _, err := client.Sink(ctx, sink.ID); err == nil {
		t.Fatal("got no error, expected one")
	}

	// UniqueWriterIdentity
	sink.ID = sinkIDs.New()
	got, err = client.CreateSinkOpt(ctx, sink, SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, sink.ID)
	// The WriterIdentity should be different.
	if got.WriterIdentity == sink.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
}

func TestUpdateSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, sink.ID)
	got, err := client.UpdateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Updating an existing sink changes it.
	sink.Filter = ""
	sink.IncludeChildren = false
	if _, err := client.UpdateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestUpdateSinkOpt(t *testing.T) {
	ctx := context.Background()
	id := sinkIDs.New()
	origSink := &Sink{
		ID:              id,
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, origSink); err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, id)

	// Updating with empty options is an error.
	_, err := client.UpdateSinkOpt(ctx, &Sink{ID: id, Destination: testSinkDestination}, SinkOptions{})
	if err == nil {
		t.Errorf("got %v, want nil", err)
	}

	// Update selected fields.
	got, err := client.UpdateSinkOpt(ctx, &Sink{ID: id}, SinkOptions{
		UpdateFilter:          true,
		UpdateIncludeChildren: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	want := *origSink
	want.Filter = ""
	want.IncludeChildren = false
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Update writer identity.
	got, err = client.UpdateSinkOpt(ctx, &Sink{ID: id, Filter: "foo"},
		SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	if got.WriterIdentity == want.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
	want.WriterIdentity = got.WriterIdentity
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestListSinks(t *testing.T) {
	ctx := context.Background()
	var sinks []*Sink
	want := map[string]*Sink{}
	for i := 0; i < 4; i++ {
		s := &Sink{
			ID:             sinkIDs.New(),
			Destination:    testSinkDestination,
			Filter:         testFilter,
			WriterIdentity: "serviceAccount:cloud-logs@system.gserviceaccount.com",
		}
		sinks = append(sinks, s)
		want[s.ID] = s
	}
	for _, s := range sinks {
		if _, err := client.CreateSink(ctx, s); err != nil {
			t.Fatalf("Create(%q): %v", s.ID, err)
		}
		defer client.DeleteSink(ctx, s.ID)
	}

	got := map[string]*Sink{}
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		// If tests run simultaneously, we may have more sinks than we
		// created. So only check for our own.
		if _, ok := want[s.ID]; ok {
			got[s.ID] = s
		}
	}
	if !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
}
<commit_msg>logging: Use /dev/null as test filter<commit_after>// Copyright 2016 Google LLC
//
// 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.

// TODO(jba): document in CONTRIBUTING.md that service account must be given "Logs Configuration Writer" IAM role for sink tests to pass.
// TODO(jba): [cont] (1) From top left menu, go to IAM & Admin. (2) In Roles dropdown for acct, select Logging > Logs Configuration Writer. (3) Save.
// TODO(jba): Also, cloud-logs@google.com must have Owner permission on the GCS bucket named for the test project.

package logadmin

import (
	"context"
	"log"
	"testing"
	"time"

	"cloud.google.com/go/internal/testutil"
	"cloud.google.com/go/internal/uid"
	ltest "cloud.google.com/go/logging/internal/testing"
	"cloud.google.com/go/storage"
	"google.golang.org/api/iterator"
	"google.golang.org/api/option"
)

var sinkIDs = uid.NewSpace("GO-CLIENT-TEST-SINK", nil)

// testFilter matches no logs to avoid exporting anything into the sinks.
const testFilter = `textPayload="/dev/null"`

var testSinkDestination string

// Called just before TestMain calls m.Run.
// Returns a cleanup function to be called after the tests finish.
func initSinks(ctx context.Context) func() {
	// Create a unique GCS bucket so concurrent tests don't interfere with each other.
	bucketIDs := uid.NewSpace(testProjectID+"-log-sink", nil)
	testBucket := bucketIDs.New()
	testSinkDestination = "storage.googleapis.com/" + testBucket
	var storageClient *storage.Client
	if integrationTest {
		// Create a unique bucket as a sink destination, and give the cloud logging account
		// owner right.
		ts := testutil.TokenSource(ctx, storage.ScopeFullControl)
		var err error
		storageClient, err = storage.NewClient(ctx, option.WithTokenSource(ts))
		if err != nil {
			log.Fatalf("new storage client: %v", err)
		}
		bucket := storageClient.Bucket(testBucket)
		if err := bucket.Create(ctx, testProjectID, nil); err != nil {
			log.Fatalf("creating storage bucket %q: %v", testBucket, err)
		}
		log.Printf("successfully created bucket %s", testBucket)
		if err := bucket.ACL().Set(ctx, "group-cloud-logs@google.com", storage.RoleOwner); err != nil {
			log.Fatalf("setting owner role: %v", err)
		}
	}
	// Clean up from aborted tests.
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			log.Printf("listing sinks: %v", err)
			break
		}
		if sinkIDs.Older(s.ID, 24*time.Hour) {
			client.DeleteSink(ctx, s.ID) // ignore error
		}
	}
	if integrationTest {
		for _, bn := range bucketNames(ctx, storageClient) {
			if bucketIDs.Older(bn, 24*time.Hour) {
				storageClient.Bucket(bn).Delete(ctx) // ignore error
			}
		}
		return func() {
			if err := storageClient.Bucket(testBucket).Delete(ctx); err != nil {
				log.Printf("deleting %q: %v", testBucket, err)
			}
			storageClient.Close()
		}
	}
	return func() {}
}

// Collect the name of all buckets for the test project.
func bucketNames(ctx context.Context, client *storage.Client) []string {
	var names []string
	it := client.Buckets(ctx, testProjectID)
loop:
	for {
		b, err := it.Next()
		switch err {
		case nil:
			names = append(names, b.Name)
		case iterator.Done:
			break loop
		default:
			log.Printf("listing buckets: %v", err)
			break loop
		}
	}
	return names
}

func TestCreateDeleteSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
	}
	got, err := client.CreateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, sink.ID)
	sink.WriterIdentity = ltest.SharedServiceAccount
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}

	if err := client.DeleteSink(ctx, sink.ID); err != nil {
		t.Fatal(err)
	}

	if _, err := client.Sink(ctx, sink.ID); err == nil {
		t.Fatal("got no error, expected one")
	}

	// UniqueWriterIdentity
	sink.ID = sinkIDs.New()
	got, err = client.CreateSinkOpt(ctx, sink, SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, sink.ID)
	// The WriterIdentity should be different.
	if got.WriterIdentity == sink.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
}

func TestUpdateSink(t *testing.T) {
	ctx := context.Background()
	sink := &Sink{
		ID:              sinkIDs.New(),
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, sink.ID)
	got, err := client.UpdateSink(ctx, sink)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Updating an existing sink changes it.
	sink.Filter = ""
	sink.IncludeChildren = false
	if _, err := client.UpdateSink(ctx, sink); err != nil {
		t.Fatal(err)
	}
	got, err = client.Sink(ctx, sink.ID)
	if err != nil {
		t.Fatal(err)
	}
	if want := sink; !testutil.Equal(got, want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestUpdateSinkOpt(t *testing.T) {
	ctx := context.Background()
	id := sinkIDs.New()
	origSink := &Sink{
		ID:              id,
		Destination:     testSinkDestination,
		Filter:          testFilter,
		IncludeChildren: true,
		WriterIdentity:  ltest.SharedServiceAccount,
	}

	if _, err := client.CreateSink(ctx, origSink); err != nil {
		t.Fatal(err)
	}
	defer client.DeleteSink(ctx, id)

	// Updating with empty options is an error.
	_, err := client.UpdateSinkOpt(ctx, &Sink{ID: id, Destination: testSinkDestination}, SinkOptions{})
	if err == nil {
		t.Errorf("got %v, want nil", err)
	}

	// Update selected fields.
	got, err := client.UpdateSinkOpt(ctx, &Sink{ID: id}, SinkOptions{
		UpdateFilter:          true,
		UpdateIncludeChildren: true,
	})
	if err != nil {
		t.Fatal(err)
	}
	want := *origSink
	want.Filter = ""
	want.IncludeChildren = false
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}

	// Update writer identity.
	got, err = client.UpdateSinkOpt(ctx, &Sink{ID: id, Filter: "foo"},
		SinkOptions{UniqueWriterIdentity: true})
	if err != nil {
		t.Fatal(err)
	}
	if got.WriterIdentity == want.WriterIdentity {
		t.Errorf("got %s, want something different", got.WriterIdentity)
	}
	want.WriterIdentity = got.WriterIdentity
	if !testutil.Equal(got, &want) {
		t.Errorf("got\n%+v\nwant\n%+v", got, want)
	}
}

func TestListSinks(t *testing.T) {
	ctx := context.Background()
	var sinks []*Sink
	want := map[string]*Sink{}
	for i := 0; i < 4; i++ {
		s := &Sink{
			ID:             sinkIDs.New(),
			Destination:    testSinkDestination,
			Filter:         testFilter,
			WriterIdentity: "serviceAccount:cloud-logs@system.gserviceaccount.com",
		}
		sinks = append(sinks, s)
		want[s.ID] = s
	}
	for _, s := range sinks {
		if _, err := client.CreateSink(ctx, s); err != nil {
			t.Fatalf("Create(%q): %v", s.ID, err)
		}
		defer client.DeleteSink(ctx, s.ID)
	}

	got := map[string]*Sink{}
	it := client.Sinks(ctx)
	for {
		s, err := it.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			t.Fatal(err)
		}
		// If tests run simultaneously, we may have more sinks than we
		// created. So only check for our own.
		if _, ok := want[s.ID]; ok {
			got[s.ID] = s
		}
	}
	if !testutil.Equal(got, want) {
		t.Errorf("got %+v, want %+v", got, want)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"math"
	"os"
	"runtime"
	"time"

	"github.com/ChristopherRabotin/smd"
)

func norm(v []float64) float64 {
	return math.Sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])
}

func main() {
	CheckEnvVars()
	runtime.GOMAXPROCS(3)

	start := time.Date(2016, 3, 14, 9, 31, 0, 0, time.UTC) // ExoMars launch date.
	//end := start.Add(time.Duration(-1) * time.Nanosecond)  // Propagate until waypoint reached.
	end := time.Date(2018, 07, 13, 0, 0, 0, 0, time.UTC)

	/* Let's propagate out of Mars at a guessed date of 7 months after launch date from Earth.
	Note that we only output the CSV because we don't need to visualize this.
	*/
	endM := end.Add(time.Duration(4 * 30.5 * 24))
	scMars := SpacecraftFromMars("IM")
	scMars.LogInfo()
	astroM := smd.NewMission(scMars, InitialMarsOrbit(), end, endM, smd.GaussianVOP, smd.Perturbations{}, smd.ExportConfig{Filename: "IM", AsCSV: false, Cosmo: false, Timestamp: false})
	astroM.Propagate()
	// Convert the position to heliocentric.
	astroM.Orbit.ToXCentric(smd.Sun, astroM.CurrentDT)
	target := astroM.Orbit
	//	target := smd.NewOrbitFromOE(226090298.679, 0.088, 26.195, 3.516, 326.494, 278.358, smd.Sun)
	//	fmt.Printf("target orbit: %s\n", target)
	sc := SpacecraftFromEarth("IE", *target)
	sc.LogInfo()
	astro := smd.NewMission(sc, InitialEarthOrbit(), start, end, smd.GaussianVOP, smd.Perturbations{}, smd.ExportConfig{Filename: "IE", AsCSV: true, Cosmo: true, Timestamp: false})
	astro.Propagate()

}

// CheckEnvVars checks that all the environment variables required are set, without checking their value. It will panic if one is missing.
func CheckEnvVars() {
	envvars := []string{"VSOP87", "DATAOUT"}
	for _, envvar := range envvars {
		if os.Getenv(envvar) == "" {
			panic(fmt.Errorf("environment variable `%s` is missing or empty,", envvar))
		}
	}
}
<commit_msg>Keeping fixed Mars and changing departure date<commit_after>package main

import (
	"fmt"
	"math"
	"os"
	"runtime"
	"time"

	"github.com/ChristopherRabotin/smd"
)

func norm(v []float64) float64 {
	return math.Sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])
}

func main() {
	CheckEnvVars()
	runtime.GOMAXPROCS(3)

	//start := time.Date(2016, 3, 14, 9, 31, 0, 0, time.UTC) // ExoMars launch date.
	start := time.Date(2016, 1, 10, 0, 0, 0, 0, time.UTC) // ExoMars launch date.
	estArrival := time.Date(2017, 7, 24, 0, 0, 0, 0, time.UTC)

	/*
		Algo for TOF targeting:
		1. Propagate from random date
		2. Check when hit SOI
		3. Compute duration needed to reach that point
		4. From mean anomaly, compute the time needed for Mars to reach that point (time = DT)
		4a. This leads to knowing how much in advance we are.
		5. Repeat full propagation (including Mars departure) with that DT difference.
	*/

	/* Let's propagate out of Mars at a guessed date of 7 months after launch date from Earth.
	Note that we only output the CSV because we don't need to visualize this.
	*/
	// *WITH THE CURRENT DATE AND TIME*, it takes one month and five days to leave the SOI. So let's propagate only for that time.
	marsEndDT := estArrival.Add(time.Duration(31*24) * time.Hour)
	scMars := SpacecraftFromMars("IM")
	scMars.LogInfo()
	astroM := smd.NewMission(scMars, InitialMarsOrbit(), estArrival, marsEndDT, smd.GaussianVOP, smd.Perturbations{}, smd.ExportConfig{Filename: "IM", AsCSV: false, Cosmo: false, Timestamp: false})
	astroM.Propagate()
	// Convert the position to heliocentric.
	astroM.Orbit.ToXCentric(smd.Sun, astroM.CurrentDT)
	target := astroM.Orbit

	for incr := 0; incr < 9; incr++ {
		actualStart := start.Add(time.Duration(incr * 31 * 24)) // Adding two week periods
		fmt.Printf("===== %s =====\n", actualStart)
		name := fmt.Sprintf("IE-%d%d%d", actualStart.Year(), actualStart.Month(), actualStart.Day())
		sc := SpacecraftFromEarth(name, *target)
		sc.LogInfo()
		// Don't propagate longer than 10 months, it should only take about 8 anyway.
		maxDT := estArrival.Add(time.Duration(10*31*24) * time.Hour)
		astro := smd.NewMission(sc, InitialEarthOrbit(), start, maxDT, smd.GaussianVOP, smd.Perturbations{}, smd.ExportConfig{Filename: name, AsCSV: true, Cosmo: true, Timestamp: false})
		astro.Propagate()
	}

}

// CheckEnvVars checks that all the environment variables required are set, without checking their value. It will panic if one is missing.
func CheckEnvVars() {
	envvars := []string{"VSOP87", "DATAOUT"}
	for _, envvar := range envvars {
		if os.Getenv(envvar) == "" {
			panic(fmt.Errorf("environment variable `%s` is missing or empty,", envvar))
		}
	}
}
<|endoftext|>


<commit_before>package que

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
)

// WorkFunc is a function that performs a Job. If an error is returned, the job
// is reenqueued with exponential backoff.
type WorkFunc func(j *Job) error

// WorkMap is a map of Job names to WorkFuncs that are used to perform Jobs of a
// given type.
type WorkMap map[string]WorkFunc

// Worker is a single worker that pulls jobs off the specified Queue. If no Job
// is found, the Worker will sleep for Interval seconds.
type Worker struct {
	// Interval is the amount of time that this Worker should sleep before trying
	// to find another Job.
	Interval time.Duration

	// Queue is the name of the queue to pull Jobs off of. The default value, "",
	// is usable and is the default for both que-go and the ruby que library.
	Queue string

	c *Client
	m WorkMap

	mu   sync.Mutex
	done bool
	ch   chan struct{}
}

var defaultWakeInterval = 5 * time.Second

func init() {
	if v := os.Getenv("QUE_WAKE_INTERVAL"); v != "" {
		if newInt, err := strconv.Atoi(v); err == nil {
			defaultWakeInterval = time.Duration(newInt) * time.Second
		}
	}
}

// NewWorker returns a Worker that fetches Jobs from the Client and executes
// them using WorkMap. If the type of Job is not registered in the WorkMap, it's
// considered an error and the job is re-enqueued with a backoff.
//
// Workers default to an Interval of 5 seconds, which can be overridden by
// setting the environment variable QUE_WAKE_INTERVAL. The default Queue is the
// nameless queue "", which can be overridden by setting QUE_QUEUE. Either of
// these settings can be changed on the returned Worker before it is started
// with Work().
func NewWorker(c *Client, m WorkMap) *Worker {
	return &Worker{
		Interval: defaultWakeInterval,
		Queue:    os.Getenv("QUE_QUEUE"),
		c:        c,
		m:        m,
		ch:       make(chan struct{}),
	}
}

// Work pulls jobs off the Worker's Queue at its Interval. This function only
// returns after Shutdown() is called, so it should be run in its own goroutine.
func (w *Worker) Work() {
	for {
		select {
		case <-w.ch:
			log.Println("worker done")
			return
		case <-time.After(w.Interval):
			for {
				if didWork := w.WorkOne(); !didWork {
					break // didn't do any work, go back to sleep
				}
			}
		}
	}
}

func (w *Worker) WorkOne() (didWork bool) {
	j, err := w.c.LockJob(w.Queue)
	if err != nil {
		log.Printf("attempting to lock job: %v", err)
		return
	}
	if j == nil {
		return // no job was available
	}
	defer j.Done()
	defer recoverPanic(j)

	didWork = true

	wf, ok := w.m[j.Type]
	if !ok {
		msg := fmt.Sprintf("unknown job type: %q", j.Type)
		log.Println(msg)
		if err = j.Error(msg); err != nil {
			log.Printf("attempting to save error on job %d: %v", j.ID, err)
		}
		return
	}

	if err = wf(j); err != nil {
		j.Error(err.Error())
		return
	}

	if err = j.Delete(); err != nil {
		log.Printf("attempting to delete job %d: %v", j.ID, err)
	}
	log.Printf("event=job_worked job_id=%d job_type=%s", j.ID, j.Type)
	return
}

// Shutdown tells the worker to finish processing its current job and then stop.
// There is currently no timeout for in-progress jobs. This function blocks
// until the Worker has stopped working. It should only be called on an active
// Worker.
func (w *Worker) Shutdown() {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.done {
		return
	}

	log.Println("worker shutting down gracefully...")
	w.ch <- struct{}{}
	w.done = true
	close(w.ch)
}

// recoverPanic tries to handle panics in job execution.
// A stacktrace is stored into Job last_error.
func recoverPanic(j *Job) {
	if r := recover(); r != nil {
		// record an error on the job with panic message and stacktrace
		stackBuf := make([]byte, 1024)
		n := runtime.Stack(stackBuf, false)

		buf := &bytes.Buffer{}
		fmt.Fprintf(buf, "%v\n", r)
		fmt.Fprintln(buf, string(stackBuf[:n]))
		fmt.Fprintln(buf, "[...]")
		stacktrace := buf.String()
		log.Printf("event=panic job_id=%d job_type=%s\n%s", j.ID, j.Type, stacktrace)
		if err := j.Error(stacktrace); err != nil {
			log.Printf("attempting to save error on job %d: %v", j.ID, err)
		}
	}
}

// WorkerPool is a pool of Workers, each working jobs from the queue Queue
// at the specified Interval using the WorkMap.
type WorkerPool struct {
	WorkMap  WorkMap
	Interval time.Duration
	Queue    string

	c       *Client
	workers []*Worker
	mu      sync.Mutex
	done    bool
}

// NewWorkerPool creates a new WorkerPool with count workers using the Client c.
func NewWorkerPool(c *Client, wm WorkMap, count int) *WorkerPool {
	return &WorkerPool{
		c:        c,
		WorkMap:  wm,
		Interval: defaultWakeInterval,
		workers:  make([]*Worker, count),
	}
}

// Start starts all of the Workers in the WorkerPool.
func (w *WorkerPool) Start() {
	w.mu.Lock()
	defer w.mu.Unlock()

	for i := range w.workers {
		w.workers[i] = NewWorker(w.c, w.WorkMap)
		w.workers[i].Interval = w.Interval
		w.workers[i].Queue = w.Queue
		go w.workers[i].Work()
	}
}

// Shutdown sends a Shutdown signal to each of the Workers in the WorkerPool and
// waits for them all to finish shutting down.
func (w *WorkerPool) Shutdown() {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.done {
		return
	}
	var wg sync.WaitGroup
	wg.Add(len(w.workers))

	for _, worker := range w.workers {
		go func(worker *Worker) {
			worker.Shutdown()
			wg.Done()
		}(worker)
	}
	wg.Wait()
	w.done = true
}
<commit_msg>Stop worker when Shutdown() is called, even when jobs are available<commit_after>package que

import (
	"bytes"
	"fmt"
	"log"
	"os"
	"runtime"
	"strconv"
	"sync"
	"time"
)

// WorkFunc is a function that performs a Job. If an error is returned, the job
// is reenqueued with exponential backoff.
type WorkFunc func(j *Job) error

// WorkMap is a map of Job names to WorkFuncs that are used to perform Jobs of a
// given type.
type WorkMap map[string]WorkFunc

// Worker is a single worker that pulls jobs off the specified Queue. If no Job
// is found, the Worker will sleep for Interval seconds.
type Worker struct {
	// Interval is the amount of time that this Worker should sleep before trying
	// to find another Job.
	Interval time.Duration

	// Queue is the name of the queue to pull Jobs off of. The default value, "",
	// is usable and is the default for both que-go and the ruby que library.
	Queue string

	c *Client
	m WorkMap

	mu   sync.Mutex
	done bool
	ch   chan struct{}
}

var defaultWakeInterval = 5 * time.Second

func init() {
	if v := os.Getenv("QUE_WAKE_INTERVAL"); v != "" {
		if newInt, err := strconv.Atoi(v); err == nil {
			defaultWakeInterval = time.Duration(newInt) * time.Second
		}
	}
}

// NewWorker returns a Worker that fetches Jobs from the Client and executes
// them using WorkMap. If the type of Job is not registered in the WorkMap, it's
// considered an error and the job is re-enqueued with a backoff.
//
// Workers default to an Interval of 5 seconds, which can be overridden by
// setting the environment variable QUE_WAKE_INTERVAL. The default Queue is the
// nameless queue "", which can be overridden by setting QUE_QUEUE. Either of
// these settings can be changed on the returned Worker before it is started
// with Work().
func NewWorker(c *Client, m WorkMap) *Worker {
	return &Worker{
		Interval: defaultWakeInterval,
		Queue:    os.Getenv("QUE_QUEUE"),
		c:        c,
		m:        m,
		ch:       make(chan struct{}),
	}
}

// Work pulls jobs off the Worker's Queue at its Interval. This function only
// returns after Shutdown() is called, so it should be run in its own goroutine.
func (w *Worker) Work() {
	defer log.Println("worker done")
	for {
		// Try to work a job
		if w.WorkOne() {
			// Since we just did work, non-blocking check whether we should exit
			select {
			case <-w.ch:
				return
			default:
				// continue in loop
			}
		} else {
			// No work found, block until exit or timer expires
			select {
			case <-w.ch:
				return
			case <-time.After(w.Interval):
				// continue in loop
			}
		}
	}
}

func (w *Worker) WorkOne() (didWork bool) {
	j, err := w.c.LockJob(w.Queue)
	if err != nil {
		log.Printf("attempting to lock job: %v", err)
		return
	}
	if j == nil {
		return // no job was available
	}
	defer j.Done()
	defer recoverPanic(j)

	didWork = true

	wf, ok := w.m[j.Type]
	if !ok {
		msg := fmt.Sprintf("unknown job type: %q", j.Type)
		log.Println(msg)
		if err = j.Error(msg); err != nil {
			log.Printf("attempting to save error on job %d: %v", j.ID, err)
		}
		return
	}

	if err = wf(j); err != nil {
		j.Error(err.Error())
		return
	}

	if err = j.Delete(); err != nil {
		log.Printf("attempting to delete job %d: %v", j.ID, err)
	}
	log.Printf("event=job_worked job_id=%d job_type=%s", j.ID, j.Type)
	return
}

// Shutdown tells the worker to finish processing its current job and then stop.
// There is currently no timeout for in-progress jobs. This function blocks
// until the Worker has stopped working. It should only be called on an active
// Worker.
func (w *Worker) Shutdown() {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.done {
		return
	}

	log.Println("worker shutting down gracefully...")
	w.ch <- struct{}{}
	w.done = true
	close(w.ch)
}

// recoverPanic tries to handle panics in job execution.
// A stacktrace is stored into Job last_error.
func recoverPanic(j *Job) {
	if r := recover(); r != nil {
		// record an error on the job with panic message and stacktrace
		stackBuf := make([]byte, 1024)
		n := runtime.Stack(stackBuf, false)

		buf := &bytes.Buffer{}
		fmt.Fprintf(buf, "%v\n", r)
		fmt.Fprintln(buf, string(stackBuf[:n]))
		fmt.Fprintln(buf, "[...]")
		stacktrace := buf.String()
		log.Printf("event=panic job_id=%d job_type=%s\n%s", j.ID, j.Type, stacktrace)
		if err := j.Error(stacktrace); err != nil {
			log.Printf("attempting to save error on job %d: %v", j.ID, err)
		}
	}
}

// WorkerPool is a pool of Workers, each working jobs from the queue Queue
// at the specified Interval using the WorkMap.
type WorkerPool struct {
	WorkMap  WorkMap
	Interval time.Duration
	Queue    string

	c       *Client
	workers []*Worker
	mu      sync.Mutex
	done    bool
}

// NewWorkerPool creates a new WorkerPool with count workers using the Client c.
func NewWorkerPool(c *Client, wm WorkMap, count int) *WorkerPool {
	return &WorkerPool{
		c:        c,
		WorkMap:  wm,
		Interval: defaultWakeInterval,
		workers:  make([]*Worker, count),
	}
}

// Start starts all of the Workers in the WorkerPool.
func (w *WorkerPool) Start() {
	w.mu.Lock()
	defer w.mu.Unlock()

	for i := range w.workers {
		w.workers[i] = NewWorker(w.c, w.WorkMap)
		w.workers[i].Interval = w.Interval
		w.workers[i].Queue = w.Queue
		go w.workers[i].Work()
	}
}

// Shutdown sends a Shutdown signal to each of the Workers in the WorkerPool and
// waits for them all to finish shutting down.
func (w *WorkerPool) Shutdown() {
	w.mu.Lock()
	defer w.mu.Unlock()

	if w.done {
		return
	}
	var wg sync.WaitGroup
	wg.Add(len(w.workers))

	for _, worker := range w.workers {
		go func(worker *Worker) {
			worker.Shutdown()
			wg.Done()
		}(worker)
	}
	wg.Wait()
	w.done = true
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var exp map[string](*regexp.Regexp) = map[string](*regexp.Regexp){}

type ServiceEnvs map[string]string

type Services map[string]ServiceEnvs

type MuninMetric struct {
	Label string
	Type  string
	Draw  string
	Value string
}

type MuninPlugin struct {
	PluginPath    string
	PluginConfDir string
	GraphTitle    string
	GraphName     string
	MuninMetrics  map[string](*MuninMetric)
}

func getExp(expstr string) *(regexp.Regexp) {
	if exp[expstr] == nil {
		exp[expstr] = regexp.MustCompile(expstr)
	}
	return exp[expstr]
}

func getEnvSettingsReader(s *Services, plg string, reader io.Reader) {
	var service string
	scanner := bufio.NewScanner(reader)

	for scanner.Scan() {
		line := strings.TrimRight(scanner.Text(), "\n")
		line = getExp("([^\\\\])#.*?$").ReplaceAllString(line, "$1")
		line = getExp("\\\\#").ReplaceAllString(line, "#")
		line = getExp("(^\\s+|\\s+$)").ReplaceAllString(line, "")

		service_m := getExp("^\\s*\\[([^\\]]*?)(\\*?)\\]\\s*$").FindStringSubmatch(line)
		if service_m != nil {
			if service_m[2] == "" && plg == service_m[1] { // perfect match
				service = service_m[1]
			} else if service_m[2] == "*" && len(service_m[1]) <= len(plg) && plg[0:len(service_m[1])] == service_m[1] { // left-hand match
				service = service_m[1] + service_m[2]
			} else {
				service = ""
			}
			continue
		}

		if service == "" {
			continue
		}

		env_m := getExp("^\\s*env\\.(\\w+)\\s+(.+)$").FindStringSubmatch(line)
		if env_m != nil && service != "" {
			if (*s)[service] == nil {
				(*s)[service] = ServiceEnvs{}
			}
			(*s)[service][env_m[1]] = env_m[2]
		}
	}
}

func getEnvSettingsFile(s *Services, plg string, file string) {
	fp, err := os.Open(file)
	if err != nil {
		return
	}
	defer fp.Close()

	getEnvSettingsReader(s, plg, fp)
}

func compileEnvPairs(s *Services, plg string) *map[string]string {
	servenvs := *s

	// ordered services
	srvs := make([]string, 0, len(servenvs))
	for k := range servenvs {
		if k == plg {
			continue
		}
		srvs = append(srvs, k)
	}
	sort.Strings(srvs)
	if servenvs[plg] != nil {
		srvs = append(srvs, plg)
	}

	// apply envs
	envs := make(map[string]string)
	for _, srv := range srvs {
		for k, v := range servenvs[srv] {
			envs[k] = v
		}
	}

	return &envs
}

func setPluginEnvironments(plg string, confdir string) {
	files, err := ioutil.ReadDir(confdir)
	if err != nil {
		log.Fatalln(err)
	}

	filenames := make([]string, 0, len(files))
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		filenames = append(filenames, f.Name())
	}
	sort.Strings(filenames)

	servenvs := make(Services)
	for _, f := range filenames {
		getEnvSettingsFile(&servenvs, plg, path.Join(confdir, f))
	}

	for k, v := range *compileEnvPairs(&servenvs, plg) {
		os.Setenv(k, v)
	}
}

func parsePluginConfig(str string, m *map[string](*MuninMetric), title *string) {
	for _, line := range strings.Split(str, "\n") {
		line_m := getExp("^([^ ]+) +(.*)$").FindStringSubmatch(line)
		if line_m == nil {
			continue
		}

		if line_m[1] == "graph_title" {
			(*title) = line_m[2]
			continue
		}

		key_m := getExp("^([^\\.]+)\\.(.*)$").FindStringSubmatch(line_m[1])
		if key_m == nil {
			continue
		}

		if (*m)[key_m[1]] == nil {
			(*m)[key_m[1]] = &(MuninMetric{})
		}
		met := (*m)[key_m[1]]

		switch key_m[2] {
		case "label":
			met.Label = line_m[2]
		case "type":
			met.Type = line_m[2]
		case "draw":
			met.Draw = line_m[2]
		}
	}
}

func parsePluginVals(str string, m *map[string](*MuninMetric)) {
	for _, line := range strings.Split(str, "\n") {
		line_m := getExp("^([^ ]+) +(.*)$").FindStringSubmatch(line)
		if line_m == nil {
			continue
		}

		key_m := getExp("^([^\\.]+)\\.(.*)$").FindStringSubmatch(line_m[1])
		if key_m == nil {
			continue
		}

		if (*m)[key_m[1]] == nil {
			(*m)[key_m[1]] = &(MuninMetric{})
		}
		met := (*m)[key_m[1]]

		switch key_m[2] {
		case "value":
			met.Value = line_m[2]
		}
	}
}

func removeUselessMetrics(m *map[string](*MuninMetric)) {
	// remove metrics which have an empty Value
	for name, mmet := range *m {
		if mmet.Value == "" {
			delete(*m, name)
		}
	}
}

func (p *MuninPlugin) Prepare() error {
	var err error

	if p.PluginConfDir != "" {
		setPluginEnvironments(path.Base(p.PluginPath), p.PluginConfDir)
	}

	out_config, err := exec.Command(p.PluginPath, "config").Output()
	if err != nil {
		return errors.New(fmt.Sprintf("%s: %s", err, out_config))
	}

	out_vals, err := exec.Command(p.PluginPath).Output()
	if err != nil {
		return errors.New(fmt.Sprintf("%s: %s", err, out_vals))
	}

	p.MuninMetrics = make(map[string](*MuninMetric))
	parsePluginConfig(string(out_config), &p.MuninMetrics, &p.GraphTitle)
	parsePluginVals(string(out_vals), &p.MuninMetrics)
	removeUselessMetrics(&p.MuninMetrics)

	return nil
}

func (p MuninPlugin) FetchMetrics() (map[string]float64, error) {
	stat := make(map[string]float64, len(p.MuninMetrics))
	for name, mmet := range p.MuninMetrics {
		parsed, err := strconv.ParseFloat(mmet.Value, 64)
		if err != nil {
			log.Printf("Failed to parse value of %s: %s", name, err)
			continue
		}

		stat[name] = parsed
	}

	return stat, nil
}

func (p MuninPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := make([](mp.Metrics), 0, len(p.MuninMetrics))
	for name, mmet := range p.MuninMetrics {
		met := mp.Metrics{Name: name}
		if mmet.Label != "" {
			met.Label = mmet.Label
		}
		if mmet.Draw == "STACK" {
			met.Stacked = true
		}
		switch mmet.Type {
		case "COUNTER", "DERIVE", "ABSOLUTE":
			met.Diff = true
		}

		metrics = append(metrics, met)
	}

	return map[string](mp.Graphs){p.GraphName: mp.Graphs{
		Label:   p.GraphTitle,
		Unit:    "float",
		Metrics: metrics,
	}}
}

func main() {
	optPluginPath := flag.String("plugin", "", "Munin plugin path")
	optPluginConfDir := flag.String("plugin-conf-d", "", "Munin plugin-conf.d path")
	optGraphName := flag.String("name", "", "Graph name")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var munin MuninPlugin
	if *optPluginPath == "" {
		log.Fatalln("Munin plugin path is required")
	}
	munin.PluginPath = *optPluginPath
	munin.PluginConfDir = *optPluginConfDir
	if *optGraphName == "" {
		munin.GraphName = "munin." + path.Base(munin.PluginPath)
	} else {
		munin.GraphName = *optGraphName
	}

	err := munin.Prepare()
	if err != nil {
		log.Fatalln(err)
	}

	helper := mp.NewMackerelPlugin(munin)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-munin-%s", munin.GraphName)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>avoid empty Label<commit_after>package main

import (
	"bufio"
	"errors"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"io"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"path"
	"regexp"
	"sort"
	"strconv"
	"strings"
)

var exp map[string](*regexp.Regexp) = map[string](*regexp.Regexp){}

type ServiceEnvs map[string]string

type Services map[string]ServiceEnvs

type MuninMetric struct {
	Label string
	Type  string
	Draw  string
	Value string
}

type MuninPlugin struct {
	PluginPath    string
	PluginConfDir string
	GraphTitle    string
	GraphName     string
	MuninMetrics  map[string](*MuninMetric)
}

func getExp(expstr string) *(regexp.Regexp) {
	if exp[expstr] == nil {
		exp[expstr] = regexp.MustCompile(expstr)
	}
	return exp[expstr]
}

func getEnvSettingsReader(s *Services, plg string, reader io.Reader) {
	var service string
	scanner := bufio.NewScanner(reader)

	for scanner.Scan() {
		line := strings.TrimRight(scanner.Text(), "\n")
		line = getExp("([^\\\\])#.*?$").ReplaceAllString(line, "$1")
		line = getExp("\\\\#").ReplaceAllString(line, "#")
		line = getExp("(^\\s+|\\s+$)").ReplaceAllString(line, "")

		service_m := getExp("^\\s*\\[([^\\]]*?)(\\*?)\\]\\s*$").FindStringSubmatch(line)
		if service_m != nil {
			if service_m[2] == "" && plg == service_m[1] { // perfect match
				service = service_m[1]
			} else if service_m[2] == "*" && len(service_m[1]) <= len(plg) && plg[0:len(service_m[1])] == service_m[1] { // left-hand match
				service = service_m[1] + service_m[2]
			} else {
				service = ""
			}
			continue
		}

		if service == "" {
			continue
		}

		env_m := getExp("^\\s*env\\.(\\w+)\\s+(.+)$").FindStringSubmatch(line)
		if env_m != nil && service != "" {
			if (*s)[service] == nil {
				(*s)[service] = ServiceEnvs{}
			}
			(*s)[service][env_m[1]] = env_m[2]
		}
	}
}

func getEnvSettingsFile(s *Services, plg string, file string) {
	fp, err := os.Open(file)
	if err != nil {
		return
	}
	defer fp.Close()

	getEnvSettingsReader(s, plg, fp)
}

func compileEnvPairs(s *Services, plg string) *map[string]string {
	servenvs := *s

	// ordered services
	srvs := make([]string, 0, len(servenvs))
	for k := range servenvs {
		if k == plg {
			continue
		}
		srvs = append(srvs, k)
	}
	sort.Strings(srvs)
	if servenvs[plg] != nil {
		srvs = append(srvs, plg)
	}

	// apply envs
	envs := make(map[string]string)
	for _, srv := range srvs {
		for k, v := range servenvs[srv] {
			envs[k] = v
		}
	}

	return &envs
}

func setPluginEnvironments(plg string, confdir string) {
	files, err := ioutil.ReadDir(confdir)
	if err != nil {
		log.Fatalln(err)
	}

	filenames := make([]string, 0, len(files))
	for _, f := range files {
		if f.IsDir() {
			continue
		}
		filenames = append(filenames, f.Name())
	}
	sort.Strings(filenames)

	servenvs := make(Services)
	for _, f := range filenames {
		getEnvSettingsFile(&servenvs, plg, path.Join(confdir, f))
	}

	for k, v := range *compileEnvPairs(&servenvs, plg) {
		os.Setenv(k, v)
	}
}

func parsePluginConfig(str string, m *map[string](*MuninMetric), title *string) {
	for _, line := range strings.Split(str, "\n") {
		line_m := getExp("^([^ ]+) +(.*)$").FindStringSubmatch(line)
		if line_m == nil {
			continue
		}

		if line_m[1] == "graph_title" {
			(*title) = line_m[2]
			continue
		}

		key_m := getExp("^([^\\.]+)\\.(.*)$").FindStringSubmatch(line_m[1])
		if key_m == nil {
			continue
		}

		if (*m)[key_m[1]] == nil {
			(*m)[key_m[1]] = &(MuninMetric{})
		}
		met := (*m)[key_m[1]]

		switch key_m[2] {
		case "label":
			met.Label = line_m[2]
		case "type":
			met.Type = line_m[2]
		case "draw":
			met.Draw = line_m[2]
		}
	}
}

func parsePluginVals(str string, m *map[string](*MuninMetric)) {
	for _, line := range strings.Split(str, "\n") {
		line_m := getExp("^([^ ]+) +(.*)$").FindStringSubmatch(line)
		if line_m == nil {
			continue
		}

		key_m := getExp("^([^\\.]+)\\.(.*)$").FindStringSubmatch(line_m[1])
		if key_m == nil {
			continue
		}

		if (*m)[key_m[1]] == nil {
			(*m)[key_m[1]] = &(MuninMetric{})
		}
		met := (*m)[key_m[1]]

		switch key_m[2] {
		case "value":
			met.Value = line_m[2]
		}
	}
}

func removeUselessMetrics(m *map[string](*MuninMetric)) {
	// remove metrics which have an empty Value
	for name, mmet := range *m {
		if mmet.Value == "" {
			delete(*m, name)
		}
	}
}

func (p *MuninPlugin) Prepare() error {
	var err error

	if p.PluginConfDir != "" {
		setPluginEnvironments(path.Base(p.PluginPath), p.PluginConfDir)
	}

	out_config, err := exec.Command(p.PluginPath, "config").Output()
	if err != nil {
		return errors.New(fmt.Sprintf("%s: %s", err, out_config))
	}

	out_vals, err := exec.Command(p.PluginPath).Output()
	if err != nil {
		return errors.New(fmt.Sprintf("%s: %s", err, out_vals))
	}

	p.MuninMetrics = make(map[string](*MuninMetric))
	parsePluginConfig(string(out_config), &p.MuninMetrics, &p.GraphTitle)
	parsePluginVals(string(out_vals), &p.MuninMetrics)
	removeUselessMetrics(&p.MuninMetrics)

	return nil
}

func (p MuninPlugin) FetchMetrics() (map[string]float64, error) {
	stat := make(map[string]float64, len(p.MuninMetrics))
	for name, mmet := range p.MuninMetrics {
		parsed, err := strconv.ParseFloat(mmet.Value, 64)
		if err != nil {
			log.Printf("Failed to parse value of %s: %s", name, err)
			continue
		}

		stat[name] = parsed
	}

	return stat, nil
}

func (p MuninPlugin) GraphDefinition() map[string](mp.Graphs) {
	metrics := make([](mp.Metrics), 0, len(p.MuninMetrics))
	for name, mmet := range p.MuninMetrics {
		met := mp.Metrics{Name: name}
		if mmet.Label == "" {
			met.Label = name
		} else {
			met.Label = mmet.Label
		}
		if mmet.Draw == "STACK" {
			met.Stacked = true
		}
		switch mmet.Type {
		case "COUNTER", "DERIVE", "ABSOLUTE":
			met.Diff = true
		}

		metrics = append(metrics, met)
	}

	return map[string](mp.Graphs){p.GraphName: mp.Graphs{
		Label:   p.GraphTitle,
		Unit:    "float",
		Metrics: metrics,
	}}
}

func main() {
	optPluginPath := flag.String("plugin", "", "Munin plugin path")
	optPluginConfDir := flag.String("plugin-conf-d", "", "Munin plugin-conf.d path")
	optGraphName := flag.String("name", "", "Graph name")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var munin MuninPlugin
	if *optPluginPath == "" {
		log.Fatalln("Munin plugin path is required")
	}
	munin.PluginPath = *optPluginPath
	munin.PluginConfDir = *optPluginConfDir
	if *optGraphName == "" {
		munin.GraphName = "munin." + path.Base(munin.PluginPath)
	} else {
		munin.GraphName = *optGraphName
	}

	err := munin.Prepare()
	if err != nil {
		log.Fatalln(err)
	}

	helper := mp.NewMackerelPlugin(munin)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-munin-%s", munin.GraphName)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>// Copyright 2013 The Prometheus 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 model

import (
	"fmt"
	"regexp"
	"strings"
)

const (
	// ExporterLabelPrefix is the label name prefix to prepend if a
	// synthetic label is already present in the exported metrics.
	ExporterLabelPrefix LabelName = "exporter_"

	// MetricNameLabel is the label name indicating the metric name of a
	// timeseries.
	MetricNameLabel LabelName = "__name__"

	// AddressLabel is the name of the label that holds the address of
	// a scrape target.
	AddressLabel LabelName = "__address__"

	// MetricsPathLabel is the name of the label that holds the path on which to
	// scrape a target.
	MetricsPathLabel LabelName = "__metrics_path__"

	// ReservedLabelPrefix is a prefix which is not legal in user-supplied
	// label names.
	ReservedLabelPrefix = "__"

	// MetaLabelPrefix is a prefix for labels that provide meta information.
	// Labels with this prefix are used for intermediate label processing and
	// will not be attached to time series.
	MetaLabelPrefix = "__meta_"

	// JobLabel is the label name indicating the job from which a timeseries
	// was scraped.
	JobLabel LabelName = "job"

	// InstanceLabel is the label name used for the instance label.
	InstanceLabel LabelName = "instance"

	// BucketLabel is used for the label that defines the upper bound of a
	// bucket of a histogram ("le" -> "less or equal").
	BucketLabel = "le"

	// QuantileLabel is used for the label that defines the quantile in a
	// summary.
	QuantileLabel = "quantile"
)

// A LabelName is a key for a LabelSet or Metric.  It has a value associated
// therewith.
type LabelName string

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (ln *LabelName) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	if err := unmarshal(&s); err != nil {
		return err
	}
	if !labelNameRE.MatchString(s) {
		return fmt.Errorf("%q is not a valid label name", s)
	}
	*ln = LabelName(s)
	return nil
}

// LabelNames is a sortable LabelName slice. In implements sort.Interface.
type LabelNames []LabelName

func (l LabelNames) Len() int {
	return len(l)
}

func (l LabelNames) Less(i, j int) bool {
	return l[i] < l[j]
}

func (l LabelNames) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l LabelNames) String() string {
	labelStrings := make([]string, 0, len(l))
	for _, label := range l {
		labelStrings = append(labelStrings, string(label))
	}
	return strings.Join(labelStrings, ", ")
}
<commit_msg>Add LabelName regular expression<commit_after>// Copyright 2013 The Prometheus 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 model

import (
	"fmt"
	"regexp"
	"strings"
)

const (
	// ExporterLabelPrefix is the label name prefix to prepend if a
	// synthetic label is already present in the exported metrics.
	ExporterLabelPrefix LabelName = "exporter_"

	// MetricNameLabel is the label name indicating the metric name of a
	// timeseries.
	MetricNameLabel LabelName = "__name__"

	// AddressLabel is the name of the label that holds the address of
	// a scrape target.
	AddressLabel LabelName = "__address__"

	// MetricsPathLabel is the name of the label that holds the path on which to
	// scrape a target.
	MetricsPathLabel LabelName = "__metrics_path__"

	// ReservedLabelPrefix is a prefix which is not legal in user-supplied
	// label names.
	ReservedLabelPrefix = "__"

	// MetaLabelPrefix is a prefix for labels that provide meta information.
	// Labels with this prefix are used for intermediate label processing and
	// will not be attached to time series.
	MetaLabelPrefix = "__meta_"

	// JobLabel is the label name indicating the job from which a timeseries
	// was scraped.
	JobLabel LabelName = "job"

	// InstanceLabel is the label name used for the instance label.
	InstanceLabel LabelName = "instance"

	// BucketLabel is used for the label that defines the upper bound of a
	// bucket of a histogram ("le" -> "less or equal").
	BucketLabel = "le"

	// QuantileLabel is used for the label that defines the quantile in a
	// summary.
	QuantileLabel = "quantile"
)

var labelNameRE = regexp.MustCompile("^[a-zA-Z_][a-zA-Z0-9_]*$")

// A LabelName is a key for a LabelSet or Metric.  It has a value associated
// therewith.
type LabelName string

// UnmarshalYAML implements the yaml.Unmarshaler interface.
func (ln *LabelName) UnmarshalYAML(unmarshal func(interface{}) error) error {
	var s string
	if err := unmarshal(&s); err != nil {
		return err
	}
	if !labelNameRE.MatchString(s) {
		return fmt.Errorf("%q is not a valid label name", s)
	}
	*ln = LabelName(s)
	return nil
}

// LabelNames is a sortable LabelName slice. In implements sort.Interface.
type LabelNames []LabelName

func (l LabelNames) Len() int {
	return len(l)
}

func (l LabelNames) Less(i, j int) bool {
	return l[i] < l[j]
}

func (l LabelNames) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l LabelNames) String() string {
	labelStrings := make([]string, 0, len(l))
	for _, label := range l {
		labelStrings = append(labelStrings, string(label))
	}
	return strings.Join(labelStrings, ", ")
}
<|endoftext|>


<commit_before>package srslog

import (
	"crypto/tls"
	"strings"
	"sync"
)

// A Writer is a connection to a syslog server.
type Writer struct {
	priority  Priority
	tag       string
	hostname  string
	network   string
	raddr     string
	tlsConfig *tls.Config
	framer    Framer
	formatter Formatter

	mu   sync.RWMutex // guards conn
	conn serverConn
}

// getConn provides access to the internal conn, protected by a mutex. The
// conn is threadsafe, so it can be used while unlocked, but we want to avoid
// race conditions on grabbing a reference to it.
func (w *Writer) getConn() serverConn {
	w.mu.RLock()
	defer w.mu.RUnlock()
	return w.conn
}

// setConn updates the internal conn, protected by a mutex.
func (w *Writer) setConn(c serverConn) {
	w.mu.Lock()
	w.conn = c
	w.mu.Unlock()
}

// connect makes a connection to the syslog server.
func (w *Writer) connect() (serverConn, error) {
	conn := w.getConn()
	if conn != nil {
		// ignore err from close, it makes sense to continue anyway
		conn.close()
		w.setConn(nil)
	}

	var hostname string
	var err error
	dialer := w.getDialer()
	conn, hostname, err = dialer.Call()
	if err == nil {
		w.setConn(conn)
		w.hostname = hostname

		return conn, nil
	} else {
		return nil, err
	}
}

// SetFormatter changes the formatter function for subsequent messages.
func (w *Writer) SetFormatter(f Formatter) {
	w.formatter = f
}

// SetFramer changes the framer function for subsequent messages.
func (w *Writer) SetFramer(f Framer) {
	w.framer = f
}

// Write sends a log message to the syslog daemon using the default priority
// passed into `srslog.New` or the `srslog.Dial*` functions.
func (w *Writer) Write(b []byte) (int, error) {
	return w.writeAndRetry(w.priority, string(b))
}

// WriteWithPriority sends a log message with a custom priority
func (w *Writer) WriteWithPriority(p Priority, b []byte) (int, error) {
	return w.writeAndRetry(p, string(b))
}

// Close closes a connection to the syslog daemon.
func (w *Writer) Close() error {
	conn := w.getConn()
	if conn != nil {
		err := conn.close()
		w.setConn(nil)
		return err
	}
	return nil
}

// Emerg logs a message with severity LOG_EMERG; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Emerg(m string) (err error) {
	_, err = w.writeAndRetry(LOG_EMERG, m)
	return err
}

// Alert logs a message with severity LOG_ALERT; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Alert(m string) (err error) {
	_, err = w.writeAndRetry(LOG_ALERT, m)
	return err
}

// Crit logs a message with severity LOG_CRIT; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Crit(m string) (err error) {
	_, err = w.writeAndRetry(LOG_CRIT, m)
	return err
}

// Err logs a message with severity LOG_ERR; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Err(m string) (err error) {
	_, err = w.writeAndRetry(LOG_ERR, m)
	return err
}

// Warning logs a message with severity LOG_WARNING; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Warning(m string) (err error) {
	_, err = w.writeAndRetry(LOG_WARNING, m)
	return err
}

// Notice logs a message with severity LOG_NOTICE; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Notice(m string) (err error) {
	_, err = w.writeAndRetry(LOG_NOTICE, m)
	return err
}

// Info logs a message with severity LOG_INFO; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Info(m string) (err error) {
	_, err = w.writeAndRetry(LOG_INFO, m)
	return err
}

// Debug logs a message with severity LOG_DEBUG; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Debug(m string) (err error) {
	_, err = w.writeAndRetry(LOG_DEBUG, m)
	return err
}

func (w *Writer) writeAndRetry(p Priority, s string) (int, error) {
	pr := (w.priority & facilityMask) | (p & severityMask)

	conn := w.getConn()
	if conn != nil {
		if n, err := w.write(conn, pr, s); err == nil {
			return n, err
		}
	}

	var err error
	if conn, err = w.connect(); err != nil {
		return 0, err
	}
	return w.write(conn, pr, s)
}

// write generates and writes a syslog formatted string. It formats the
// message based on the current Formatter and Framer.
func (w *Writer) write(conn serverConn, p Priority, msg string) (int, error) {
	// ensure it ends in a \n
	if !strings.HasSuffix(msg, "\n") {
		msg += "\n"
	}

	err := conn.writeString(w.framer, w.formatter, p, w.hostname, w.tag, msg)
	if err != nil {
		return 0, err
	}
	// Note: return the length of the input, not the number of
	// bytes printed by Fprintf, because this must behave like
	// an io.Writer.
	return len(msg), nil
}
<commit_msg>Don't use defer in hot path.<commit_after>package srslog

import (
	"crypto/tls"
	"strings"
	"sync"
)

// A Writer is a connection to a syslog server.
type Writer struct {
	priority  Priority
	tag       string
	hostname  string
	network   string
	raddr     string
	tlsConfig *tls.Config
	framer    Framer
	formatter Formatter

	mu   sync.RWMutex // guards conn
	conn serverConn
}

// getConn provides access to the internal conn, protected by a mutex. The
// conn is threadsafe, so it can be used while unlocked, but we want to avoid
// race conditions on grabbing a reference to it.
func (w *Writer) getConn() serverConn {
	w.mu.RLock()
	conn := w.conn
	w.mu.RUnlock()
	return conn
}

// setConn updates the internal conn, protected by a mutex.
func (w *Writer) setConn(c serverConn) {
	w.mu.Lock()
	w.conn = c
	w.mu.Unlock()
}

// connect makes a connection to the syslog server.
func (w *Writer) connect() (serverConn, error) {
	conn := w.getConn()
	if conn != nil {
		// ignore err from close, it makes sense to continue anyway
		conn.close()
		w.setConn(nil)
	}

	var hostname string
	var err error
	dialer := w.getDialer()
	conn, hostname, err = dialer.Call()
	if err == nil {
		w.setConn(conn)
		w.hostname = hostname

		return conn, nil
	} else {
		return nil, err
	}
}

// SetFormatter changes the formatter function for subsequent messages.
func (w *Writer) SetFormatter(f Formatter) {
	w.formatter = f
}

// SetFramer changes the framer function for subsequent messages.
func (w *Writer) SetFramer(f Framer) {
	w.framer = f
}

// Write sends a log message to the syslog daemon using the default priority
// passed into `srslog.New` or the `srslog.Dial*` functions.
func (w *Writer) Write(b []byte) (int, error) {
	return w.writeAndRetry(w.priority, string(b))
}

// WriteWithPriority sends a log message with a custom priority
func (w *Writer) WriteWithPriority(p Priority, b []byte) (int, error) {
	return w.writeAndRetry(p, string(b))
}

// Close closes a connection to the syslog daemon.
func (w *Writer) Close() error {
	conn := w.getConn()
	if conn != nil {
		err := conn.close()
		w.setConn(nil)
		return err
	}
	return nil
}

// Emerg logs a message with severity LOG_EMERG; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Emerg(m string) (err error) {
	_, err = w.writeAndRetry(LOG_EMERG, m)
	return err
}

// Alert logs a message with severity LOG_ALERT; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Alert(m string) (err error) {
	_, err = w.writeAndRetry(LOG_ALERT, m)
	return err
}

// Crit logs a message with severity LOG_CRIT; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Crit(m string) (err error) {
	_, err = w.writeAndRetry(LOG_CRIT, m)
	return err
}

// Err logs a message with severity LOG_ERR; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Err(m string) (err error) {
	_, err = w.writeAndRetry(LOG_ERR, m)
	return err
}

// Warning logs a message with severity LOG_WARNING; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Warning(m string) (err error) {
	_, err = w.writeAndRetry(LOG_WARNING, m)
	return err
}

// Notice logs a message with severity LOG_NOTICE; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Notice(m string) (err error) {
	_, err = w.writeAndRetry(LOG_NOTICE, m)
	return err
}

// Info logs a message with severity LOG_INFO; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Info(m string) (err error) {
	_, err = w.writeAndRetry(LOG_INFO, m)
	return err
}

// Debug logs a message with severity LOG_DEBUG; this overrides the default
// priority passed to `srslog.New` and the `srslog.Dial*` functions.
func (w *Writer) Debug(m string) (err error) {
	_, err = w.writeAndRetry(LOG_DEBUG, m)
	return err
}

func (w *Writer) writeAndRetry(p Priority, s string) (int, error) {
	pr := (w.priority & facilityMask) | (p & severityMask)

	conn := w.getConn()
	if conn != nil {
		if n, err := w.write(conn, pr, s); err == nil {
			return n, err
		}
	}

	var err error
	if conn, err = w.connect(); err != nil {
		return 0, err
	}
	return w.write(conn, pr, s)
}

// write generates and writes a syslog formatted string. It formats the
// message based on the current Formatter and Framer.
func (w *Writer) write(conn serverConn, p Priority, msg string) (int, error) {
	// ensure it ends in a \n
	if !strings.HasSuffix(msg, "\n") {
		msg += "\n"
	}

	err := conn.writeString(w.framer, w.formatter, p, w.hostname, w.tag, msg)
	if err != nil {
		return 0, err
	}
	// Note: return the length of the input, not the number of
	// bytes printed by Fprintf, because this must behave like
	// an io.Writer.
	return len(msg), nil
}
<|endoftext|>


<commit_before>//
// Copyright © 2011-2013 Guy M. Allard
//
// 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 stompngo

import (
	"bufio"
	"fmt"
	"strconv"
	"time"
)

/*
	Logical network writer.  Read wiredata structures from the communication
	channel, and put them on the wire.
*/
func (c *Connection) writer() {
	q := false
	for {

		select {
		case d := <-c.output:
			c.wireWrite(d)
		case q = <-c.wsd:
			break
		}

		if q {
			break
		}

	}
	c.log("writer shutdown", time.Now())
}

/*
	Connection logical write.
*/
func (c *Connection) wireWrite(d wiredata) {
	f := &d.frame
	switch f.Command {
	case "\n": // HeartBeat frame
		if _, e := fmt.Fprintf(c.wtr, "%s", f.Command); e != nil {
			d.errchan <- e
			return
		}
	default: // Other frames
		if e := f.writeFrame(c.wtr, c.Protocol()); e != nil {
			d.errchan <- e
			return
		}
		if e := c.wtr.WriteByte('\x00'); e != nil {
			d.errchan <- e
			return
		}
	}
	if e := c.wtr.Flush(); e != nil {
		d.errchan <- e
		return
	}
	//
	if c.hbd != nil {
		c.hbd.ls = time.Now().UnixNano() // Latest good send
	}
	c.mets.tfw += 1             // Frame written count
	c.mets.tbw += f.Size(false) // Bytes written count
	//
	d.errchan <- nil
	return
}

/*
	Frame physical write.
*/
func (f *Frame) writeFrame(w *bufio.Writer, l string) error {
	// Write the frame Command
	if _, e := fmt.Fprintf(w, "%s\n", f.Command); e != nil {
		return e
	}
	// Content length - Always add it if client does not suppress it and
	// does not supply it.
	if _, ok := f.Headers.Contains("suppress-content-length"); !ok {
		if _, clok := f.Headers.Contains("content-length"); !clok {
			l := strconv.Itoa(len(f.Body))
			f.Headers = f.Headers.Add("content-length", l)
		}
	}
	// Write the frame Headers
	for i := 0; i < len(f.Headers); i += 2 {
		if l > SPL_10 && f.Command != CONNECT {
			f.Headers[i] = encode(f.Headers[i])
			f.Headers[i+1] = encode(f.Headers[i+1])
		}
		_, e := fmt.Fprintf(w, "%s:%s\n", f.Headers[i], f.Headers[i+1])
		if e != nil {
			return e
		}
	}
	// Write the last Header LF
	if _, e := fmt.Fprintf(w, "\n"); e != nil {
		return e
	}
	// Write the body
	if _, e := w.Write(f.Body); e != nil {
		return e
	}
	return nil
}
<commit_msg>Avoid "fmt" in writer<commit_after>//
// Copyright © 2011-2013 Guy M. Allard
//
// 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 stompngo

import (
	"bufio"
	"strconv"
	"time"
)

/*
	Logical network writer.  Read wiredata structures from the communication
	channel, and put them on the wire.
*/
func (c *Connection) writer() {
	q := false
	for {

		select {
		case d := <-c.output:
			c.wireWrite(d)
		case q = <-c.wsd:
			break
		}

		if q {
			break
		}

	}
	c.log("writer shutdown", time.Now())
}

/*
	Connection logical write.
*/
func (c *Connection) wireWrite(d wiredata) {
	f := &d.frame
	switch f.Command {
	case "\n": // HeartBeat frame
		if _, e := c.wtr.WriteString(f.Command); e != nil {
			d.errchan <- e
			return
		}
	default: // Other frames
		if e := f.writeFrame(c.wtr, c.Protocol()); e != nil {
			d.errchan <- e
			return
		}
		if e := c.wtr.WriteByte('\x00'); e != nil {
			d.errchan <- e
			return
		}
	}
	if e := c.wtr.Flush(); e != nil {
		d.errchan <- e
		return
	}
	//
	if c.hbd != nil {
		c.hbd.ls = time.Now().UnixNano() // Latest good send
	}
	c.mets.tfw += 1             // Frame written count
	c.mets.tbw += f.Size(false) // Bytes written count
	//
	d.errchan <- nil
	return
}

/*
	Frame physical write.
*/
func (f *Frame) writeFrame(w *bufio.Writer, l string) error {
	// Write the frame Command
	if _, e := w.WriteString(f.Command + "\n"); e != nil {
		return e
	}
	// Content length - Always add it if client does not suppress it and
	// does not supply it.
	if _, ok := f.Headers.Contains("suppress-content-length"); !ok {
		if _, clok := f.Headers.Contains("content-length"); !clok {
			l := strconv.Itoa(len(f.Body))
			f.Headers = f.Headers.Add("content-length", l)
		}
	}
	// Write the frame Headers
	for i := 0; i < len(f.Headers); i += 2 {
		if l > SPL_10 && f.Command != CONNECT {
			f.Headers[i] = encode(f.Headers[i])
			f.Headers[i+1] = encode(f.Headers[i+1])
		}
		_, e := w.WriteString(f.Headers[i] + ":" + f.Headers[i+1] + "\n")
		if e != nil {
			return e
		}
	}
	// Write the last Header LF
	if e := w.WriteByte('\n'); e != nil {
		return e
	}
	// Write the body
	if len(f.Body) != 0 { // Foolish to write 0 length data
		if _, e := w.Write(f.Body); e != nil {
			return e
		}
	}
	return nil
}
<|endoftext|>