<commit_before>package main

import (
	"fmt"
	"sync"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

type ControllerOptions struct {
	Namespace string
	Matcher   Matcher
}

type (
	ContainerEnterFunc func(
		pod *v1.Pod,
		container *v1.Container,
		initialAddPhase bool) bool

	ContainerExitFunc func(pod *v1.Pod,
		container *v1.Container)

	ContainerErrorFunc func(pod *v1.Pod,
		container *v1.Container, err error)
)

type Callbacks struct {
	OnEvent LogEventFunc
	OnEnter ContainerEnterFunc
	OnExit  ContainerExitFunc
	OnError ContainerErrorFunc
}

type Controller struct {
	ControllerOptions
	clientset *kubernetes.Clientset
	tailers   map[string]*ContainerTailer
	callbacks Callbacks
	sync.Mutex
}

func NewController(
	clientset *kubernetes.Clientset,
	options ControllerOptions,
	callbacks Callbacks) *Controller {
	return &Controller{
		ControllerOptions: options,
		clientset:         clientset,
		tailers:           map[string]*ContainerTailer{},
		callbacks:         callbacks,
	}
}

func (ctl *Controller) Run() {
	podListWatcher := cache.NewListWatchFromClient(
		ctl.clientset.CoreV1Client.RESTClient(), "pods", ctl.Namespace, fields.Everything())

	obj, err := podListWatcher.List(metav1.ListOptions{})
	if err != nil {
		panic(err)
	}
	if podList, ok := obj.(*v1.PodList); ok {
		for _, pod := range podList.Items {
			ctl.onInitialAdd(&pod)
		}
	}

	_, informer := cache.NewIndexerInformer(
		podListWatcher, &v1.Pod{}, 0, cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				if pod, ok := obj.(*v1.Pod); ok {
					ctl.onAdd(pod)
				}
			},
			UpdateFunc: func(old interface{}, new interface{}) {
				if pod, ok := new.(*v1.Pod); ok {
					ctl.onUpdate(pod)
				}
			},
			DeleteFunc: func(obj interface{}) {
				if pod, ok := obj.(*v1.Pod); ok {
					ctl.onDelete(pod)
				}
			},
		}, cache.Indexers{})

	stopCh := make(chan struct{}, 1)
	go informer.Run(stopCh)
	<-stopCh
}

func (ctl *Controller) onInitialAdd(pod *v1.Pod) {
	for _, container := range pod.Spec.InitContainers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, true)
		}
	}
	for _, container := range pod.Spec.Containers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, true)
		}
	}
}

func (ctl *Controller) onAdd(pod *v1.Pod) {
	for _, container := range pod.Spec.InitContainers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, false)
		}
	}
	for _, container := range pod.Spec.Containers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, false)
		}
	}
}

func (ctl *Controller) onUpdate(pod *v1.Pod) {
	ctl.onUpdateWithContainers(pod, pod.Spec.Containers,
		pod.Status.ContainerStatuses)
	ctl.onUpdateWithContainers(pod, pod.Spec.InitContainers,
		pod.Status.InitContainerStatuses)
}

func (ctl *Controller) onUpdateWithContainers(pod *v1.Pod,
	containers []v1.Container,
	containerStatuses []v1.ContainerStatus) {
	for _, containerStatus := range containerStatuses {
		var container *v1.Container
		for i, c := range containers {
			if c.Name == containerStatus.Name {
				container = &containers[i]
				break
			}
		}
		if container == nil {
			// Should be impossible; means there's a status for a container that isn't
			// part of the spec
			continue
		}

		if ctl.shouldIncludeContainer(pod, container) {
			ctl.addContainer(pod, container, false)
		} else {
			ctl.deleteContainer(pod, container)
		}
	}
}

func (ctl *Controller) onDelete(pod *v1.Pod) {
	for _, container := range pod.Spec.Containers {
		ctl.deleteContainer(pod, &container)
	}
}

func (ctl *Controller) shouldIncludePod(pod *v1.Pod) bool {
	if !ctl.Matcher.Match(pod) {
		return false
	}
	if pod.Status.Phase != v1.PodRunning &&
		pod.Status.Phase != v1.PodPending {
		return false
	}
	return true
}

func (ctl *Controller) shouldIncludeContainer(pod *v1.Pod, container *v1.Container) bool {
	if !ctl.shouldIncludePod(pod) {
		return false
	}
	var status *v1.ContainerStatus
	for _, s := range pod.Status.ContainerStatuses {
		if s.Name == container.Name {
			status = &s
			break
		}
	}
	if status == nil {
		for _, s := range pod.Status.InitContainerStatuses {
			if s.Name == container.Name {
				status = &s
				break
			}
		}
		if status == nil {
			return false
		}
	}
	if status.State.Waiting != nil || status.State.Terminated != nil ||
		status.State.Running == nil {
		return false
	}
	return true
}

func (ctl *Controller) addContainer(
	pod *v1.Pod,
	container *v1.Container,
	initialAdd bool) {
	ctl.Lock()
	defer ctl.Unlock()

	key := buildKey(pod, container)
	if _, ok := ctl.tailers[key]; ok {
		return
	}

	if !ctl.callbacks.OnEnter(pod, container, initialAdd) {
		return
	}

	targetPod, targetContainer := *pod, *container // Copy to avoid mutation

	var fromTimestamp *time.Time
	if initialAdd {
		// Don't show any history, but add a small amount of buffer to
		// account for clock skew
		now := time.Now().Add(time.Second * -5)
		fromTimestamp = &now
	} else {
		for _, status := range pod.Status.ContainerStatuses {
			if status.Name == container.Name && status.State.Running != nil {
				startTime := status.State.Running.StartedAt.Time
				fromTimestamp = &startTime
				break
			}
		}
	}

	tailer := NewContainerTailer(ctl.clientset, targetPod, targetContainer,
		ctl.callbacks.OnEvent, fromTimestamp)
	ctl.tailers[key] = tailer

	go func() {
		tailer.Run(func(err error) {
			ctl.callbacks.OnError(&targetPod, &targetContainer, err)
		})
	}()
}

func (ctl *Controller) deleteContainer(pod *v1.Pod, container *v1.Container) {
	ctl.Lock()
	defer ctl.Unlock()

	key := buildKey(pod, container)
	if tailer, ok := ctl.tailers[key]; ok {
		delete(ctl.tailers, key)
		tailer.Stop()
		ctl.callbacks.OnExit(pod, container)
	}
}

func buildKey(pod *v1.Pod, container *v1.Container) string {
	return fmt.Sprintf("%s/%s/%s", pod.Namespace, pod.Name, container.Name)
}
<commit_msg>Fix edge case where we might use the wrong timestamp if a newly discovered pod has multiple containers that started at different times.<commit_after>package main

import (
	"fmt"
	"sync"
	"time"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/tools/cache"
)

type ControllerOptions struct {
	Namespace string
	Matcher   Matcher
}

type (
	ContainerEnterFunc func(
		pod *v1.Pod,
		container *v1.Container,
		initialAddPhase bool) bool

	ContainerExitFunc func(pod *v1.Pod,
		container *v1.Container)

	ContainerErrorFunc func(pod *v1.Pod,
		container *v1.Container, err error)
)

type Callbacks struct {
	OnEvent LogEventFunc
	OnEnter ContainerEnterFunc
	OnExit  ContainerExitFunc
	OnError ContainerErrorFunc
}

type Controller struct {
	ControllerOptions
	clientset *kubernetes.Clientset
	tailers   map[string]*ContainerTailer
	callbacks Callbacks
	sync.Mutex
}

func NewController(
	clientset *kubernetes.Clientset,
	options ControllerOptions,
	callbacks Callbacks) *Controller {
	return &Controller{
		ControllerOptions: options,
		clientset:         clientset,
		tailers:           map[string]*ContainerTailer{},
		callbacks:         callbacks,
	}
}

func (ctl *Controller) Run() {
	podListWatcher := cache.NewListWatchFromClient(
		ctl.clientset.CoreV1Client.RESTClient(), "pods", ctl.Namespace, fields.Everything())

	obj, err := podListWatcher.List(metav1.ListOptions{})
	if err != nil {
		panic(err)
	}
	if podList, ok := obj.(*v1.PodList); ok {
		for _, pod := range podList.Items {
			ctl.onInitialAdd(&pod)
		}
	}

	_, informer := cache.NewIndexerInformer(
		podListWatcher, &v1.Pod{}, 0, cache.ResourceEventHandlerFuncs{
			AddFunc: func(obj interface{}) {
				if pod, ok := obj.(*v1.Pod); ok {
					ctl.onAdd(pod)
				}
			},
			UpdateFunc: func(old interface{}, new interface{}) {
				if pod, ok := new.(*v1.Pod); ok {
					ctl.onUpdate(pod)
				}
			},
			DeleteFunc: func(obj interface{}) {
				if pod, ok := obj.(*v1.Pod); ok {
					ctl.onDelete(pod)
				}
			},
		}, cache.Indexers{})

	stopCh := make(chan struct{}, 1)
	go informer.Run(stopCh)
	<-stopCh
}

func (ctl *Controller) onInitialAdd(pod *v1.Pod) {
	for _, container := range pod.Spec.InitContainers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, true)
		}
	}
	for _, container := range pod.Spec.Containers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, true)
		}
	}
}

func (ctl *Controller) onAdd(pod *v1.Pod) {
	for _, container := range pod.Spec.InitContainers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, false)
		}
	}
	for _, container := range pod.Spec.Containers {
		if ctl.shouldIncludeContainer(pod, &container) {
			ctl.addContainer(pod, &container, false)
		}
	}
}

func (ctl *Controller) onUpdate(pod *v1.Pod) {
	ctl.onUpdateWithContainers(pod, pod.Spec.Containers,
		pod.Status.ContainerStatuses)
	ctl.onUpdateWithContainers(pod, pod.Spec.InitContainers,
		pod.Status.InitContainerStatuses)
}

func (ctl *Controller) onUpdateWithContainers(pod *v1.Pod,
	containers []v1.Container,
	containerStatuses []v1.ContainerStatus) {
	for _, containerStatus := range containerStatuses {
		var container *v1.Container
		for i, c := range containers {
			if c.Name == containerStatus.Name {
				container = &containers[i]
				break
			}
		}
		if container == nil {
			// Should be impossible; means there's a status for a container that isn't
			// part of the spec
			continue
		}

		if ctl.shouldIncludeContainer(pod, container) {
			ctl.addContainer(pod, container, false)
		} else {
			ctl.deleteContainer(pod, container)
		}
	}
}

func (ctl *Controller) onDelete(pod *v1.Pod) {
	for _, container := range pod.Spec.Containers {
		ctl.deleteContainer(pod, &container)
	}
}

func (ctl *Controller) shouldIncludePod(pod *v1.Pod) bool {
	if !ctl.Matcher.Match(pod) {
		return false
	}
	if pod.Status.Phase != v1.PodRunning &&
		pod.Status.Phase != v1.PodPending {
		return false
	}
	return true
}

func (ctl *Controller) shouldIncludeContainer(pod *v1.Pod, container *v1.Container) bool {
	if !ctl.shouldIncludePod(pod) {
		return false
	}
	var status *v1.ContainerStatus
	for _, s := range pod.Status.ContainerStatuses {
		if s.Name == container.Name {
			status = &s
			break
		}
	}
	if status == nil {
		for _, s := range pod.Status.InitContainerStatuses {
			if s.Name == container.Name {
				status = &s
				break
			}
		}
		if status == nil {
			return false
		}
	}
	if status.State.Waiting != nil || status.State.Terminated != nil ||
		status.State.Running == nil {
		return false
	}
	return true
}

func (ctl *Controller) addContainer(
	pod *v1.Pod,
	container *v1.Container,
	initialAdd bool) {
	ctl.Lock()
	defer ctl.Unlock()

	key := buildKey(pod, container)
	if _, ok := ctl.tailers[key]; ok {
		return
	}

	if !ctl.callbacks.OnEnter(pod, container, initialAdd) {
		return
	}

	targetPod, targetContainer := *pod, *container // Copy to avoid mutation

	var fromTimestamp *time.Time
	if initialAdd {
		// Don't show any history, but add a small amount of buffer to
		// account for clock skew
		now := time.Now().Add(time.Second * -5)
		fromTimestamp = &now
	} else {
		for _, status := range pod.Status.ContainerStatuses {
			if status.Name == container.Name && status.State.Running != nil {
				startTime := status.State.Running.StartedAt.Time
				if fromTimestamp == nil || startTime.Before(*fromTimestamp) {
					fromTimestamp = &startTime
				}
			}
		}
	}

	tailer := NewContainerTailer(ctl.clientset, targetPod, targetContainer,
		ctl.callbacks.OnEvent, fromTimestamp)
	ctl.tailers[key] = tailer

	go func() {
		tailer.Run(func(err error) {
			ctl.callbacks.OnError(&targetPod, &targetContainer, err)
		})
	}()
}

func (ctl *Controller) deleteContainer(pod *v1.Pod, container *v1.Container) {
	ctl.Lock()
	defer ctl.Unlock()

	key := buildKey(pod, container)
	if tailer, ok := ctl.tailers[key]; ok {
		delete(ctl.tailers, key)
		tailer.Stop()
		ctl.callbacks.OnExit(pod, container)
	}
}

func buildKey(pod *v1.Pod, container *v1.Container) string {
	return fmt.Sprintf("%s/%s/%s", pod.Namespace, pod.Name, container.Name)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"memcached.connections": mp.Graphs{
		Label: "Memcached Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "curr_connections", Label: "Connections", Diff: false},
		},
	},
	"memcached.cmd": mp.Graphs{
		Label: "Memcached Command",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "cmd_get", Label: "Get", Diff: true},
			mp.Metrics{Key: "cmd_set", Label: "Set", Diff: true},
			mp.Metrics{Key: "cmd_flush", Label: "Flush", Diff: true},
			mp.Metrics{Key: "cmd_touch", Label: "Touch", Diff: true},
		},
	},
	"memcached.hitmiss": mp.Graphs{
		Label: "Memcached Hits/Misses",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "get_hits", Label: "Get Hits", Diff: true},
			mp.Metrics{Key: "get_misses", Label: "Get Misses", Diff: true},
			mp.Metrics{Key: "delete_hits", Label: "Delete Hits", Diff: true},
			mp.Metrics{Key: "delete_misses", Label: "Delete Misses", Diff: true},
			mp.Metrics{Key: "incr_hits", Label: "Incr Hits", Diff: true},
			mp.Metrics{Key: "incr_misses", Label: "Incr Misses", Diff: true},
			mp.Metrics{Key: "cas_hits", Label: "Cas Hits", Diff: true},
			mp.Metrics{Key: "cas_misses", Label: "Cas Misses", Diff: true},
			mp.Metrics{Key: "touch_hits", Label: "Touch Hits", Diff: true},
			mp.Metrics{Key: "touch_misses", Label: "Touch Misses", Diff: true},
		},
	},
	"memcached.evictions": mp.Graphs{
		Label: "Memcached Evictions",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "evictions", Label: "Evictions", Diff: true},
		},
	},
	"memcached.unfetched": mp.Graphs{
		Label: "Memcached Unfetched",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "expired_unfetched", Label: "Expired unfetched", Diff: true},
			mp.Metrics{Key: "evicted_unfetched", Label: "Evicted unfetched", Diff: true},
		},
	},
	"memcached.rusage": mp.Graphs{
		Label: "Memcached Resouce Usage",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "rusage_user", Label: "User", Diff: true},
			mp.Metrics{Key: "rusage_system", Label: "System", Diff: true},
		},
	},
	"memcached.bytes": mp.Graphs{
		Label: "Memcached Traffics",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Key: "bytes_read", Label: "Read", Diff: true},
			mp.Metrics{Key: "bytes_written", Label: "Write", Diff: true},
		},
	},
}

type MemcachedPlugin struct {
	Target   string
	Tempfile string
}

func (m MemcachedPlugin) FetchData() (map[string]float64, error) {
	conn, err := net.Dial("tcp", m.Target)
	if err != nil {
		return nil, err
	}
	fmt.Fprintln(conn, "stats")
	scanner := bufio.NewScanner(conn)
	stat := make(map[string]float64)

	for scanner.Scan() {
		line := scanner.Text()
		s := string(line)
		if s == "END" {
			return stat, nil
		}

		res := strings.Split(s, " ")
		if res[0] == "STAT" {
			stat[res[1]], err = strconv.ParseFloat(res[2], 64)
			if err != nil {
				log.Println("FetchData:", err)
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return stat, err
	}
	return nil, nil
}

func (m MemcachedPlugin) GetGraphDefinition() map[string](mp.Graphs) {
	return graphdef
}

func (m MemcachedPlugin) GetTempfilename() string {
	return m.Tempfile
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "11211", "Port")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var memcached MemcachedPlugin

	memcached.Target = fmt.Sprintf("%s:%s", *optHost, *optPort)
	if *optTempfile != "" {
		memcached.Tempfile = *optTempfile
	} else {
		memcached.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-memcached-%s-%s", *optHost, *optPort)
	}

	helper := mp.MackerelPlugin{memcached}

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

import (
	"bufio"
	"flag"
	"fmt"
	mp "github.com/mackerelio/go-mackerel-plugin"
	"log"
	"net"
	"os"
	"strconv"
	"strings"
)

var graphdef map[string](mp.Graphs) = map[string](mp.Graphs){
	"memcached.connections": mp.Graphs{
		Label: "Memcached Connections",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "curr_connections", Label: "Connections", Diff: false},
		},
	},
	"memcached.cmd": mp.Graphs{
		Label: "Memcached Command",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "cmd_get", Label: "Get", Diff: true},
			mp.Metrics{Name: "cmd_set", Label: "Set", Diff: true},
			mp.Metrics{Name: "cmd_flush", Label: "Flush", Diff: true},
			mp.Metrics{Name: "cmd_touch", Label: "Touch", Diff: true},
		},
	},
	"memcached.hitmiss": mp.Graphs{
		Label: "Memcached Hits/Misses",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "get_hits", Label: "Get Hits", Diff: true},
			mp.Metrics{Name: "get_misses", Label: "Get Misses", Diff: true},
			mp.Metrics{Name: "delete_hits", Label: "Delete Hits", Diff: true},
			mp.Metrics{Name: "delete_misses", Label: "Delete Misses", Diff: true},
			mp.Metrics{Name: "incr_hits", Label: "Incr Hits", Diff: true},
			mp.Metrics{Name: "incr_misses", Label: "Incr Misses", Diff: true},
			mp.Metrics{Name: "cas_hits", Label: "Cas Hits", Diff: true},
			mp.Metrics{Name: "cas_misses", Label: "Cas Misses", Diff: true},
			mp.Metrics{Name: "touch_hits", Label: "Touch Hits", Diff: true},
			mp.Metrics{Name: "touch_misses", Label: "Touch Misses", Diff: true},
		},
	},
	"memcached.evictions": mp.Graphs{
		Label: "Memcached Evictions",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "evictions", Label: "Evictions", Diff: true},
		},
	},
	"memcached.unfetched": mp.Graphs{
		Label: "Memcached Unfetched",
		Unit:  "integer",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "expired_unfetched", Label: "Expired unfetched", Diff: true},
			mp.Metrics{Name: "evicted_unfetched", Label: "Evicted unfetched", Diff: true},
		},
	},
	"memcached.rusage": mp.Graphs{
		Label: "Memcached Resouce Usage",
		Unit:  "float",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "rusage_user", Label: "User", Diff: true},
			mp.Metrics{Name: "rusage_system", Label: "System", Diff: true},
		},
	},
	"memcached.bytes": mp.Graphs{
		Label: "Memcached Traffics",
		Unit:  "bytes",
		Metrics: [](mp.Metrics){
			mp.Metrics{Name: "bytes_read", Label: "Read", Diff: true},
			mp.Metrics{Name: "bytes_written", Label: "Write", Diff: true},
		},
	},
}

type MemcachedPlugin struct {
	Target   string
	Tempfile string
}

func (m MemcachedPlugin) FetchMetrics() (map[string]float64, error) {
	conn, err := net.Dial("tcp", m.Target)
	if err != nil {
		return nil, err
	}
	fmt.Fprintln(conn, "stats")
	scanner := bufio.NewScanner(conn)
	stat := make(map[string]float64)

	for scanner.Scan() {
		line := scanner.Text()
		s := string(line)
		if s == "END" {
			return stat, nil
		}

		res := strings.Split(s, " ")
		if res[0] == "STAT" {
			stat[res[1]], err = strconv.ParseFloat(res[2], 64)
			if err != nil {
				log.Println("FetchMetrics:", err)
			}
		}
	}
	if err := scanner.Err(); err != nil {
		return stat, err
	}
	return nil, nil
}

func (m MemcachedPlugin) GraphDefinition() map[string](mp.Graphs) {
	return graphdef
}

func main() {
	optHost := flag.String("host", "localhost", "Hostname")
	optPort := flag.String("port", "11211", "Port")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var memcached MemcachedPlugin
	helper := mp.NewMackerelPlugin(memcached)

	memcached.Target = fmt.Sprintf("%s:%s", *optHost, *optPort)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-memcached-%s-%s", *optHost, *optPort)
	}

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


<commit_before>package vegeta

import (
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"strconv"
	"strings"
	"text/tabwriter"

	lttb "github.com/dgryski/go-lttb"
	"github.com/lucasb-eyer/go-colorful"
)

// A Report represents the state a Reporter uses to write out its reports.
type Report interface {
	// Add adds a given *Result to a Report.
	Add(*Result)
}

// Closer wraps the optional Report Close method.
type Closer interface {
	// Close permantently closes a Report, running any necessary book keeping.
	Close()
}

// A Reporter function writes out reports to the given io.Writer or returns an
// error in case of failure.
type Reporter func(io.Writer) error

// Report is a convenience method wrapping the Reporter function type.
func (rep Reporter) Report(w io.Writer) error { return rep(w) }

// NewHistogramReporter returns a Reporter that writes out a Histogram as
// aligned, formatted text.
func NewHistogramReporter(h *Histogram) Reporter {
	return func(w io.Writer) (err error) {
		tw := tabwriter.NewWriter(w, 0, 8, 2, ' ', tabwriter.StripEscape)
		if _, err = fmt.Fprintf(tw, "Bucket\t\t#\t%%\tHistogram\n"); err != nil {
			return err
		}

		for i, count := range h.Counts {
			ratio := float64(count) / float64(h.Total)
			lo, hi := h.Buckets.Nth(i)
			pad := strings.Repeat("#", int(ratio*75))
			_, err = fmt.Fprintf(tw, "[%s,\t%s]\t%d\t%.2f%%\t%s\n", lo, hi, count, ratio*100, pad)
			if err != nil {
				return nil
			}
		}

		return tw.Flush()
	}
}

// NewTextReporter returns a Reporter that writes out Metrics as aligned,
// formatted text.
func NewTextReporter(m *Metrics) Reporter {
	const fmtstr = "Requests\t[total, rate]\t%d, %.2f\n" +
		"Duration\t[total, attack, wait]\t%s, %s, %s\n" +
		"Latencies\t[mean, 50, 95, 99, max]\t%s, %s, %s, %s, %s\n" +
		"Bytes In\t[total, mean]\t%d, %.2f\n" +
		"Bytes Out\t[total, mean]\t%d, %.2f\n" +
		"Success\t[ratio]\t%.2f%%\n" +
		"Status Codes\t[code:count]\t"

	return func(w io.Writer) (err error) {
		tw := tabwriter.NewWriter(w, 0, 8, 2, ' ', tabwriter.StripEscape)
		if _, err = fmt.Fprintf(tw, fmtstr,
			m.Requests, m.Rate,
			m.Duration+m.Wait, m.Duration, m.Wait,
			m.Latencies.Mean, m.Latencies.P50, m.Latencies.P95, m.Latencies.P99, m.Latencies.Max,
			m.BytesIn.Total, m.BytesIn.Mean,
			m.BytesOut.Total, m.BytesOut.Mean,
			m.Success*100,
		); err != nil {
			return err
		}

		for code, count := range m.StatusCodes {
			if _, err = fmt.Fprintf(tw, "%s:%d  ", code, count); err != nil {
				return err
			}
		}

		if _, err = fmt.Fprintln(tw, "\nError Set:"); err != nil {
			return err
		}

		for _, e := range m.Errors {
			if _, err = fmt.Fprintln(tw, e); err != nil {
				return err
			}
		}

		return tw.Flush()
	}
}

// NewJSONReporter returns a Reporter that writes out Metrics as JSON.
func NewJSONReporter(m *Metrics) Reporter {
	return func(w io.Writer) error {
		return json.NewEncoder(w).Encode(m)
	}
}

// NewPlotReporter returns a Reporter that writes a self-contained
// HTML page with an interactive plot of the latencies of Requests, built with
// http://dygraphs.com/
func NewPlotReporter(title string, rs *Results) Reporter {
	return func(w io.Writer) (err error) {
		series := make(map[string][]Results, len(*rs))
		for _, r := range *rs {
			idx := 0
			if r.Error == "" {
				idx++
			}

			if len(series[r.Attack]) == 0 {
				series[r.Attack] = make([]Results, 2)
			}

			series[r.Attack][idx] = append(series[r.Attack][idx], r)
		}

		samples := make(map[string][][]lttb.Point, len(series))
		for attack, results := range series {
			for i := range results {
				sort.Sort(results[i])
				points := make([]lttb.Point, 0, len(results[i]))
				for _, r := range results[i] {
					points = append(points, lttb.Point{
						X: r.Timestamp.Sub(results[i][0].Timestamp).Seconds(),
						Y: r.Latency.Seconds() * 1000,
					})
				}
				samples[attack] = append(samples[attack], lttb.LTTB(points, 1000))
			}
		}

		const count = 2 // OK and Errors
		i, offsets := 0, make(map[string]int, len(series))
		for name := range series {
			offsets[name] = 1 + i*count
			i++
		}

		_, err = fmt.Fprintf(w, plotsTemplateHead, title, asset(dygraphs), asset(html2canvas))
		if err != nil {
			return err
		}

		const nan = "NaN"

		data := make([]string, 1+len(series)*count)
		for attack, results := range samples {
			for idx, points := range results {
				for i, p := range points {
					for j := range data {
						data[j] = nan
					}

					offset := offsets[attack] + idx
					data[0] = strconv.FormatFloat(p.X, 'f', -1, 32)
					data[offset] = strconv.FormatFloat(p.Y, 'f', -1, 32)

					s := "[" + strings.Join(data, ",") + "]"

					if i < len(*rs)-1 {
						s += ","
					}

					if _, err = io.WriteString(w, s); err != nil {
						return err
					}
				}
			}
		}

		labels := make([]string, len(data))
		labels[0] = strconv.Quote("Seconds")

		for attack, offset := range offsets {
			labels[offset] = strconv.Quote(attack + " - ERR")
			labels[offset+1] = strconv.Quote(attack + " - OK")
		}

		colors := make([]string, 0, len(labels)-1)
		palette, err := colorful.HappyPalette(cap(colors))
		if err != nil {
			return err
		}

		for _, color := range palette {
			colors = append(colors, strconv.Quote(color.Hex()))
		}

		_, err = fmt.Fprintf(w, plotsTemplateTail, title, strings.Join(labels, ","), strings.Join(colors, ","))
		return err
	}
}

const (
	plotsTemplateHead = `<!doctype html>

  <title>%s</title>
  <meta charset="utf-8">

  <button id="download">Download as PNG</button>

`
)
<commit_msg>reporters: Fix incorrect allocation optimisation<commit_after>package vegeta

import (
	"encoding/json"
	"fmt"
	"io"
	"sort"
	"strconv"
	"strings"
	"text/tabwriter"

	lttb "github.com/dgryski/go-lttb"
	"github.com/lucasb-eyer/go-colorful"
)

// A Report represents the state a Reporter uses to write out its reports.
type Report interface {
	// Add adds a given *Result to a Report.
	Add(*Result)
}

// Closer wraps the optional Report Close method.
type Closer interface {
	// Close permantently closes a Report, running any necessary book keeping.
	Close()
}

// A Reporter function writes out reports to the given io.Writer or returns an
// error in case of failure.
type Reporter func(io.Writer) error

// Report is a convenience method wrapping the Reporter function type.
func (rep Reporter) Report(w io.Writer) error { return rep(w) }

// NewHistogramReporter returns a Reporter that writes out a Histogram as
// aligned, formatted text.
func NewHistogramReporter(h *Histogram) Reporter {
	return func(w io.Writer) (err error) {
		tw := tabwriter.NewWriter(w, 0, 8, 2, ' ', tabwriter.StripEscape)
		if _, err = fmt.Fprintf(tw, "Bucket\t\t#\t%%\tHistogram\n"); err != nil {
			return err
		}

		for i, count := range h.Counts {
			ratio := float64(count) / float64(h.Total)
			lo, hi := h.Buckets.Nth(i)
			pad := strings.Repeat("#", int(ratio*75))
			_, err = fmt.Fprintf(tw, "[%s,\t%s]\t%d\t%.2f%%\t%s\n", lo, hi, count, ratio*100, pad)
			if err != nil {
				return nil
			}
		}

		return tw.Flush()
	}
}

// NewTextReporter returns a Reporter that writes out Metrics as aligned,
// formatted text.
func NewTextReporter(m *Metrics) Reporter {
	const fmtstr = "Requests\t[total, rate]\t%d, %.2f\n" +
		"Duration\t[total, attack, wait]\t%s, %s, %s\n" +
		"Latencies\t[mean, 50, 95, 99, max]\t%s, %s, %s, %s, %s\n" +
		"Bytes In\t[total, mean]\t%d, %.2f\n" +
		"Bytes Out\t[total, mean]\t%d, %.2f\n" +
		"Success\t[ratio]\t%.2f%%\n" +
		"Status Codes\t[code:count]\t"

	return func(w io.Writer) (err error) {
		tw := tabwriter.NewWriter(w, 0, 8, 2, ' ', tabwriter.StripEscape)
		if _, err = fmt.Fprintf(tw, fmtstr,
			m.Requests, m.Rate,
			m.Duration+m.Wait, m.Duration, m.Wait,
			m.Latencies.Mean, m.Latencies.P50, m.Latencies.P95, m.Latencies.P99, m.Latencies.Max,
			m.BytesIn.Total, m.BytesIn.Mean,
			m.BytesOut.Total, m.BytesOut.Mean,
			m.Success*100,
		); err != nil {
			return err
		}

		for code, count := range m.StatusCodes {
			if _, err = fmt.Fprintf(tw, "%s:%d  ", code, count); err != nil {
				return err
			}
		}

		if _, err = fmt.Fprintln(tw, "\nError Set:"); err != nil {
			return err
		}

		for _, e := range m.Errors {
			if _, err = fmt.Fprintln(tw, e); err != nil {
				return err
			}
		}

		return tw.Flush()
	}
}

// NewJSONReporter returns a Reporter that writes out Metrics as JSON.
func NewJSONReporter(m *Metrics) Reporter {
	return func(w io.Writer) error {
		return json.NewEncoder(w).Encode(m)
	}
}

// NewPlotReporter returns a Reporter that writes a self-contained
// HTML page with an interactive plot of the latencies of Requests, built with
// http://dygraphs.com/
func NewPlotReporter(title string, rs *Results) Reporter {
	return func(w io.Writer) (err error) {
		series := map[string][]Results{}
		for _, r := range *rs {
			idx := 0
			if r.Error == "" {
				idx++
			}

			if len(series[r.Attack]) == 0 {
				series[r.Attack] = make([]Results, 2)
			}

			series[r.Attack][idx] = append(series[r.Attack][idx], r)
		}

		samples := make(map[string][][]lttb.Point, len(series))
		for attack, results := range series {
			for i := range results {
				sort.Sort(results[i])
				points := make([]lttb.Point, 0, len(results[i]))
				for _, r := range results[i] {
					points = append(points, lttb.Point{
						X: r.Timestamp.Sub(results[i][0].Timestamp).Seconds(),
						Y: r.Latency.Seconds() * 1000,
					})
				}
				samples[attack] = append(samples[attack], lttb.LTTB(points, 1000))
			}
		}

		const count = 2 // OK and Errors
		i, offsets := 0, make(map[string]int, len(series))
		for name := range series {
			offsets[name] = 1 + i*count
			i++
		}

		_, err = fmt.Fprintf(w, plotsTemplateHead, title, asset(dygraphs), asset(html2canvas))
		if err != nil {
			return err
		}

		const nan = "NaN"

		data := make([]string, 1+len(series)*count)
		for attack, results := range samples {
			for idx, points := range results {
				for i, p := range points {
					for j := range data {
						data[j] = nan
					}

					offset := offsets[attack] + idx
					data[0] = strconv.FormatFloat(p.X, 'f', -1, 32)
					data[offset] = strconv.FormatFloat(p.Y, 'f', -1, 32)

					s := "[" + strings.Join(data, ",") + "]"

					if i < len(*rs)-1 {
						s += ","
					}

					if _, err = io.WriteString(w, s); err != nil {
						return err
					}
				}
			}
		}

		labels := make([]string, len(data))
		labels[0] = strconv.Quote("Seconds")

		for attack, offset := range offsets {
			labels[offset] = strconv.Quote(attack + " - ERR")
			labels[offset+1] = strconv.Quote(attack + " - OK")
		}

		colors := make([]string, 0, len(labels)-1)
		palette, err := colorful.HappyPalette(cap(colors))
		if err != nil {
			return err
		}

		for _, color := range palette {
			colors = append(colors, strconv.Quote(color.Hex()))
		}

		_, err = fmt.Fprintf(w, plotsTemplateTail, title, strings.Join(labels, ","), strings.Join(colors, ","))
		return err
	}
}

const (
	plotsTemplateHead = `<!doctype html>

  <title>%s</title>
  <meta charset="utf-8">

  <button id="download">Download as PNG</button>

`
)
<|endoftext|>


<commit_before>package execution

import (
	"pfi/sensorbee/sensorbee/bql/udf"
	"pfi/sensorbee/sensorbee/core"
	"pfi/sensorbee/sensorbee/data"
)

type groupbyExecutionPlan struct {
	streamRelationStreamExecutionPlan
}

// tmpGroupData is an intermediate data structure to represent
// a set of rows that have the same values for GROUP BY columns.
type tmpGroupData struct {
	// this is the group (e.g. [1, "toy"]), where the values are
	// in order of the items in the GROUP BY clause
	group data.Array
	// for each aggregate function, we hold an array with the
	// input values.
	aggData map[string][]data.Value
	// as per our assumptions about grouping, the non-aggregation
	// data should be identical within every group
	nonAggData data.Map
}

// CanBuildGroupbyExecutionPlan checks whether the given statement
// allows to use an groupbyExecutionPlan.
func CanBuildGroupbyExecutionPlan(lp *LogicalPlan, reg udf.FunctionRegistry) bool {
	return lp.GroupingStmt
}

// groupbyExecutionPlan is a simple plan that follows the
// theoretical processing model. It supports only statements
// that use aggregation.
//
// After each tuple arrives,
// - compute the contents of the current window using the
//   specified window size/type,
// - perform a SELECT query on that data,
// - compute the data that need to be emitted by comparison with
//   the previous run's results.
func NewGroupbyExecutionPlan(lp *LogicalPlan, reg udf.FunctionRegistry) (ExecutionPlan, error) {
	underlying, err := newStreamRelationStreamExecutionPlan(lp, reg)
	if err != nil {
		return nil, err
	}
	return &groupbyExecutionPlan{
		*underlying,
	}, nil
}

// Process takes an input tuple and returns a slice of Map values that
// correspond to the results of the query represented by this execution
// plan. Note that the order of items in the returned slice is undefined
// and cannot be relied on.
func (ep *groupbyExecutionPlan) Process(input *core.Tuple) ([]data.Map, error) {
	return ep.process(input, ep.performQueryOnBuffer)
}

// performQueryOnBuffer executes a SELECT query on the data of the tuples
// currently stored in the buffer. The query results (which is a set of
// data.Value, not core.Tuple) is stored in ep.curResults. The data
// that was stored in ep.curResults before this method was called is
// moved to ep.prevResults. Note that the order of values in ep.curResults
// is undefined.
//
// In case of an error the contents of ep.curResults will still be
// the same as before the call (so that the next run performs as
// if no error had happened), but the contents of ep.curResults are
// undefined.
//
// Currently performQueryOnBuffer can only perform SELECT ... WHERE ...
// queries without aggregate functions, GROUP BY, or HAVING clauses.
func (ep *groupbyExecutionPlan) performQueryOnBuffer() error {
	// reuse the allocated memory
	output := ep.prevResults[0:0]
	// remember the previous results
	ep.prevResults = ep.curResults

	rollback := func() {
		// NB. ep.prevResults currently points to an slice with
		//     results from the previous run. ep.curResults points
		//     to the same slice. output points to a different slice
		//     with a different underlying array.
		//     in the next run, output will be reusing the underlying
		//     storage of the current ep.prevResults to hold results.
		//     therefore when we leave this function we must make
		//     sure that ep.prevResults and ep.curResults have
		//     different underlying arrays or ISTREAM/DSTREAM will
		//     return wrong results.
		ep.prevResults = output
	}

	// collect a list of all aggregate parameter evaluators in all
	// projections. this is necessary to avoid duplicate evaluation
	// if the same parameter is used in multiple aggregation funcs.
	allAggEvaluators := map[string]Evaluator{}
	for _, proj := range ep.projections {
		for key, agg := range proj.aggrEvals {
			allAggEvaluators[key] = agg
		}
	}

	// groups holds one item for every combination of values that
	// appear in the GROUP BY clause
	groups := []tmpGroupData{}

	// findOrCreateGroup looks up the group that has the given
	// groupValues in the `groups` list. if there is no such
	// group, a new one is created and a copy of the given map
	// is used as a representative of this group's values.
	findOrCreateGroup := func(groupValues []data.Value, d data.Map) (*tmpGroupData, error) {
		eq := Equal(binOp{}).(*compBinOp).cmpOp
		groupValuesArr := data.Array(groupValues)
		// find the correct group
		groupIdx := -1
		for i, groupData := range groups {
			equals, err := eq(groupData.group, groupValuesArr)
			if err != nil {
				return nil, err
			}
			if equals {
				groupIdx = i
				break
			}
		}
		// if there is no such group, create one
		if groupIdx < 0 {
			newGroup := tmpGroupData{
				// the values that make up this group
				groupValues,
				// the input values of the aggregate functions
				map[string][]data.Value{},
				// a representative set of values for this group for later evaluation
				// TODO actually we don't need the whole map,
				//      just the parts common to the whole group
				d.Copy(),
			}
			// initialize the map with the aggregate function inputs
			for _, proj := range ep.projections {
				for key := range proj.aggrEvals {
					newGroup.aggData[key] = make([]data.Value, 0, 1)
				}
			}
			groups = append(groups, newGroup)
			groupIdx = len(groups) - 1
		}

		// return a pointer to the (found or created) group
		return &groups[groupIdx], nil
	}

	// we need to make a cross product of the data in all buffers,
	// combine it to get an input like
	//  {"streamA": {data}, "streamB": {data}, "streamC": {data}}
	// and then run filter/projections on each of this items

	dataHolder := data.Map{}

	// function to evaluate filter on the input data and do the computations
	// that are required on each input tuple. (those computations differ
	// depending on whether we are in grouping mode or not.)
	evalItem := func(d data.Map) error {
		// evaluate filter condition and convert to bool
		if ep.filter != nil {
			filterResult, err := ep.filter.Eval(d)
			if err != nil {
				return err
			}
			filterResultBool, err := data.ToBool(filterResult)
			if err != nil {
				return err
			}
			// if it evaluated to false, do not further process this tuple
			// (ToBool also evalutes the NULL value to false, so we don't
			// need to treat this specially)
			if !filterResultBool {
				return nil
			}
		}

		// now compute the expressions in the GROUP BY to find the correct
		// group to append to
		itemGroupValues := make([]data.Value, len(ep.groupList))
		for i, eval := range ep.groupList {
			// ordinary "flat" expression
			value, err := eval.Eval(d)
			if err != nil {
				return err
			}
			itemGroupValues[i] = value
		}

		itemGroup, err := findOrCreateGroup(itemGroupValues, d)
		if err != nil {
			return err
		}

		// now compute all the input data for the aggregate functions,
		// e.g. for `SELECT count(a) + max(b/2)`, compute `a` and `b/2`
		for key, agg := range allAggEvaluators {
			value, err := agg.Eval(d)
			if err != nil {
				return err
			}
			// store this value in the output map
			itemGroup.aggData[key] = append(itemGroup.aggData[key], value)
		}
		return nil
	}

	evalGroup := func(group *tmpGroupData) error {
		result := data.Map(make(map[string]data.Value, len(ep.projections)))
		// collect input for aggregate functions into an array
		// within each group
		for key, _ := range allAggEvaluators {
			group.nonAggData[key] = data.Array(group.aggData[key])
			delete(group.aggData, key)
		}
		// evaluate HAVING condition, if there is one
		for _, proj := range ep.projections {
			if proj.alias == ":having:" {
				havingResult, err := proj.evaluator.Eval(group.nonAggData)
				if err != nil {
					return err
				}
				havingResultBool, err := data.ToBool(havingResult)
				if err != nil {
					return err
				}
				// if it evaluated to false, do not further process this group
				// (ToBool also evalutes the NULL value to false, so we don't
				// need to treat this specially)
				if !havingResultBool {
					return nil
				}
				break
			}
		}
		// now evaluate all other projections
		for _, proj := range ep.projections {
			if proj.alias == ":having:" {
				continue
			}
			// now evaluate this projection on the flattened data
			value, err := proj.evaluator.Eval(group.nonAggData)
			if err != nil {
				return err
			}
			if err := assignOutputValue(result, proj.alias, value); err != nil {
				return err
			}
		}
		output = append(output, result)
		return nil
	}

	evalNoGroup := func() error {
		// if we have an empty group list *and* a GROUP BY clause,
		// we have to return an empty result (because there are no
		// rows with "the same values"). but if the list is empty and
		// we *don't* have a GROUP BY clause, then we need to compute
		// all foldables and aggregates with an empty input
		if len(ep.groupList) > 0 {
			return nil
		}
		input := data.Map{}
		result := data.Map(make(map[string]data.Value, len(ep.projections)))
		for _, proj := range ep.projections {
			// collect input for aggregate functions
			if proj.hasAggregate {
				for key := range proj.aggrEvals {
					input[key] = data.Array{}
				}
			}
			// now evaluate this projection on the flattened data.
			// note that input has *only* the keys of the empty
			// arrays, no other columns, but we cannot have other
			// columns involved in the projection (since we know
			// that GROUP BY is empty).
			value, err := proj.evaluator.Eval(input)
			if err != nil {
				return err
			}
			if err := assignOutputValue(result, proj.alias, value); err != nil {
				return err
			}
		}
		output = append(output, result)
		return nil
	}

	// Note: `ep.buffers` is a map, so iterating over its keys may yield
	// different results in every run of the program. We cannot expect
	// a consistent order in which evalItem is run on the items of the
	// cartesian product.
	allStreams := make([]string, 0, len(ep.buffers))
	for key := range ep.buffers {
		allStreams = append(allStreams, key)
	}
	if err := ep.processCartesianProduct(dataHolder, allStreams, evalItem); err != nil {
		rollback()
		return err
	}

	// if we arrive here, then the input for the aggregation functions
	// is in the `group` list and we need to compute aggregation and output.
	// TODO deal with the case of an empty list
	for _, group := range groups {
		if err := evalGroup(&group); err != nil {
			rollback()
			return err
		}
	}
	if len(groups) == 0 {
		if err := evalNoGroup(); err != nil {
			rollback()
			return err
		}
	}

	ep.curResults = output
	return nil
}
<commit_msg>use a map lookup for grouping, not slice+DeepEqual<commit_after>package execution

import (
	"pfi/sensorbee/sensorbee/bql/udf"
	"pfi/sensorbee/sensorbee/core"
	"pfi/sensorbee/sensorbee/data"
)

type groupbyExecutionPlan struct {
	streamRelationStreamExecutionPlan
}

// tmpGroupData is an intermediate data structure to represent
// a set of rows that have the same values for GROUP BY columns.
type tmpGroupData struct {
	// this is the group (e.g. [1, "toy"]), where the values are
	// in order of the items in the GROUP BY clause
	group data.Array
	// for each aggregate function, we hold an array with the
	// input values.
	aggData map[string][]data.Value
	// as per our assumptions about grouping, the non-aggregation
	// data should be identical within every group
	nonAggData data.Map
}

// CanBuildGroupbyExecutionPlan checks whether the given statement
// allows to use an groupbyExecutionPlan.
func CanBuildGroupbyExecutionPlan(lp *LogicalPlan, reg udf.FunctionRegistry) bool {
	return lp.GroupingStmt
}

// groupbyExecutionPlan is a simple plan that follows the
// theoretical processing model. It supports only statements
// that use aggregation.
//
// After each tuple arrives,
// - compute the contents of the current window using the
//   specified window size/type,
// - perform a SELECT query on that data,
// - compute the data that need to be emitted by comparison with
//   the previous run's results.
func NewGroupbyExecutionPlan(lp *LogicalPlan, reg udf.FunctionRegistry) (ExecutionPlan, error) {
	underlying, err := newStreamRelationStreamExecutionPlan(lp, reg)
	if err != nil {
		return nil, err
	}
	return &groupbyExecutionPlan{
		*underlying,
	}, nil
}

// Process takes an input tuple and returns a slice of Map values that
// correspond to the results of the query represented by this execution
// plan. Note that the order of items in the returned slice is undefined
// and cannot be relied on.
func (ep *groupbyExecutionPlan) Process(input *core.Tuple) ([]data.Map, error) {
	return ep.process(input, ep.performQueryOnBuffer)
}

// performQueryOnBuffer executes a SELECT query on the data of the tuples
// currently stored in the buffer. The query results (which is a set of
// data.Value, not core.Tuple) is stored in ep.curResults. The data
// that was stored in ep.curResults before this method was called is
// moved to ep.prevResults. Note that the order of values in ep.curResults
// is undefined.
//
// In case of an error the contents of ep.curResults will still be
// the same as before the call (so that the next run performs as
// if no error had happened), but the contents of ep.curResults are
// undefined.
//
// Currently performQueryOnBuffer can only perform SELECT ... WHERE ...
// queries without aggregate functions, GROUP BY, or HAVING clauses.
func (ep *groupbyExecutionPlan) performQueryOnBuffer() error {
	// reuse the allocated memory
	output := ep.prevResults[0:0]
	// remember the previous results
	ep.prevResults = ep.curResults

	rollback := func() {
		// NB. ep.prevResults currently points to an slice with
		//     results from the previous run. ep.curResults points
		//     to the same slice. output points to a different slice
		//     with a different underlying array.
		//     in the next run, output will be reusing the underlying
		//     storage of the current ep.prevResults to hold results.
		//     therefore when we leave this function we must make
		//     sure that ep.prevResults and ep.curResults have
		//     different underlying arrays or ISTREAM/DSTREAM will
		//     return wrong results.
		ep.prevResults = output
	}

	// collect a list of all aggregate parameter evaluators in all
	// projections. this is necessary to avoid duplicate evaluation
	// if the same parameter is used in multiple aggregation funcs.
	allAggEvaluators := map[string]Evaluator{}
	for _, proj := range ep.projections {
		for key, agg := range proj.aggrEvals {
			allAggEvaluators[key] = agg
		}
	}

	// groups holds one item for every combination of values that
	// appear in the GROUP BY clause
	groups := map[data.HashValue]*tmpGroupData{}
	// we also keep a list of group keys so that we can still loop
	// over them in the order they were added
	groupKeys := []data.HashValue{}

	// findOrCreateGroup looks up the group that has the given
	// groupValues in the `groups`map. if there is no such
	// group, a new one is created and a copy of the given map
	// is used as a representative of this group's values.
	findOrCreateGroup := func(groupValues []data.Value, d data.Map) (*tmpGroupData, error) {
		groupHash := data.Hash(data.Array(groupValues))
		// find the correct group
		group, exists := groups[groupHash]
		// if there is no such group, create one
		if !exists {
			newGroup := &tmpGroupData{
				// the values that make up this group
				groupValues,
				// the input values of the aggregate functions
				map[string][]data.Value{},
				// a representative set of values for this group for later evaluation
				// TODO actually we don't need the whole map,
				//      just the parts common to the whole group
				d.Copy(),
			}
			// initialize the map with the aggregate function inputs
			for _, proj := range ep.projections {
				for key := range proj.aggrEvals {
					newGroup.aggData[key] = make([]data.Value, 0, 1)
				}
			}
			groups[groupHash] = newGroup
			group = newGroup
			groupKeys = append(groupKeys, groupHash)
		}

		// return a pointer to the (found or created) group
		return group, nil
	}

	// we need to make a cross product of the data in all buffers,
	// combine it to get an input like
	//  {"streamA": {data}, "streamB": {data}, "streamC": {data}}
	// and then run filter/projections on each of this items

	dataHolder := data.Map{}

	// function to evaluate filter on the input data and do the computations
	// that are required on each input tuple. (those computations differ
	// depending on whether we are in grouping mode or not.)
	evalItem := func(d data.Map) error {
		// evaluate filter condition and convert to bool
		if ep.filter != nil {
			filterResult, err := ep.filter.Eval(d)
			if err != nil {
				return err
			}
			filterResultBool, err := data.ToBool(filterResult)
			if err != nil {
				return err
			}
			// if it evaluated to false, do not further process this tuple
			// (ToBool also evalutes the NULL value to false, so we don't
			// need to treat this specially)
			if !filterResultBool {
				return nil
			}
		}

		// now compute the expressions in the GROUP BY to find the correct
		// group to append to
		// TODO there is actually no need to allocate this array again
		//      and again, or even have an array since we can update the
		//      group's hash value incrementally
		itemGroupValues := make([]data.Value, len(ep.groupList))
		for i, eval := range ep.groupList {
			// ordinary "flat" expression
			value, err := eval.Eval(d)
			if err != nil {
				return err
			}
			itemGroupValues[i] = value
		}

		itemGroup, err := findOrCreateGroup(itemGroupValues, d)
		if err != nil {
			return err
		}

		// now compute all the input data for the aggregate functions,
		// e.g. for `SELECT count(a) + max(b/2)`, compute `a` and `b/2`
		for key, agg := range allAggEvaluators {
			value, err := agg.Eval(d)
			if err != nil {
				return err
			}
			// store this value in the output map
			itemGroup.aggData[key] = append(itemGroup.aggData[key], value)
		}
		return nil
	}

	evalGroup := func(group *tmpGroupData) error {
		result := data.Map(make(map[string]data.Value, len(ep.projections)))
		// collect input for aggregate functions into an array
		// within each group
		for key, _ := range allAggEvaluators {
			group.nonAggData[key] = data.Array(group.aggData[key])
			delete(group.aggData, key)
		}
		// evaluate HAVING condition, if there is one
		for _, proj := range ep.projections {
			if proj.alias == ":having:" {
				havingResult, err := proj.evaluator.Eval(group.nonAggData)
				if err != nil {
					return err
				}
				havingResultBool, err := data.ToBool(havingResult)
				if err != nil {
					return err
				}
				// if it evaluated to false, do not further process this group
				// (ToBool also evalutes the NULL value to false, so we don't
				// need to treat this specially)
				if !havingResultBool {
					return nil
				}
				break
			}
		}
		// now evaluate all other projections
		for _, proj := range ep.projections {
			if proj.alias == ":having:" {
				continue
			}
			// now evaluate this projection on the flattened data
			value, err := proj.evaluator.Eval(group.nonAggData)
			if err != nil {
				return err
			}
			if err := assignOutputValue(result, proj.alias, value); err != nil {
				return err
			}
		}
		output = append(output, result)
		return nil
	}

	evalNoGroup := func() error {
		// if we have an empty group list *and* a GROUP BY clause,
		// we have to return an empty result (because there are no
		// rows with "the same values"). but if the list is empty and
		// we *don't* have a GROUP BY clause, then we need to compute
		// all foldables and aggregates with an empty input
		if len(ep.groupList) > 0 {
			return nil
		}
		input := data.Map{}
		result := data.Map(make(map[string]data.Value, len(ep.projections)))
		for _, proj := range ep.projections {
			// collect input for aggregate functions
			if proj.hasAggregate {
				for key := range proj.aggrEvals {
					input[key] = data.Array{}
				}
			}
			// now evaluate this projection on the flattened data.
			// note that input has *only* the keys of the empty
			// arrays, no other columns, but we cannot have other
			// columns involved in the projection (since we know
			// that GROUP BY is empty).
			value, err := proj.evaluator.Eval(input)
			if err != nil {
				return err
			}
			if err := assignOutputValue(result, proj.alias, value); err != nil {
				return err
			}
		}
		output = append(output, result)
		return nil
	}

	// Note: `ep.buffers` is a map, so iterating over its keys may yield
	// different results in every run of the program. We cannot expect
	// a consistent order in which evalItem is run on the items of the
	// cartesian product.
	allStreams := make([]string, 0, len(ep.buffers))
	for key := range ep.buffers {
		allStreams = append(allStreams, key)
	}
	if err := ep.processCartesianProduct(dataHolder, allStreams, evalItem); err != nil {
		rollback()
		return err
	}

	// if we arrive here, then the input for the aggregation functions
	// is in the `group` list and we need to compute aggregation and output.
	// NB. we do not directly loop over the `groups` map to avoid random order.
	for _, groupKey := range groupKeys {
		group := groups[groupKey]
		if err := evalGroup(group); err != nil {
			rollback()
			return err
		}
	}
	if len(groups) == 0 {
		if err := evalNoGroup(); err != nil {
			rollback()
			return err
		}
	}

	ep.curResults = output
	return nil
}
<|endoftext|>


<commit_before>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type HardwareConfig

package common

import (
	"context"
	"fmt"

	"github.com/hashicorp/packer/builder/vsphere/driver"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

type HardwareConfig struct {
	// Number of CPU sockets.
	CPUs int32 `mapstructure:"CPUs"`
	// Number of CPU cores per socket.
	CpuCores int32 `mapstructure:"cpu_cores"`
	// Amount of reserved CPU resources in MHz.
	CPUReservation int64 `mapstructure:"CPU_reservation"`
	// Upper limit of available CPU resources in MHz.
	CPULimit int64 `mapstructure:"CPU_limit"`
	// Enable CPU hot plug setting for virtual machine. Defaults to `false`.
	CpuHotAddEnabled bool `mapstructure:"CPU_hot_plug"`
	// Amount of RAM in MB.
	RAM int64 `mapstructure:"RAM"`
	// Amount of reserved RAM in MB.
	RAMReservation int64 `mapstructure:"RAM_reservation"`
	// Reserve all available RAM. Defaults to `false`. Cannot be used together
	// with `RAM_reservation`.
	RAMReserveAll bool `mapstructure:"RAM_reserve_all"`
	// Enable RAM hot plug setting for virtual machine. Defaults to `false`.
	MemoryHotAddEnabled bool `mapstructure:"RAM_hot_plug"`
	// Amount of video memory in KB.
	VideoRAM int64 `mapstructure:"video_ram"`
	// vGPU profile for accelerated graphics. See [NVIDIA GRID vGPU documentation](https://docs.nvidia.com/grid/latest/grid-vgpu-user-guide/index.html#configure-vmware-vsphere-vm-with-vgpu)
	// for examples of profile names. Defaults to none.
	VGPUProfile string `mapstructure:"vgpu_profile"`
	// Enable nested hardware virtualization for VM. Defaults to `false`.
	NestedHV bool `mapstructure:"NestedHV"`
	// Set the Firmware for virtual machine. Supported values: `bios`, `efi` or `efi-secure`. Defaults to `bios`.
	Firmware string `mapstructure:"firmware"`
	// During the boot, force entry into the BIOS setup screen. Defaults to `false`.
	ForceBIOSSetup bool `mapstructure:"force_bios_setup"`
}

func (c *HardwareConfig) Prepare() []error {
	var errs []error

	if c.RAMReservation > 0 && c.RAMReserveAll != false {
		errs = append(errs, fmt.Errorf("'RAM_reservation' and 'RAM_reserve_all' cannot be used together"))
	}

	if c.Firmware != "" && c.Firmware != "bios" && c.Firmware != "efi" && c.Firmware != "efi-secure" {
		errs = append(errs, fmt.Errorf("'firmware' must be '', 'bios', 'efi' or 'efi-secure'"))
	}

	return errs
}

type StepConfigureHardware struct {
	Config *HardwareConfig
}

func (s *StepConfigureHardware) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)
	vm := state.Get("vm").(*driver.VirtualMachine)

	if *s.Config != (HardwareConfig{}) {
		ui.Say("Customizing hardware...")

		err := vm.Configure(&driver.HardwareConfig{
			CPUs:                s.Config.CPUs,
			CpuCores:            s.Config.CpuCores,
			CPUReservation:      s.Config.CPUReservation,
			CPULimit:            s.Config.CPULimit,
			RAM:                 s.Config.RAM,
			RAMReservation:      s.Config.RAMReservation,
			RAMReserveAll:       s.Config.RAMReserveAll,
			NestedHV:            s.Config.NestedHV,
			CpuHotAddEnabled:    s.Config.CpuHotAddEnabled,
			MemoryHotAddEnabled: s.Config.MemoryHotAddEnabled,
			VideoRAM:            s.Config.VideoRAM,
			VGPUProfile:         s.Config.VGPUProfile,
			Firmware:            s.Config.Firmware,
			ForceBIOSSetup:      s.Config.ForceBIOSSetup,
		})
		if err != nil {
			state.Put("error", err)
			return multistep.ActionHalt
		}
	}

	return multistep.ActionContinue
}

func (s *StepConfigureHardware) Cleanup(multistep.StateBag) {}
<commit_msg>fix struct comment for generation<commit_after>//go:generate struct-markdown
//go:generate mapstructure-to-hcl2 -type HardwareConfig

package common

import (
	"context"
	"fmt"

	"github.com/hashicorp/packer/builder/vsphere/driver"
	"github.com/hashicorp/packer/helper/multistep"
	"github.com/hashicorp/packer/packer"
)

type HardwareConfig struct {
	// Number of CPU cores.
	CPUs int32 `mapstructure:"CPUs"`
	// Number of CPU cores per socket.
	CpuCores int32 `mapstructure:"cpu_cores"`
	// Amount of reserved CPU resources in MHz.
	CPUReservation int64 `mapstructure:"CPU_reservation"`
	// Upper limit of available CPU resources in MHz.
	CPULimit int64 `mapstructure:"CPU_limit"`
	// Enable CPU hot plug setting for virtual machine. Defaults to `false`.
	CpuHotAddEnabled bool `mapstructure:"CPU_hot_plug"`
	// Amount of RAM in MB.
	RAM int64 `mapstructure:"RAM"`
	// Amount of reserved RAM in MB.
	RAMReservation int64 `mapstructure:"RAM_reservation"`
	// Reserve all available RAM. Defaults to `false`. Cannot be used together
	// with `RAM_reservation`.
	RAMReserveAll bool `mapstructure:"RAM_reserve_all"`
	// Enable RAM hot plug setting for virtual machine. Defaults to `false`.
	MemoryHotAddEnabled bool `mapstructure:"RAM_hot_plug"`
	// Amount of video memory in KB.
	VideoRAM int64 `mapstructure:"video_ram"`
	// vGPU profile for accelerated graphics. See [NVIDIA GRID vGPU documentation](https://docs.nvidia.com/grid/latest/grid-vgpu-user-guide/index.html#configure-vmware-vsphere-vm-with-vgpu)
	// for examples of profile names. Defaults to none.
	VGPUProfile string `mapstructure:"vgpu_profile"`
	// Enable nested hardware virtualization for VM. Defaults to `false`.
	NestedHV bool `mapstructure:"NestedHV"`
	// Set the Firmware for virtual machine. Supported values: `bios`, `efi` or `efi-secure`. Defaults to `bios`.
	Firmware string `mapstructure:"firmware"`
	// During the boot, force entry into the BIOS setup screen. Defaults to `false`.
	ForceBIOSSetup bool `mapstructure:"force_bios_setup"`
}

func (c *HardwareConfig) Prepare() []error {
	var errs []error

	if c.RAMReservation > 0 && c.RAMReserveAll != false {
		errs = append(errs, fmt.Errorf("'RAM_reservation' and 'RAM_reserve_all' cannot be used together"))
	}

	if c.Firmware != "" && c.Firmware != "bios" && c.Firmware != "efi" && c.Firmware != "efi-secure" {
		errs = append(errs, fmt.Errorf("'firmware' must be '', 'bios', 'efi' or 'efi-secure'"))
	}

	return errs
}

type StepConfigureHardware struct {
	Config *HardwareConfig
}

func (s *StepConfigureHardware) Run(_ context.Context, state multistep.StateBag) multistep.StepAction {
	ui := state.Get("ui").(packer.Ui)
	vm := state.Get("vm").(*driver.VirtualMachine)

	if *s.Config != (HardwareConfig{}) {
		ui.Say("Customizing hardware...")

		err := vm.Configure(&driver.HardwareConfig{
			CPUs:                s.Config.CPUs,
			CpuCores:            s.Config.CpuCores,
			CPUReservation:      s.Config.CPUReservation,
			CPULimit:            s.Config.CPULimit,
			RAM:                 s.Config.RAM,
			RAMReservation:      s.Config.RAMReservation,
			RAMReserveAll:       s.Config.RAMReserveAll,
			NestedHV:            s.Config.NestedHV,
			CpuHotAddEnabled:    s.Config.CpuHotAddEnabled,
			MemoryHotAddEnabled: s.Config.MemoryHotAddEnabled,
			VideoRAM:            s.Config.VideoRAM,
			VGPUProfile:         s.Config.VGPUProfile,
			Firmware:            s.Config.Firmware,
			ForceBIOSSetup:      s.Config.ForceBIOSSetup,
		})
		if err != nil {
			state.Put("error", err)
			return multistep.ActionHalt
		}
	}

	return multistep.ActionContinue
}

func (s *StepConfigureHardware) Cleanup(multistep.StateBag) {}
<|endoftext|>


<commit_before>package queue

import (
	"github.com/nabeken/aws-go-sqs/queue/option"

	"github.com/stripe/aws-go/aws"
	"github.com/stripe/aws-go/gen/sqs"
)

type Queue struct {
	*sqs.SQS
	URL aws.StringValue
}

func New(s *sqs.SQS, name string) (*Queue, error) {
	u, err := GetQueueURL(s, name)
	if err != nil {
		return nil, err
	}

	return &Queue{
		SQS: s,
		URL: u,
	}, nil
}

func (q *Queue) SendMessage(body string, opts ...option.SendMessageRequest) error {
	req := &sqs.SendMessageRequest{
		MessageBody: aws.String(body),
		QueueURL:    q.URL,
	}

	for _, f := range opts {
		f(req)
	}

	_, err := q.SQS.SendMessage(req)
	return err
}

func (q *Queue) ReceiveMessage(opts ...option.ReceiveMessageRequest) ([]sqs.Message, error) {
	req := &sqs.ReceiveMessageRequest{
		QueueURL: q.URL,
	}

	for _, f := range opts {
		f(req)
	}

	resp, err := q.SQS.ReceiveMessage(req)
	if err != nil {
		return nil, err
	}
	return resp.Messages, nil
}

func (q *Queue) DeleteMessage(receiptHandle aws.StringValue) error {
	return q.SQS.DeleteMessage(&sqs.DeleteMessageRequest{
		QueueURL:      q.URL,
		ReceiptHandle: receiptHandle,
	})
}

func GetQueueURL(s *sqs.SQS, name string) (aws.StringValue, error) {
	req := &sqs.GetQueueURLRequest{
		QueueName: aws.String(name),
	}

	resp, err := s.GetQueueURL(req)
	if err != nil {
		return nil, err
	}
	return resp.QueueURL, nil
}
<commit_msg>queue: Add PurgeQueue()<commit_after>package queue

import (
	"github.com/nabeken/aws-go-sqs/queue/option"

	"github.com/stripe/aws-go/aws"
	"github.com/stripe/aws-go/gen/sqs"
)

type Queue struct {
	*sqs.SQS
	URL aws.StringValue
}

func New(s *sqs.SQS, name string) (*Queue, error) {
	u, err := GetQueueURL(s, name)
	if err != nil {
		return nil, err
	}

	return &Queue{
		SQS: s,
		URL: u,
	}, nil
}

func (q *Queue) SendMessage(body string, opts ...option.SendMessageRequest) error {
	req := &sqs.SendMessageRequest{
		MessageBody: aws.String(body),
		QueueURL:    q.URL,
	}

	for _, f := range opts {
		f(req)
	}

	_, err := q.SQS.SendMessage(req)
	return err
}

func (q *Queue) ReceiveMessage(opts ...option.ReceiveMessageRequest) ([]sqs.Message, error) {
	req := &sqs.ReceiveMessageRequest{
		QueueURL: q.URL,
	}

	for _, f := range opts {
		f(req)
	}

	resp, err := q.SQS.ReceiveMessage(req)
	if err != nil {
		return nil, err
	}
	return resp.Messages, nil
}

func (q *Queue) DeleteMessage(receiptHandle aws.StringValue) error {
	return q.SQS.DeleteMessage(&sqs.DeleteMessageRequest{
		QueueURL:      q.URL,
		ReceiptHandle: receiptHandle,
	})
}

func (q *Queue) PurgeQueue() error {
	return q.SQS.PurgeQueue(&sqs.PurgeQueueRequest{
		QueueURL: q.URL,
	})
}

func GetQueueURL(s *sqs.SQS, name string) (aws.StringValue, error) {
	req := &sqs.GetQueueURLRequest{
		QueueName: aws.String(name),
	}

	resp, err := s.GetQueueURL(req)
	if err != nil {
		return nil, err
	}
	return resp.QueueURL, nil
}
<|endoftext|>


<commit_before>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package queue

import (
	"encoding/gob"
	"errors"
	"io"
	"net"
	"sync/atomic"
	"time"
)

// The size of buffered channels created by ChannelFromWriter.
const ChanSize = 32

// Message represents the message stored in the queue.
//
// A message is specified by an action and a slice of strings, representing
// arguments to the action.
//
// For example, the action "regenerate apprc" could receive one argument: the
// name of the app for which the apprc file will be regenerate.
type Message struct {
	Action string
	Args   []string
	Visits int
}

// ChannelFromWriter returns a channel from a given io.WriteCloser.
//
// Every time a Message is sent to the channel, it gets written to the writer
// in gob format.  ChannelFromWriter also returns a channel for errors in
// writtings. You can use a select for error checking:
//
//     ch, errCh := ChannelFromWriter(w)
//     // use ch
//     select {
//     case err := <-errCh:
//         // threat the error
//     case time.After(5e9):
//         // no error after 5 seconds
//     }
//
// Please notice that there is no deadline for the writting. You can obviously
// ignore errors, if they are not significant for you.
//
// Whenever you close the message channel (and you should, to make it clear
// that you will not send any messages to the channel anymore), the error
// channel will get automatically closed, so the WriteCloser.
//
// Both channels are buffered by ChanSize.
func ChannelFromWriter(w io.WriteCloser) (chan<- Message, <-chan error) {
	msgChan := make(chan Message, ChanSize)
	errChan := make(chan error, ChanSize)
	go write(w, msgChan, errChan)
	return msgChan, errChan
}

// write reads messages from ch and write them to w, in gob format.
//
// If clients close ch, write will close errCh.
func write(w io.WriteCloser, ch <-chan Message, errCh chan<- error) {
	defer close(errCh)
	defer w.Close()
	encoder := gob.NewEncoder(w)
	for msg := range ch {
		if err := encoder.Encode(msg); err != nil {
			errCh <- err
		}
	}
}

// Server is the server that hosts the queue. It receives messages and
// process them.
type Server struct {
	listener net.Listener
	pairs    chan pair
	closed   int32
}

// StartServer starts a new queue server from a local address.
//
// The address must be a TCP address, in the format host:port (for example,
// [::1]:8080 or 192.168.254.10:2020).
func StartServer(laddr string) (*Server, error) {
	var (
		server Server
		err    error
	)
	server.listener, err = net.Listen("tcp", laddr)
	if err != nil {
		return nil, errors.New("Could not start server: " + err.Error())
	}
	server.pairs = make(chan pair, ChanSize)
	go server.loop()
	return &server, nil
}

// handle handles a new client, sending errors to the qs.errors channel and
// received messages to qs.messages.
func (qs *Server) handle(conn net.Conn) {
	var err error
	decoder := gob.NewDecoder(conn)
	for err == nil {
		var msg Message
		err = decoder.Decode(&msg)
		if atomic.LoadInt32(&qs.closed) == 0 {
			qs.pairs <- pair{message: msg, err: err}
		}
	}
}

// loop accepts connection forever, and uses read to read messages from it,
// decoding them to a channel of messages.
func (qs *Server) loop() {
	for atomic.LoadInt32(&qs.closed) == 0 {
		conn, err := qs.listener.Accept()
		if err != nil {
			if e, ok := err.(*net.OpError); ok && !e.Temporary() {
				return
			}
			continue
		}
		go qs.handle(conn)
	}
}

// Message returns the first available message in the queue, or an error if it
// fails to read the message, or times out while waiting for the message.
//
// If timeout is negative, this method will wait nearly forever for the
// arriving of a message or an error.
func (qs *Server) Message(timeout time.Duration) (Message, error) {
	var (
		msg Message
		err error
	)
	if timeout < 0 {
		timeout = 1 << 62
	}
	select {
	case pair := <-qs.pairs:
		if pair.err == io.EOF {
			pair.err = errors.New("EOF: client disconnected.")
		}
		msg = pair.message
		err = pair.err
	case <-time.After(timeout):
		err = errors.New("Timed out waiting for the message.")
	}
	return msg, err
}

// PutBack puts a message back in the queue. It should be used when a message
// got using Message method cannot be processed yet. You put it back in the
// queue for processing later.
func (qs *Server) PutBack(message Message) {
	if atomic.LoadInt32(&qs.closed) == 0 {
		message.Visits++
		qs.pairs <- pair{message: message}
	}
}

// Addr returns the address of the server.
func (qs *Server) Addr() string {
	return qs.listener.Addr().String()
}

// Close closes the server, closing the underlying listener.
func (qs *Server) Close() error {
	if !atomic.CompareAndSwapInt32(&qs.closed, 0, 1) {
		return errors.New("Server already closed.")
	}
	err := qs.listener.Close()
	qs.pairs <- pair{err: errors.New("Server is closed.")}
	return err
}

// Dial is used to connect to a queue server.
//
// It returns three values: the channel to which messages should be sent, the
// channel where the client will get errors from the server during writing of
// messages and an error, that will be non-nil in case of failure to connect to
// the queue server.
//
// Whenever the message channel gets closed, the connection with the remote
// server will be closed.
func Dial(addr string) (chan<- Message, <-chan error, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, nil, errors.New("Could not dial to " + addr + ": " + err.Error())
	}
	messages, errors := ChannelFromWriter(conn)
	return messages, errors, nil
}

// pair is a pair of message and error.
type pair struct {
	message Message
	err     error
}
<commit_msg>queue: using a channel to prevent Message from leaking<commit_after>// Copyright 2012 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package queue

import (
	"encoding/gob"
	"errors"
	"io"
	"net"
	"sync/atomic"
	"time"
)

// The size of buffered channels created by ChannelFromWriter.
const ChanSize = 32

// Message represents the message stored in the queue.
//
// A message is specified by an action and a slice of strings, representing
// arguments to the action.
//
// For example, the action "regenerate apprc" could receive one argument: the
// name of the app for which the apprc file will be regenerate.
type Message struct {
	Action string
	Args   []string
	Visits int
}

// ChannelFromWriter returns a channel from a given io.WriteCloser.
//
// Every time a Message is sent to the channel, it gets written to the writer
// in gob format.  ChannelFromWriter also returns a channel for errors in
// writtings. You can use a select for error checking:
//
//     ch, errCh := ChannelFromWriter(w)
//     // use ch
//     select {
//     case err := <-errCh:
//         // threat the error
//     case time.After(5e9):
//         // no error after 5 seconds
//     }
//
// Please notice that there is no deadline for the writting. You can obviously
// ignore errors, if they are not significant for you.
//
// Whenever you close the message channel (and you should, to make it clear
// that you will not send any messages to the channel anymore), the error
// channel will get automatically closed, so the WriteCloser.
//
// Both channels are buffered by ChanSize.
func ChannelFromWriter(w io.WriteCloser) (chan<- Message, <-chan error) {
	msgChan := make(chan Message, ChanSize)
	errChan := make(chan error, ChanSize)
	go write(w, msgChan, errChan)
	return msgChan, errChan
}

// write reads messages from ch and write them to w, in gob format.
//
// If clients close ch, write will close errCh.
func write(w io.WriteCloser, ch <-chan Message, errCh chan<- error) {
	defer close(errCh)
	defer w.Close()
	encoder := gob.NewEncoder(w)
	for msg := range ch {
		if err := encoder.Encode(msg); err != nil {
			errCh <- err
		}
	}
}

// Server is the server that hosts the queue. It receives messages and
// process them.
type Server struct {
	listener net.Listener
	pairs    chan pair
	close    chan int
	closed   int32
}

// StartServer starts a new queue server from a local address.
//
// The address must be a TCP address, in the format host:port (for example,
// [::1]:8080 or 192.168.254.10:2020).
func StartServer(laddr string) (*Server, error) {
	var (
		server Server
		err    error
	)
	server.listener, err = net.Listen("tcp", laddr)
	if err != nil {
		return nil, errors.New("Could not start server: " + err.Error())
	}
	server.pairs = make(chan pair, ChanSize)
	server.close = make(chan int, 1)
	go server.loop()
	return &server, nil
}

// handle handles a new client, sending errors to the qs.errors channel and
// received messages to qs.messages.
func (qs *Server) handle(conn net.Conn) {
	var err error
	decoder := gob.NewDecoder(conn)
	for err == nil {
		var msg Message
		err = decoder.Decode(&msg)
		if atomic.LoadInt32(&qs.closed) == 0 {
			qs.pairs <- pair{message: msg, err: err}
		}
	}
}

// loop accepts connection forever, and uses read to read messages from it,
// decoding them to a channel of messages.
func (qs *Server) loop() {
	for atomic.LoadInt32(&qs.closed) == 0 {
		conn, err := qs.listener.Accept()
		if err != nil {
			if e, ok := err.(*net.OpError); ok && !e.Temporary() {
				return
			}
			continue
		}
		go qs.handle(conn)
	}
}

// Message returns the first available message in the queue, or an error if it
// fails to read the message, or times out while waiting for the message.
//
// If timeout is negative, this method will wait nearly forever for the
// arriving of a message or an error.
func (qs *Server) Message(timeout time.Duration) (Message, error) {
	var (
		msg Message
		err error
	)
	if timeout < 0 {
		timeout = 1 << 62
	}
	select {
	case pair := <-qs.pairs:
		if pair.err == io.EOF {
			pair.err = errors.New("EOF: client disconnected.")
		}
		msg = pair.message
		err = pair.err
	case <-qs.close:
	case <-time.After(timeout):
		err = errors.New("Timed out waiting for the message.")
	}
	return msg, err
}

// PutBack puts a message back in the queue. It should be used when a message
// got using Message method cannot be processed yet. You put it back in the
// queue for processing later.
func (qs *Server) PutBack(message Message) {
	if atomic.LoadInt32(&qs.closed) == 0 {
		message.Visits++
		qs.pairs <- pair{message: message}
	}
}

// Addr returns the address of the server.
func (qs *Server) Addr() string {
	return qs.listener.Addr().String()
}

// Close closes the server, closing the underlying listener.
func (qs *Server) Close() error {
	if !atomic.CompareAndSwapInt32(&qs.closed, 0, 1) {
		return errors.New("Server already closed.")
	}
	err := qs.listener.Close()
	qs.pairs <- pair{err: errors.New("Server is closed.")}
	qs.close <- 1
	return err
}

// Dial is used to connect to a queue server.
//
// It returns three values: the channel to which messages should be sent, the
// channel where the client will get errors from the server during writing of
// messages and an error, that will be non-nil in case of failure to connect to
// the queue server.
//
// Whenever the message channel gets closed, the connection with the remote
// server will be closed.
func Dial(addr string) (chan<- Message, <-chan error, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, nil, errors.New("Could not dial to " + addr + ": " + err.Error())
	}
	messages, errors := ChannelFromWriter(conn)
	return messages, errors, nil
}

// pair is a pair of message and error.
type pair struct {
	message Message
	err     error
}
<|endoftext|>