<commit_before>package detailed

import (
	"fmt"
	"strings"

	"github.com/weaveworks/scope/probe/docker"
	"github.com/weaveworks/scope/probe/endpoint"
	"github.com/weaveworks/scope/probe/host"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/scope/probe/process"
	"github.com/weaveworks/scope/render"
	"github.com/weaveworks/scope/report"
)

// Shapes that are allowed
const (
	ImageNameNone = "<none>"

	// Keys we use to render container names
	AmazonECSContainerNameLabel  = "com.amazonaws.ecs.container-name"
	KubernetesContainerNameLabel = "io.kubernetes.container.name"
	MarathonAppIDEnv             = "MARATHON_APP_ID"
)

// NodeSummaryGroup is a topology-typed group of children for a Node.
type NodeSummaryGroup struct {
	ID         string        `json:"id"`
	Label      string        `json:"label"`
	Nodes      []NodeSummary `json:"nodes"`
	TopologyID string        `json:"topologyId"`
	Columns    []Column      `json:"columns"`
}

// Copy returns a value copy of the NodeSummaryGroup
func (g NodeSummaryGroup) Copy() NodeSummaryGroup {
	result := NodeSummaryGroup{
		TopologyID: g.TopologyID,
		Label:      g.Label,
		Columns:    g.Columns,
	}
	for _, node := range g.Nodes {
		result.Nodes = append(result.Nodes, node.Copy())
	}
	return result
}

// Column provides special json serialization for column ids, so they include
// their label for the frontend.
type Column struct {
	ID          string `json:"id"`
	Label       string `json:"label"`
	DefaultSort bool   `json:"defaultSort"`
}

// NodeSummary is summary information about a child for a Node.
type NodeSummary struct {
	ID         string               `json:"id"`
	Label      string               `json:"label"`
	LabelMinor string               `json:"label_minor"`
	Rank       string               `json:"rank"`
	Shape      string               `json:"shape,omitempty"`
	Stack      bool                 `json:"stack,omitempty"`
	Linkable   bool                 `json:"linkable,omitempty"` // Whether this node can be linked-to
	Pseudo     bool                 `json:"pseudo,omitempty"`
	Metadata   []report.MetadataRow `json:"metadata,omitempty"`
	Parents    []Parent             `json:"parents,omitempty"`
	Metrics    []report.MetricRow   `json:"metrics,omitempty"`
	Tables     []report.Table       `json:"tables,omitempty"`
	Adjacency  report.IDList        `json:"adjacency,omitempty"`
}

// MakeNodeSummary summarizes a node, if possible.
func MakeNodeSummary(r report.Report, n report.Node) (NodeSummary, bool) {
	renderers := map[string]func(NodeSummary, report.Node) (NodeSummary, bool){
		render.Pseudo:         pseudoNodeSummary,
		report.Process:        processNodeSummary,
		report.Container:      containerNodeSummary,
		report.ContainerImage: containerImageNodeSummary,
		report.Pod:            podNodeSummary,
		report.Service:        serviceNodeSummary,
		report.Deployment:     deploymentNodeSummary,
		report.ReplicaSet:     replicaSetNodeSummary,
		report.Host:           hostNodeSummary,
	}
	if renderer, ok := renderers[n.Topology]; ok {
		return renderer(baseNodeSummary(r, n), n)
	}
	if strings.HasPrefix(n.Topology, "group:") {
		return groupNodeSummary(baseNodeSummary(r, n), r, n)
	}
	return NodeSummary{}, false
}

// SummarizeMetrics returns a copy of the NodeSummary where the metrics are
// replaced with their summaries
func (n NodeSummary) SummarizeMetrics() NodeSummary {
	cp := n.Copy()
	for i, m := range cp.Metrics {
		cp.Metrics[i] = m.Summary()
	}
	return cp
}

// Copy returns a value copy of the NodeSummary
func (n NodeSummary) Copy() NodeSummary {
	result := NodeSummary{
		ID:         n.ID,
		Label:      n.Label,
		LabelMinor: n.LabelMinor,
		Rank:       n.Rank,
		Shape:      n.Shape,
		Stack:      n.Stack,
		Linkable:   n.Linkable,
		Adjacency:  n.Adjacency.Copy(),
	}
	for _, row := range n.Metadata {
		result.Metadata = append(result.Metadata, row.Copy())
	}
	for _, table := range n.Tables {
		result.Tables = append(result.Tables, table.Copy())
	}
	for _, row := range n.Metrics {
		result.Metrics = append(result.Metrics, row.Copy())
	}
	return result
}

func baseNodeSummary(r report.Report, n report.Node) NodeSummary {
	t, _ := r.Topology(n.Topology)
	return NodeSummary{
		ID:        n.ID,
		Shape:     t.GetShape(),
		Linkable:  true,
		Metadata:  NodeMetadata(r, n),
		Metrics:   NodeMetrics(r, n),
		Parents:   Parents(r, n),
		Tables:    NodeTables(r, n),
		Adjacency: n.Adjacency.Copy(),
	}
}

func pseudoNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Pseudo = true
	base.Rank = n.ID

	if template, ok := map[string]struct{ Label, LabelMinor string }{
		render.TheInternetID:      {render.InboundMajor, ""},
		render.IncomingInternetID: {render.InboundMajor, render.InboundMinor},
		render.OutgoingInternetID: {render.OutboundMajor, render.OutboundMinor},
	}[n.ID]; ok {
		base.Label = template.Label
		base.LabelMinor = template.LabelMinor
		base.Shape = report.Cloud
		return base, true
	}

	// try rendering it as an uncontained node
	if strings.HasPrefix(n.ID, render.MakePseudoNodeID(render.UncontainedID)) {
		base.Label = render.UncontainedMajor
		base.LabelMinor = report.ExtractHostID(n)
		base.Shape = report.Square
		base.Stack = true
		return base, true
	}

	// try rendering it as an unmanaged node
	if strings.HasPrefix(n.ID, render.MakePseudoNodeID(render.UnmanagedID)) {
		base.Label = render.UnmanagedMajor
		base.Shape = report.Square
		base.Stack = true
		base.LabelMinor = report.ExtractHostID(n)
		return base, true
	}

	// try rendering it as an endpoint
	if addr, ok := n.Latest.Lookup(endpoint.Addr); ok {
		base.Label = addr
		base.Shape = report.Circle
		return base, true
	}

	return NodeSummary{}, false
}

func processNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label, _ = n.Latest.Lookup(process.Name)
	base.Rank, _ = n.Latest.Lookup(process.Name)

	pid, ok := n.Latest.Lookup(process.PID)
	if !ok {
		return NodeSummary{}, false
	}
	if containerName, ok := n.Latest.Lookup(docker.ContainerName); ok {
		base.LabelMinor = fmt.Sprintf("%s (%s:%s)", report.ExtractHostID(n), containerName, pid)
	} else {
		base.LabelMinor = fmt.Sprintf("%s (%s)", report.ExtractHostID(n), pid)
	}

	_, isConnected := n.Latest.Lookup(render.IsConnected)
	base.Linkable = isConnected
	return base, true
}

func containerNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label = getRenderableContainerName(n)
	base.LabelMinor = report.ExtractHostID(n)

	if imageName, ok := n.Latest.Lookup(docker.ImageName); ok {
		base.Rank = docker.ImageNameWithoutVersion(imageName)
	}

	return base, true
}

func containerImageNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	imageName, ok := n.Latest.Lookup(docker.ImageName)
	if !ok {
		return NodeSummary{}, false
	}

	imageNameWithoutVersion := docker.ImageNameWithoutVersion(imageName)
	base.Label = imageNameWithoutVersion
	base.Rank = imageNameWithoutVersion
	base.Stack = true

	if base.Label == ImageNameNone {
		base.Label, _ = n.Latest.Lookup(docker.ImageID)
		if len(base.Label) > 12 {
			base.Label = base.Label[:12]
		}
	}

	if c, ok := n.Counters.Lookup(report.Container); ok {
		if c == 1 {
			base.LabelMinor = fmt.Sprintf("%d container", c)
		} else {
			base.LabelMinor = fmt.Sprintf("%d containers", c)
		}
	}
	return base, true
}

func podNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label, _ = n.Latest.Lookup(kubernetes.Name)
	namespace, _ := n.Latest.Lookup(kubernetes.Namespace)
	base.Rank = namespace + "/" + base.Label
	if c, ok := n.Counters.Lookup(report.Container); ok {
		if c == 1 {
			base.LabelMinor = fmt.Sprintf("%d container", c)
		} else {
			base.LabelMinor = fmt.Sprintf("%d containers", c)
		}
	}

	return base, true
}

func serviceNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label, _ = n.Latest.Lookup(kubernetes.Name)
	namespace, _ := n.Latest.Lookup(kubernetes.Namespace)
	base.Rank = namespace + "/" + base.Label
	base.Stack = true

	// Services are always just a group of pods, so there's no counting multiple
	// services which might be grouped together.
	if p, ok := n.Counters.Lookup(report.Pod); ok {
		if p == 1 {
			base.LabelMinor = fmt.Sprintf("%d pod", p)
		} else {
			base.LabelMinor = fmt.Sprintf("%d pods", p)
		}
	}

	return base, true
}

func deploymentNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label, _ = n.Latest.Lookup(kubernetes.Name)
	namespace, _ := n.Latest.Lookup(kubernetes.Namespace)
	base.Rank = namespace + "/" + base.Label
	base.Stack = true

	if p, ok := n.Counters.Lookup(report.Pod); ok {
		if p == 1 {
			base.LabelMinor = fmt.Sprintf("%d pod", p)
		} else {
			base.LabelMinor = fmt.Sprintf("%d pods", p)
		}
	}

	return base, true
}

func replicaSetNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label, _ = n.Latest.Lookup(kubernetes.Name)
	namespace, _ := n.Latest.Lookup(kubernetes.Namespace)
	base.Rank = namespace + "/" + base.Label
	base.Stack = true

	if p, ok := n.Counters.Lookup(report.Pod); ok {
		if p == 1 {
			base.LabelMinor = fmt.Sprintf("%d pod", p)
		} else {
			base.LabelMinor = fmt.Sprintf("%d pods", p)
		}
	}

	return base, true
}

func hostNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	var (
		hostname, _ = n.Latest.Lookup(host.HostName)
		parts       = strings.SplitN(hostname, ".", 2)
	)

	if len(parts) == 2 {
		base.Label, base.LabelMinor, base.Rank = parts[0], parts[1], parts[1]
	} else {
		base.Label = hostname
	}

	return base, true
}

// groupNodeSummary renders the summary for a group node. n.Topology is
// expected to be of the form: group:container:hostname
func groupNodeSummary(base NodeSummary, r report.Report, n report.Node) (NodeSummary, bool) {
	parts := strings.Split(n.Topology, ":")
	if len(parts) != 3 {
		return NodeSummary{}, false
	}

	label, ok := n.Latest.Lookup(parts[2])
	if !ok {
		return NodeSummary{}, false
	}
	base.Label, base.Rank = label, label

	t, ok := r.Topology(parts[1])
	if ok && t.Label != "" {
		if count, ok := n.Counters.Lookup(parts[1]); ok {
			if count == 1 {
				base.LabelMinor = fmt.Sprintf("%d %s", count, t.Label)
			} else {
				base.LabelMinor = fmt.Sprintf("%d %s", count, t.LabelPlural)
			}
		}
	}

	base.Shape = t.GetShape()
	base.Stack = true
	return base, true
}

type nodeSummariesByID []NodeSummary

func (s nodeSummariesByID) Len() int           { return len(s) }
func (s nodeSummariesByID) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s nodeSummariesByID) Less(i, j int) bool { return s[i].ID < s[j].ID }

// NodeSummaries is a set of NodeSummaries indexed by ID.
type NodeSummaries map[string]NodeSummary

// Summaries converts RenderableNodes into a set of NodeSummaries
func Summaries(r report.Report, rns report.Nodes) NodeSummaries {
	result := NodeSummaries{}
	for id, node := range rns {
		if summary, ok := MakeNodeSummary(r, node); ok {
			for i, m := range summary.Metrics {
				summary.Metrics[i] = m.Summary()
			}
			result[id] = summary
		}
	}
	return result
}

// Copy returns a deep value-copy of NodeSummaries
func (n NodeSummaries) Copy() NodeSummaries {
	result := NodeSummaries{}
	for k, v := range n {
		result[k] = v.Copy()
	}
	return result
}

// getRenderableContainerName obtains a user-friendly container name, to render in the UI
func getRenderableContainerName(nmd report.Node) string {
	for _, key := range []string{
		// Amazon's ecs-agent produces huge Docker container names, destructively
		// derived from mangling Container Definition names in Task
		// Definitions.
		//
		// However, the ecs-agent provides a label containing the original Container
		// Definition name.
		docker.LabelPrefix + AmazonECSContainerNameLabel,
		// Kubernetes also mangles its Docker container names and provides a
		// label with the original container name. However, note that this label
		// is only provided by Kubernetes versions >= 1.2 (see
		// https://github.com/kubernetes/kubernetes/pull/17234/ )
		docker.LabelPrefix + KubernetesContainerNameLabel,
		// Marathon doesn't set any Docker labels and this is the only meaningful
		// attribute we can find to make Scope useful without Mesos plugin
		docker.EnvPrefix + MarathonAppIDEnv,
		docker.ContainerName,
		docker.ContainerHostname,
	} {
		if label, ok := nmd.Latest.Lookup(key); ok {
			return label
		}
	}
	return ""
}
<commit_msg>Review feedback<commit_after>package detailed

import (
	"fmt"
	"strings"

	"github.com/weaveworks/scope/probe/docker"
	"github.com/weaveworks/scope/probe/endpoint"
	"github.com/weaveworks/scope/probe/host"
	"github.com/weaveworks/scope/probe/kubernetes"
	"github.com/weaveworks/scope/probe/process"
	"github.com/weaveworks/scope/render"
	"github.com/weaveworks/scope/report"
)

// Shapes that are allowed
const (
	ImageNameNone = "<none>"

	// Keys we use to render container names
	AmazonECSContainerNameLabel  = "com.amazonaws.ecs.container-name"
	KubernetesContainerNameLabel = "io.kubernetes.container.name"
	MarathonAppIDEnv             = "MARATHON_APP_ID"
)

// NodeSummaryGroup is a topology-typed group of children for a Node.
type NodeSummaryGroup struct {
	ID         string        `json:"id"`
	Label      string        `json:"label"`
	Nodes      []NodeSummary `json:"nodes"`
	TopologyID string        `json:"topologyId"`
	Columns    []Column      `json:"columns"`
}

// Copy returns a value copy of the NodeSummaryGroup
func (g NodeSummaryGroup) Copy() NodeSummaryGroup {
	result := NodeSummaryGroup{
		TopologyID: g.TopologyID,
		Label:      g.Label,
		Columns:    g.Columns,
	}
	for _, node := range g.Nodes {
		result.Nodes = append(result.Nodes, node.Copy())
	}
	return result
}

// Column provides special json serialization for column ids, so they include
// their label for the frontend.
type Column struct {
	ID          string `json:"id"`
	Label       string `json:"label"`
	DefaultSort bool   `json:"defaultSort"`
}

// NodeSummary is summary information about a child for a Node.
type NodeSummary struct {
	ID         string               `json:"id"`
	Label      string               `json:"label"`
	LabelMinor string               `json:"label_minor"`
	Rank       string               `json:"rank"`
	Shape      string               `json:"shape,omitempty"`
	Stack      bool                 `json:"stack,omitempty"`
	Linkable   bool                 `json:"linkable,omitempty"` // Whether this node can be linked-to
	Pseudo     bool                 `json:"pseudo,omitempty"`
	Metadata   []report.MetadataRow `json:"metadata,omitempty"`
	Parents    []Parent             `json:"parents,omitempty"`
	Metrics    []report.MetricRow   `json:"metrics,omitempty"`
	Tables     []report.Table       `json:"tables,omitempty"`
	Adjacency  report.IDList        `json:"adjacency,omitempty"`
}

// MakeNodeSummary summarizes a node, if possible.
func MakeNodeSummary(r report.Report, n report.Node) (NodeSummary, bool) {
	renderers := map[string]func(NodeSummary, report.Node) (NodeSummary, bool){
		render.Pseudo:         pseudoNodeSummary,
		report.Process:        processNodeSummary,
		report.Container:      containerNodeSummary,
		report.ContainerImage: containerImageNodeSummary,
		report.Pod:            podNodeSummary,
		report.Service:        serviceNodeSummary,
		report.Deployment:     deploymentNodeSummary,
		report.ReplicaSet:     replicaSetNodeSummary,
		report.Host:           hostNodeSummary,
	}
	if renderer, ok := renderers[n.Topology]; ok {
		return renderer(baseNodeSummary(r, n), n)
	}
	if strings.HasPrefix(n.Topology, "group:") {
		return groupNodeSummary(baseNodeSummary(r, n), r, n)
	}
	return NodeSummary{}, false
}

// SummarizeMetrics returns a copy of the NodeSummary where the metrics are
// replaced with their summaries
func (n NodeSummary) SummarizeMetrics() NodeSummary {
	cp := n.Copy()
	for i, m := range cp.Metrics {
		cp.Metrics[i] = m.Summary()
	}
	return cp
}

// Copy returns a value copy of the NodeSummary
func (n NodeSummary) Copy() NodeSummary {
	result := NodeSummary{
		ID:         n.ID,
		Label:      n.Label,
		LabelMinor: n.LabelMinor,
		Rank:       n.Rank,
		Shape:      n.Shape,
		Stack:      n.Stack,
		Linkable:   n.Linkable,
		Adjacency:  n.Adjacency.Copy(),
	}
	for _, row := range n.Metadata {
		result.Metadata = append(result.Metadata, row.Copy())
	}
	for _, table := range n.Tables {
		result.Tables = append(result.Tables, table.Copy())
	}
	for _, row := range n.Metrics {
		result.Metrics = append(result.Metrics, row.Copy())
	}
	return result
}

func baseNodeSummary(r report.Report, n report.Node) NodeSummary {
	t, _ := r.Topology(n.Topology)
	return NodeSummary{
		ID:        n.ID,
		Shape:     t.GetShape(),
		Linkable:  true,
		Metadata:  NodeMetadata(r, n),
		Metrics:   NodeMetrics(r, n),
		Parents:   Parents(r, n),
		Tables:    NodeTables(r, n),
		Adjacency: n.Adjacency.Copy(),
	}
}

func pseudoNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Pseudo = true
	base.Rank = n.ID

	if template, ok := map[string]struct{ Label, LabelMinor string }{
		render.TheInternetID:      {render.InboundMajor, ""},
		render.IncomingInternetID: {render.InboundMajor, render.InboundMinor},
		render.OutgoingInternetID: {render.OutboundMajor, render.OutboundMinor},
	}[n.ID]; ok {
		base.Label = template.Label
		base.LabelMinor = template.LabelMinor
		base.Shape = report.Cloud
		return base, true
	}

	// try rendering it as an uncontained node
	if strings.HasPrefix(n.ID, render.MakePseudoNodeID(render.UncontainedID)) {
		base.Label = render.UncontainedMajor
		base.LabelMinor = report.ExtractHostID(n)
		base.Shape = report.Square
		base.Stack = true
		return base, true
	}

	// try rendering it as an unmanaged node
	if strings.HasPrefix(n.ID, render.MakePseudoNodeID(render.UnmanagedID)) {
		base.Label = render.UnmanagedMajor
		base.Shape = report.Square
		base.Stack = true
		base.LabelMinor = report.ExtractHostID(n)
		return base, true
	}

	// try rendering it as an endpoint
	if addr, ok := n.Latest.Lookup(endpoint.Addr); ok {
		base.Label = addr
		base.Shape = report.Circle
		return base, true
	}

	return NodeSummary{}, false
}

func processNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label, _ = n.Latest.Lookup(process.Name)
	base.Rank, _ = n.Latest.Lookup(process.Name)

	pid, ok := n.Latest.Lookup(process.PID)
	if !ok {
		return NodeSummary{}, false
	}
	if containerName, ok := n.Latest.Lookup(docker.ContainerName); ok {
		base.LabelMinor = fmt.Sprintf("%s (%s:%s)", report.ExtractHostID(n), containerName, pid)
	} else {
		base.LabelMinor = fmt.Sprintf("%s (%s)", report.ExtractHostID(n), pid)
	}

	_, isConnected := n.Latest.Lookup(render.IsConnected)
	base.Linkable = isConnected
	return base, true
}

func containerNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base.Label = getRenderableContainerName(n)
	base.LabelMinor = report.ExtractHostID(n)

	if imageName, ok := n.Latest.Lookup(docker.ImageName); ok {
		base.Rank = docker.ImageNameWithoutVersion(imageName)
	}

	return base, true
}

func containerImageNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	imageName, ok := n.Latest.Lookup(docker.ImageName)
	if !ok {
		return NodeSummary{}, false
	}

	imageNameWithoutVersion := docker.ImageNameWithoutVersion(imageName)
	base.Label = imageNameWithoutVersion
	base.Rank = imageNameWithoutVersion
	base.Stack = true

	if base.Label == ImageNameNone {
		base.Label, _ = n.Latest.Lookup(docker.ImageID)
		if len(base.Label) > 12 {
			base.Label = base.Label[:12]
		}
	}

	if c, ok := n.Counters.Lookup(report.Container); ok {
		if c == 1 {
			base.LabelMinor = fmt.Sprintf("%d container", c)
		} else {
			base.LabelMinor = fmt.Sprintf("%d containers", c)
		}
	}
	return base, true
}

func addKubernetesLabelAndRank(base NodeSummary, n report.Node) NodeSummary {
	base.Label, _ = n.Latest.Lookup(kubernetes.Name)
	namespace, _ := n.Latest.Lookup(kubernetes.Namespace)
	base.Rank = namespace + "/" + base.Label
	return base
}

func podNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base = addKubernetesLabelAndRank(base, n)
	if c, ok := n.Counters.Lookup(report.Container); ok {
		if c == 1 {
			base.LabelMinor = fmt.Sprintf("%d container", c)
		} else {
			base.LabelMinor = fmt.Sprintf("%d containers", c)
		}
	}

	return base, true
}

func serviceNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base = addKubernetesLabelAndRank(base, n)
	base.Stack = true

	// Services are always just a group of pods, so there's no counting multiple
	// services which might be grouped together.
	if p, ok := n.Counters.Lookup(report.Pod); ok {
		if p == 1 {
			base.LabelMinor = fmt.Sprintf("%d pod", p)
		} else {
			base.LabelMinor = fmt.Sprintf("%d pods", p)
		}
	}

	return base, true
}

func deploymentNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base = addKubernetesLabelAndRank(base, n)
	base.Stack = true

	if p, ok := n.Counters.Lookup(report.Pod); ok {
		if p == 1 {
			base.LabelMinor = fmt.Sprintf("%d pod", p)
		} else {
			base.LabelMinor = fmt.Sprintf("%d pods", p)
		}
	}

	return base, true
}

func replicaSetNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	base = addKubernetesLabelAndRank(base, n)
	base.Stack = true

	if p, ok := n.Counters.Lookup(report.Pod); ok {
		if p == 1 {
			base.LabelMinor = fmt.Sprintf("%d pod", p)
		} else {
			base.LabelMinor = fmt.Sprintf("%d pods", p)
		}
	}

	return base, true
}

func hostNodeSummary(base NodeSummary, n report.Node) (NodeSummary, bool) {
	var (
		hostname, _ = n.Latest.Lookup(host.HostName)
		parts       = strings.SplitN(hostname, ".", 2)
	)

	if len(parts) == 2 {
		base.Label, base.LabelMinor, base.Rank = parts[0], parts[1], parts[1]
	} else {
		base.Label = hostname
	}

	return base, true
}

// groupNodeSummary renders the summary for a group node. n.Topology is
// expected to be of the form: group:container:hostname
func groupNodeSummary(base NodeSummary, r report.Report, n report.Node) (NodeSummary, bool) {
	parts := strings.Split(n.Topology, ":")
	if len(parts) != 3 {
		return NodeSummary{}, false
	}

	label, ok := n.Latest.Lookup(parts[2])
	if !ok {
		return NodeSummary{}, false
	}
	base.Label, base.Rank = label, label

	t, ok := r.Topology(parts[1])
	if ok && t.Label != "" {
		if count, ok := n.Counters.Lookup(parts[1]); ok {
			if count == 1 {
				base.LabelMinor = fmt.Sprintf("%d %s", count, t.Label)
			} else {
				base.LabelMinor = fmt.Sprintf("%d %s", count, t.LabelPlural)
			}
		}
	}

	base.Shape = t.GetShape()
	base.Stack = true
	return base, true
}

type nodeSummariesByID []NodeSummary

func (s nodeSummariesByID) Len() int           { return len(s) }
func (s nodeSummariesByID) Swap(i, j int)      { s[i], s[j] = s[j], s[i] }
func (s nodeSummariesByID) Less(i, j int) bool { return s[i].ID < s[j].ID }

// NodeSummaries is a set of NodeSummaries indexed by ID.
type NodeSummaries map[string]NodeSummary

// Summaries converts RenderableNodes into a set of NodeSummaries
func Summaries(r report.Report, rns report.Nodes) NodeSummaries {
	result := NodeSummaries{}
	for id, node := range rns {
		if summary, ok := MakeNodeSummary(r, node); ok {
			for i, m := range summary.Metrics {
				summary.Metrics[i] = m.Summary()
			}
			result[id] = summary
		}
	}
	return result
}

// Copy returns a deep value-copy of NodeSummaries
func (n NodeSummaries) Copy() NodeSummaries {
	result := NodeSummaries{}
	for k, v := range n {
		result[k] = v.Copy()
	}
	return result
}

// getRenderableContainerName obtains a user-friendly container name, to render in the UI
func getRenderableContainerName(nmd report.Node) string {
	for _, key := range []string{
		// Amazon's ecs-agent produces huge Docker container names, destructively
		// derived from mangling Container Definition names in Task
		// Definitions.
		//
		// However, the ecs-agent provides a label containing the original Container
		// Definition name.
		docker.LabelPrefix + AmazonECSContainerNameLabel,
		// Kubernetes also mangles its Docker container names and provides a
		// label with the original container name. However, note that this label
		// is only provided by Kubernetes versions >= 1.2 (see
		// https://github.com/kubernetes/kubernetes/pull/17234/ )
		docker.LabelPrefix + KubernetesContainerNameLabel,
		// Marathon doesn't set any Docker labels and this is the only meaningful
		// attribute we can find to make Scope useful without Mesos plugin
		docker.EnvPrefix + MarathonAppIDEnv,
		docker.ContainerName,
		docker.ContainerHostname,
	} {
		if label, ok := nmd.Latest.Lookup(key); ok {
			return label
		}
	}
	return ""
}
<|endoftext|>


<commit_before>// Package b2 implements Storage based on an Backblaze B2 bucket.
package b2

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/efarrer/iothrottler"
	"github.com/pkg/errors"
	backblaze "gopkg.in/kothar/go-backblaze.v0"

	"github.com/kopia/kopia/internal/retry"
	"github.com/kopia/kopia/repo/blob"
)

const (
	b2storageType = "b2"
)

type b2Storage struct {
	Options

	ctx context.Context

	cli    *backblaze.B2
	bucket *backblaze.Bucket

	downloadThrottler *iothrottler.IOThrottlerPool
	uploadThrottler   *iothrottler.IOThrottlerPool
}

func (s *b2Storage) GetBlob(ctx context.Context, id blob.ID, offset, length int64) ([]byte, error) {
	fileName := s.getObjectNameString(id)

	attempt := func() (interface{}, error) {
		var fileRange *backblaze.FileRange

		if length > 0 {
			fileRange = &backblaze.FileRange{
				Start: offset,
				End:   offset + length - 1,
			}
		}

		_, r, err := s.bucket.DownloadFileRangeByName(fileName, fileRange)
		if err != nil {
			return nil, err
		}
		defer r.Close() //nolint:errcheck

		throttled, err := s.downloadThrottler.AddReader(r)
		if err != nil {
			return nil, err
		}

		b, err := ioutil.ReadAll(throttled)
		if err != nil {
			return nil, err
		}

		if len(b) != int(length) && length > 0 {
			return nil, errors.Errorf("invalid length, got %v bytes, but expected %v", len(b), length)
		}

		if length == 0 {
			return []byte{}, nil
		}

		return b, nil
	}

	v, err := exponentialBackoff(ctx, fmt.Sprintf("GetBlob(%q,%v,%v)", id, offset, length), attempt)
	if err != nil {
		return nil, translateError(err)
	}

	return v.([]byte), nil
}

func (s *b2Storage) resolveFileID(fileName string) (string, error) {
	resp, err := s.bucket.ListFileVersions(fileName, "", 1)
	if err != nil {
		return "", err
	}

	if len(resp.Files) > 0 {
		if resp.Files[0].Name == fileName && resp.Files[0].Action == backblaze.Upload {
			return resp.Files[0].ID, nil
		}
	}

	return "", nil
}

func (s *b2Storage) GetMetadata(ctx context.Context, id blob.ID) (blob.Metadata, error) {
	fileName := s.getObjectNameString(id)

	attempt := func() (interface{}, error) {
		fileID, err := s.resolveFileID(fileName)
		if err != nil {
			return nil, err
		}

		fi, err := s.bucket.GetFileInfo(fileID)
		if err != nil {
			return nil, err
		}

		return blob.Metadata{
			BlobID:    id,
			Length:    fi.ContentLength,
			Timestamp: time.Unix(0, fi.UploadTimestamp*1e6),
		}, nil
	}

	v, err := exponentialBackoff(ctx, fmt.Sprintf("GetMetadata(%q)", id), attempt)
	if err != nil {
		return blob.Metadata{}, translateError(err)
	}

	return v.(blob.Metadata), nil
}

func translateError(err error) error {
	if b2err, ok := err.(*backblaze.B2Error); ok {
		if b2err.Status == http.StatusNotFound {
			// Normal "not found". That's fine.
			return blob.ErrBlobNotFound
		}

		if b2err.Status == http.StatusBadRequest && (b2err.Code == "already_hidden" || b2err.Code == "no_such_file") {
			// Special case when hiding a file that is already hidden. It's basically
			// not found.
			return blob.ErrBlobNotFound
		}
	}

	return err
}

func exponentialBackoff(ctx context.Context, desc string, att retry.AttemptFunc) (interface{}, error) {
	return retry.WithExponentialBackoff(ctx, desc, att, isRetriableError)
}

func isRetriableError(err error) bool {
	if b2err, ok := err.(*backblaze.B2Error); ok {
		switch b2err.Status {
		case http.StatusRequestTimeout:
			return true
		case http.StatusTooManyRequests:
			return true
		case http.StatusInternalServerError:
			return true
		case http.StatusServiceUnavailable:
			return true
		case http.StatusGatewayTimeout:
			return true
		}
	}

	return false
}

func (s *b2Storage) PutBlob(ctx context.Context, id blob.ID, data blob.Bytes) error {
	throttled, err := s.uploadThrottler.AddReader(ioutil.NopCloser(data.Reader()))
	if err != nil {
		return err
	}

	progressCallback := blob.ProgressCallback(ctx)
	if progressCallback != nil {
		progressCallback(string(id), 0, int64(data.Length()))
		defer progressCallback(string(id), int64(data.Length()), int64(data.Length()))
	}

	fileName := s.getObjectNameString(id)
	_, err = s.bucket.UploadFile(fileName, nil, throttled)

	return translateError(err)
}

func (s *b2Storage) SetTime(ctx context.Context, b blob.ID, t time.Time) error {
	return blob.ErrSetTimeUnsupported
}

func (s *b2Storage) DeleteBlob(ctx context.Context, id blob.ID) error {
	fileName := s.getObjectNameString(id)
	_, err := s.bucket.HideFile(fileName)
	err = translateError(err)

	if errors.Is(err, blob.ErrBlobNotFound) {
		// Deleting failed because it already is deleted? Fine.
		return nil
	}

	return err
}

func (s *b2Storage) getObjectNameString(id blob.ID) string {
	return s.Prefix + string(id)
}

func (s *b2Storage) ListBlobs(ctx context.Context, prefix blob.ID, callback func(blob.Metadata) error) error {
	const maxFileQuery = 1000

	fullPrefix := s.getObjectNameString(prefix)
	nextFile := ""

	for {
		resp, err := s.bucket.ListFileNamesWithPrefix(nextFile, maxFileQuery, fullPrefix, "")
		if err != nil {
			return err
		}

		for i := range resp.Files {
			f := &resp.Files[i]
			bm := blob.Metadata{
				BlobID:    blob.ID(f.Name[len(s.Prefix):]),
				Length:    f.ContentLength,
				Timestamp: time.Unix(0, f.UploadTimestamp*int64(time.Millisecond)),
			}

			if err := callback(bm); err != nil {
				return err
			}
		}

		nextFile = resp.NextFileName

		if nextFile == "" {
			break
		}
	}

	return nil
}

func (s *b2Storage) ConnectionInfo() blob.ConnectionInfo {
	return blob.ConnectionInfo{
		Type:   b2storageType,
		Config: &s.Options,
	}
}

func (s *b2Storage) DisplayName() string {
	return fmt.Sprintf("B2: %v", s.BucketName)
}

func (s *b2Storage) Close(ctx context.Context) error {
	return nil
}

func (s *b2Storage) String() string {
	return fmt.Sprintf("b2://%s/%s", s.BucketName, s.Prefix)
}

func toBandwidth(bytesPerSecond int) iothrottler.Bandwidth {
	if bytesPerSecond <= 0 {
		return iothrottler.Unlimited
	}

	return iothrottler.Bandwidth(bytesPerSecond) * iothrottler.BytesPerSecond
}

// New creates new B2-backed storage with specified options.
func New(ctx context.Context, opt *Options) (blob.Storage, error) {
	if opt.BucketName == "" {
		return nil, errors.New("bucket name must be specified")
	}

	cli, err := backblaze.NewB2(backblaze.Credentials{KeyID: opt.KeyID, ApplicationKey: opt.Key})
	if err != nil {
		return nil, errors.Wrap(err, "unable to create client")
	}

	downloadThrottler := iothrottler.NewIOThrottlerPool(toBandwidth(opt.MaxDownloadSpeedBytesPerSecond))
	uploadThrottler := iothrottler.NewIOThrottlerPool(toBandwidth(opt.MaxUploadSpeedBytesPerSecond))

	bucket, err := cli.Bucket(opt.BucketName)
	if err != nil {
		return nil, errors.Wrapf(err, "cannot open bucket %q", opt.BucketName)
	}

	if bucket == nil {
		return nil, errors.Errorf("bucket not found: %s", opt.BucketName)
	}

	return &b2Storage{
		Options:           *opt,
		ctx:               ctx,
		cli:               cli,
		bucket:            bucket,
		downloadThrottler: downloadThrottler,
		uploadThrottler:   uploadThrottler,
	}, nil
}

func init() {
	blob.AddSupportedStorage(
		b2storageType,
		func() interface{} {
			return &Options{}
		},
		func(ctx context.Context, o interface{}) (blob.Storage, error) {
			return New(ctx, o.(*Options))
		})
}
<commit_msg>b2: added exponential back-off around PutBlob() and DeleteBlob() (#679)<commit_after>// Package b2 implements Storage based on an Backblaze B2 bucket.
package b2

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	"github.com/efarrer/iothrottler"
	"github.com/pkg/errors"
	backblaze "gopkg.in/kothar/go-backblaze.v0"

	"github.com/kopia/kopia/internal/retry"
	"github.com/kopia/kopia/repo/blob"
)

const (
	b2storageType = "b2"
)

type b2Storage struct {
	Options

	ctx context.Context

	cli    *backblaze.B2
	bucket *backblaze.Bucket

	downloadThrottler *iothrottler.IOThrottlerPool
	uploadThrottler   *iothrottler.IOThrottlerPool
}

func (s *b2Storage) GetBlob(ctx context.Context, id blob.ID, offset, length int64) ([]byte, error) {
	fileName := s.getObjectNameString(id)

	attempt := func() (interface{}, error) {
		var fileRange *backblaze.FileRange

		if length > 0 {
			fileRange = &backblaze.FileRange{
				Start: offset,
				End:   offset + length - 1,
			}
		}

		_, r, err := s.bucket.DownloadFileRangeByName(fileName, fileRange)
		if err != nil {
			return nil, err
		}
		defer r.Close() //nolint:errcheck

		throttled, err := s.downloadThrottler.AddReader(r)
		if err != nil {
			return nil, err
		}

		b, err := ioutil.ReadAll(throttled)
		if err != nil {
			return nil, err
		}

		if len(b) != int(length) && length > 0 {
			return nil, errors.Errorf("invalid length, got %v bytes, but expected %v", len(b), length)
		}

		if length == 0 {
			return []byte{}, nil
		}

		return b, nil
	}

	v, err := exponentialBackoff(ctx, fmt.Sprintf("GetBlob(%q,%v,%v)", id, offset, length), attempt)
	if err != nil {
		return nil, translateError(err)
	}

	return v.([]byte), nil
}

func (s *b2Storage) resolveFileID(fileName string) (string, error) {
	resp, err := s.bucket.ListFileVersions(fileName, "", 1)
	if err != nil {
		return "", err
	}

	if len(resp.Files) > 0 {
		if resp.Files[0].Name == fileName && resp.Files[0].Action == backblaze.Upload {
			return resp.Files[0].ID, nil
		}
	}

	return "", nil
}

func (s *b2Storage) GetMetadata(ctx context.Context, id blob.ID) (blob.Metadata, error) {
	fileName := s.getObjectNameString(id)

	attempt := func() (interface{}, error) {
		fileID, err := s.resolveFileID(fileName)
		if err != nil {
			return nil, err
		}

		fi, err := s.bucket.GetFileInfo(fileID)
		if err != nil {
			return nil, err
		}

		return blob.Metadata{
			BlobID:    id,
			Length:    fi.ContentLength,
			Timestamp: time.Unix(0, fi.UploadTimestamp*1e6),
		}, nil
	}

	v, err := exponentialBackoff(ctx, fmt.Sprintf("GetMetadata(%q)", id), attempt)
	if err != nil {
		return blob.Metadata{}, translateError(err)
	}

	return v.(blob.Metadata), nil
}

func translateError(err error) error {
	if b2err, ok := err.(*backblaze.B2Error); ok {
		if b2err.Status == http.StatusNotFound {
			// Normal "not found". That's fine.
			return blob.ErrBlobNotFound
		}

		if b2err.Status == http.StatusBadRequest && (b2err.Code == "already_hidden" || b2err.Code == "no_such_file") {
			// Special case when hiding a file that is already hidden. It's basically
			// not found.
			return blob.ErrBlobNotFound
		}
	}

	return err
}

func exponentialBackoff(ctx context.Context, desc string, att retry.AttemptFunc) (interface{}, error) {
	return retry.WithExponentialBackoff(ctx, desc, att, isRetriableError)
}

func isRetriableError(err error) bool {
	if b2err, ok := err.(*backblaze.B2Error); ok {
		switch b2err.Status {
		case http.StatusRequestTimeout:
			return true
		case http.StatusTooManyRequests:
			return true
		case http.StatusInternalServerError:
			return true
		case http.StatusServiceUnavailable:
			return true
		case http.StatusGatewayTimeout:
			return true
		}
	}

	return false
}

func (s *b2Storage) PutBlob(ctx context.Context, id blob.ID, data blob.Bytes) error {
	progressCallback := blob.ProgressCallback(ctx)
	if progressCallback != nil {
		progressCallback(string(id), 0, int64(data.Length()))
		defer progressCallback(string(id), int64(data.Length()), int64(data.Length()))
	}

	attempt := func() (interface{}, error) {
		throttled, err := s.uploadThrottler.AddReader(ioutil.NopCloser(data.Reader()))
		if err != nil {
			return nil, err
		}

		fileName := s.getObjectNameString(id)
		_, err = s.bucket.UploadFile(fileName, nil, throttled)

		return nil, err
	}

	if _, err := exponentialBackoff(ctx, fmt.Sprintf("PutBlob(%q)", id), attempt); err != nil {
		return translateError(err)
	}

	return nil
}

func (s *b2Storage) SetTime(ctx context.Context, b blob.ID, t time.Time) error {
	return blob.ErrSetTimeUnsupported
}

func (s *b2Storage) DeleteBlob(ctx context.Context, id blob.ID) error {
	attempt := func() (interface{}, error) {
		fileName := s.getObjectNameString(id)
		return s.bucket.HideFile(fileName)
	}

	if _, err := exponentialBackoff(ctx, fmt.Sprintf("DeleteBlob(%q)", id), attempt); err != nil {
		err = translateError(err)
		if errors.Is(err, blob.ErrBlobNotFound) {
			// Deleting failed because it already is deleted? Fine.
			return nil
		}
	}

	return nil
}

func (s *b2Storage) getObjectNameString(id blob.ID) string {
	return s.Prefix + string(id)
}

func (s *b2Storage) ListBlobs(ctx context.Context, prefix blob.ID, callback func(blob.Metadata) error) error {
	const maxFileQuery = 1000

	fullPrefix := s.getObjectNameString(prefix)
	nextFile := ""

	for {
		resp, err := s.bucket.ListFileNamesWithPrefix(nextFile, maxFileQuery, fullPrefix, "")
		if err != nil {
			return err
		}

		for i := range resp.Files {
			f := &resp.Files[i]
			bm := blob.Metadata{
				BlobID:    blob.ID(f.Name[len(s.Prefix):]),
				Length:    f.ContentLength,
				Timestamp: time.Unix(0, f.UploadTimestamp*int64(time.Millisecond)),
			}

			if err := callback(bm); err != nil {
				return err
			}
		}

		nextFile = resp.NextFileName

		if nextFile == "" {
			break
		}
	}

	return nil
}

func (s *b2Storage) ConnectionInfo() blob.ConnectionInfo {
	return blob.ConnectionInfo{
		Type:   b2storageType,
		Config: &s.Options,
	}
}

func (s *b2Storage) DisplayName() string {
	return fmt.Sprintf("B2: %v", s.BucketName)
}

func (s *b2Storage) Close(ctx context.Context) error {
	return nil
}

func (s *b2Storage) String() string {
	return fmt.Sprintf("b2://%s/%s", s.BucketName, s.Prefix)
}

func toBandwidth(bytesPerSecond int) iothrottler.Bandwidth {
	if bytesPerSecond <= 0 {
		return iothrottler.Unlimited
	}

	return iothrottler.Bandwidth(bytesPerSecond) * iothrottler.BytesPerSecond
}

// New creates new B2-backed storage with specified options.
func New(ctx context.Context, opt *Options) (blob.Storage, error) {
	if opt.BucketName == "" {
		return nil, errors.New("bucket name must be specified")
	}

	cli, err := backblaze.NewB2(backblaze.Credentials{KeyID: opt.KeyID, ApplicationKey: opt.Key})
	if err != nil {
		return nil, errors.Wrap(err, "unable to create client")
	}

	downloadThrottler := iothrottler.NewIOThrottlerPool(toBandwidth(opt.MaxDownloadSpeedBytesPerSecond))
	uploadThrottler := iothrottler.NewIOThrottlerPool(toBandwidth(opt.MaxUploadSpeedBytesPerSecond))

	bucket, err := cli.Bucket(opt.BucketName)
	if err != nil {
		return nil, errors.Wrapf(err, "cannot open bucket %q", opt.BucketName)
	}

	if bucket == nil {
		return nil, errors.Errorf("bucket not found: %s", opt.BucketName)
	}

	return &b2Storage{
		Options:           *opt,
		ctx:               ctx,
		cli:               cli,
		bucket:            bucket,
		downloadThrottler: downloadThrottler,
		uploadThrottler:   uploadThrottler,
	}, nil
}

func init() {
	blob.AddSupportedStorage(
		b2storageType,
		func() interface{} {
			return &Options{}
		},
		func(ctx context.Context, o interface{}) (blob.Storage, error) {
			return New(ctx, o.(*Options))
		})
}
<|endoftext|>


<commit_before>package mesosreporter

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/client/adapter"
	"github.com/dropbox/changes-client/client/reporter"
)

var (
	// The size of the payload queue. Once it reaches this size,
	// all writes (and thus all of the exported functions from this
	// module) will become blocking.
	maxPendingReports = 64

	// Maximum number of times we retry a payload until we give up
	// and panic.
	//
	// XXX the panic occurs during the publish goroutine, which might
	// not be well characterized for properly handling the error.
	numPublishRetries = 8

	// How long we wait before retrying a payload. We always wait
	// this amount of time so the name here is a bit of a misnomer.
	backoffTimeMs = 1000
)

// All data that goes to the server is encompassed in a payload.
type ReportPayload struct {
	path string
	// A map of fields to their values. Note that the date field
	// will be automatically set when the data is sent.
	data     map[string]string
	filename string
}

// A reporter that connects and reports to a specific jobstep id.
// Each jobstep id has a number of endpoints associated with it that
// allows the reporter to update the status of logs, snapshots, etc.
//
// With each reporter there is a goroutine associated with it that
// listens to publishChannel and shutdownChannel, publishing all
// data from publishChannel to the publishUri for the jobstepID
// associated with the current build. Sending any information
// to shutdownChannel causes the goroutine to stop.
//
// Notably this means that all of the methods in this module
// are asynchronous and as a result there is a delay between
// them successfully finishing and Changes actually acknowledging
// them at the endpoint. More importantly, however, because the
// requests are sent in a separate goroutine, the methods here may
// succeed even when the endpoing requests fail.
//
// In debug mode, endpoint requests are still queued in the publish
// channel but never sent by the publishing goroutine, which allows
// the mesos reporter to run without actually connecting
// to the changes server.
type Reporter struct {
	jobstepID       string
	publishUri      string
	publishChannel  chan ReportPayload
	shutdownChannel chan struct{}
	debug           bool
}

// Utility function for posting a request to a uri. The parameters
// here, which more or less correspond to ReportPayload.data,
// are serialized to form the request body. The body is encoded
// as a MIME multipart (see RFC 2338).
//
// The file is also added a as field in the request body.
func httpPost(uri string, params map[string]string, file string) (resp *http.Response, err error) {
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	for key, val := range params {
		err = writer.WriteField(key, val)
		if err != nil {
			log.Printf("[reporter] Couldn't write field %s", key)
			return nil, err
		}
	}

	if len(file) > 0 {
		handle, err := os.Open(file)
		if err != nil {
			return nil, err
		}

		err = writer.WriteField("name", filepath.Base(file))
		if err != nil {
			return nil, err
		}

		fileField, err := writer.CreateFormFile("file", filepath.Base(file))
		if err != nil {
			return nil, err
		}

		_, err = io.Copy(fileField, handle)
		if err != nil {
			return nil, err
		}
	}

	_ = writer.Close()

	resp, err = http.Post(uri, writer.FormDataContentType(), body)

	if err != nil {
		return nil, err
	}

	// Close the Body channel immediately as we don't use it
	// and this loop can stay open for an extremely long period
	// of time
	resp.Body.Close()

	return resp, nil
}

// Utility method for sending a payload. This wraps httpPost in a framework
// nicer for the Reporter itself, as it turns the ReportPayload into its
// associated params (which corresponds to its data). We also attempt
// httpPost multiple times in order to account for flakiness in the
// network connection. This function is synchronous.
func sendPayload(r *Reporter, rp ReportPayload) {
	var (
		resp   *http.Response
		err    error
		status string
	)

	path := r.publishUri + rp.path
	if rp.data == nil {
		rp.data = make(map[string]string)
	}

	rp.data["date"] = time.Now().UTC().Format("2006-01-02T15:04:05.0Z")
	for tryCnt := 1; tryCnt <= numPublishRetries; tryCnt++ {
		log.Printf("[reporter] POST %s try: %d", path, tryCnt)
		resp, err = httpPost(path, rp.data, rp.filename)

		if resp != nil {
			status = resp.Status
		} else {
			status = "-1"
		}

		if resp != nil && resp.StatusCode/100 == 2 {
			break
		}

		var errmsg string
		if err != nil {
			errmsg = err.Error()
		} else {
			// If there wasn't an IO error, use the response body as the error message.
			var bodyData bytes.Buffer
			if _, e := bodyData.ReadFrom(resp.Body); e != nil {
				log.Printf("[reporter] Error reading POST %s response body: %s", path, e)
			}
			errmsg = bodyData.String()
		}
		log.Printf("[reporter] POST %s failed, try: %d, resp: %s, err: %s",
			path, tryCnt, status, errmsg)

		/* We are unable to publish to the endpoint.
		 * Fail fast and let the above layers handle the outage */
		if tryCnt == numPublishRetries {
			panic("Couldn't to connect to publish endpoint")
		}
		log.Printf("[reporter] Sleep for %d ms", backoffTimeMs)
		time.Sleep(time.Duration(backoffTimeMs) * time.Millisecond)
	}
}

// Continually listens to the publish channel and sends the payloads
// if we aren't in debug mode.
func transportSend(r *Reporter) {
	for rp := range r.publishChannel {
		// dont send reports when running in debug mode
		if r.debug == true {
			continue
		}

		sendPayload(r, rp)
	}
	r.shutdownChannel <- struct{}{}
}

func (r *Reporter) Init(c *client.Config) {
	log.Printf("[reporter] Construct reporter with publish uri: %s", c.Server)
	r.jobstepID = c.JobstepID
	r.publishUri = c.Server
	r.publishChannel = make(chan ReportPayload, maxPendingReports)
	r.shutdownChannel = make(chan struct{})
	r.debug = c.Debug

	// Initialize the goroutine that actually sends the requests. We spawn
	// this even when in debug mode to prevent the payloads from
	// massively queueing. Since we by default use a queue with a maximum
	// limit, once it reaches that limit writes will block causing the
	// main goroutine to halt forever.
	go transportSend(r)
}

func (r *Reporter) PushJobstepStatus(status string, result string) {
	log.Printf("[reporter] Pushing status %s", status)
	form := make(map[string]string)
	form["status"] = status
	if len(result) > 0 {
		form["result"] = result
	}

	hostname, err := os.Hostname()
	if err == nil {
		form["node"] = hostname
	}
	r.publishChannel <- ReportPayload{"/jobsteps/" + r.jobstepID + "/", form, ""}
}

func (r *Reporter) PushCommandStatus(cID string, status string, retCode int) {
	form := make(map[string]string)
	form["status"] = status
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.publishChannel <- ReportPayload{"/commands/" + cID + "/", form, ""}
}

func (r *Reporter) PushSnapshotImageStatus(iID string, status string) {
	form := make(map[string]string)
	form["status"] = status
	r.publishChannel <- ReportPayload{"/snapshotimages/" + iID + "/", form, ""}
}

func (r *Reporter) PushLogChunk(source string, payload []byte) {
	form := make(map[string]string)
	form["source"] = source
	form["text"] = string(payload)
	if r.debug {
		log.Print(string(payload))
	}
	r.publishChannel <- ReportPayload{"/jobsteps/" + r.jobstepID + "/logappend/", form, ""}
}

func (r *Reporter) PushCommandOutput(cID string, status string, retCode int, output []byte) {
	form := make(map[string]string)
	form["status"] = status
	form["output"] = string(output)
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.publishChannel <- ReportPayload{"/commands/" + cID + "/", form, ""}
}

func (r *Reporter) PublishArtifacts(cmd client.ConfigCmd, a adapter.Adapter, clientLog *client.Log) {
	if len(cmd.Artifacts) == 0 {
		clientLog.Writeln("==> Skipping artifact collection")
		return
	}

	clientLog.Writeln(fmt.Sprintf("==> Collecting artifacts matching %s", cmd.Artifacts))

	matches, err := a.CollectArtifacts(cmd.Artifacts, clientLog)
	if err != nil {
		clientLog.Writeln(fmt.Sprintf("==> ERROR: " + err.Error()))
		return
	}

	for _, artifact := range matches {
		clientLog.Writeln(fmt.Sprintf("==> Found: %s", artifact))
	}

	r.pushArtifacts(matches)
}

func (r *Reporter) pushArtifacts(artifacts []string) {
	// TODO: PushArtifacts is synchronous due to races with Adapter.Shutdown(), but
	// really what we'd want to do is just say "wait until channel empty, ok continue"
	for _, artifact := range artifacts {
		sendPayload(r, ReportPayload{"/jobsteps/" + r.jobstepID + "/artifacts/", nil, artifact})
	}
}

// Close the publish and shutdown channels, which causes the inner goroutines to
// terminate, thus cleaning up what is created by Init.
func (r *Reporter) Shutdown() {
	close(r.publishChannel)
	<-r.shutdownChannel
	close(r.shutdownChannel)
	log.Print("[reporter] Shutdown complete")
}

func init() {
	flag.IntVar(&maxPendingReports, "max_pending_reports", 64, "Backlog size")
	flag.IntVar(&numPublishRetries, "num_publish_retries", 8,
		"Number of times to retry")
	flag.IntVar(&maxPendingReports, "backoff_time_ms", 1000,
		"Time to wait between two consecutive retries")

	reporter.Register("mesos", &Reporter{})
}
<commit_msg>Make sure artifact retry failures get to Sentry<commit_after>package mesosreporter

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"log"
	"mime/multipart"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"time"

	"github.com/dropbox/changes-client/client"
	"github.com/dropbox/changes-client/client/adapter"
	"github.com/dropbox/changes-client/client/reporter"
	"github.com/dropbox/changes-client/common/sentry"
)

var (
	// The size of the payload queue. Once it reaches this size,
	// all writes (and thus all of the exported functions from this
	// module) will become blocking.
	maxPendingReports = 64

	// Maximum number of times we retry a payload until we give up
	// and panic.
	//
	// XXX the panic occurs during the publish goroutine, which might
	// not be well characterized for properly handling the error.
	numPublishRetries = 8

	// How long we wait before retrying a payload. We always wait
	// this amount of time so the name here is a bit of a misnomer.
	backoffTimeMs = 1000
)

// All data that goes to the server is encompassed in a payload.
type ReportPayload struct {
	path string
	// A map of fields to their values. Note that the date field
	// will be automatically set when the data is sent.
	data     map[string]string
	filename string
}

// A reporter that connects and reports to a specific jobstep id.
// Each jobstep id has a number of endpoints associated with it that
// allows the reporter to update the status of logs, snapshots, etc.
//
// With each reporter there is a goroutine associated with it that
// listens to publishChannel and shutdownChannel, publishing all
// data from publishChannel to the publishUri for the jobstepID
// associated with the current build. Sending any information
// to shutdownChannel causes the goroutine to stop.
//
// Notably this means that all of the methods in this module
// are asynchronous and as a result there is a delay between
// them successfully finishing and Changes actually acknowledging
// them at the endpoint. More importantly, however, because the
// requests are sent in a separate goroutine, the methods here may
// succeed even when the endpoing requests fail.
//
// In debug mode, endpoint requests are still queued in the publish
// channel but never sent by the publishing goroutine, which allows
// the mesos reporter to run without actually connecting
// to the changes server.
type Reporter struct {
	jobstepID       string
	publishUri      string
	publishChannel  chan ReportPayload
	shutdownChannel chan struct{}
	debug           bool
}

// Utility function for posting a request to a uri. The parameters
// here, which more or less correspond to ReportPayload.data,
// are serialized to form the request body. The body is encoded
// as a MIME multipart (see RFC 2338).
//
// The file is also added a as field in the request body.
func httpPost(uri string, params map[string]string, file string) (resp *http.Response, err error) {
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)

	for key, val := range params {
		err = writer.WriteField(key, val)
		if err != nil {
			log.Printf("[reporter] Couldn't write field %s", key)
			return nil, err
		}
	}

	if len(file) > 0 {
		handle, err := os.Open(file)
		if err != nil {
			return nil, err
		}

		err = writer.WriteField("name", filepath.Base(file))
		if err != nil {
			return nil, err
		}

		fileField, err := writer.CreateFormFile("file", filepath.Base(file))
		if err != nil {
			return nil, err
		}

		_, err = io.Copy(fileField, handle)
		if err != nil {
			return nil, err
		}
	}

	_ = writer.Close()

	resp, err = http.Post(uri, writer.FormDataContentType(), body)

	if err != nil {
		return nil, err
	}

	// Close the Body channel immediately as we don't use it
	// and this loop can stay open for an extremely long period
	// of time
	resp.Body.Close()

	return resp, nil
}

// Utility method for sending a payload. This wraps httpPost in a framework
// nicer for the Reporter itself, as it turns the ReportPayload into its
// associated params (which corresponds to its data). We also attempt
// httpPost multiple times in order to account for flakiness in the
// network connection. This function is synchronous.
func sendPayload(r *Reporter, rp ReportPayload) {
	var (
		resp   *http.Response
		err    error
		status string
	)

	path := r.publishUri + rp.path
	if rp.data == nil {
		rp.data = make(map[string]string)
	}

	rp.data["date"] = time.Now().UTC().Format("2006-01-02T15:04:05.0Z")
	for tryCnt := 1; tryCnt <= numPublishRetries; tryCnt++ {
		log.Printf("[reporter] POST %s try: %d", path, tryCnt)
		resp, err = httpPost(path, rp.data, rp.filename)

		if resp != nil {
			status = resp.Status
		} else {
			status = "-1"
		}

		if resp != nil && resp.StatusCode/100 == 2 {
			break
		}

		var errmsg string
		if err != nil {
			errmsg = err.Error()
		} else {
			// If there wasn't an IO error, use the response body as the error message.
			var bodyData bytes.Buffer
			if _, e := bodyData.ReadFrom(resp.Body); e != nil {
				log.Printf("[reporter] Error reading POST %s response body: %s", path, e)
			}
			errmsg = bodyData.String()
		}
		log.Printf("[reporter] POST %s failed, try: %d, resp: %s, err: %s",
			path, tryCnt, status, errmsg)

		/* We are unable to publish to the endpoint.
		 * Fail fast and let the above layers handle the outage */
		if tryCnt == numPublishRetries {
			const msg = "Couldn't to connect to publish endpoint"
			if client := sentry.GetClient(); client != nil {
				client.CaptureMessageAndWait(msg, nil)
			}
			// TODO: Report an infra failure rather than panicing.
			panic(msg)
		}
		log.Printf("[reporter] Sleep for %d ms", backoffTimeMs)
		time.Sleep(time.Duration(backoffTimeMs) * time.Millisecond)
	}
}

// Continually listens to the publish channel and sends the payloads
// if we aren't in debug mode.
func transportSend(r *Reporter) {
	for rp := range r.publishChannel {
		// dont send reports when running in debug mode
		if r.debug == true {
			continue
		}

		sendPayload(r, rp)
	}
	r.shutdownChannel <- struct{}{}
}

func (r *Reporter) Init(c *client.Config) {
	log.Printf("[reporter] Construct reporter with publish uri: %s", c.Server)
	r.jobstepID = c.JobstepID
	r.publishUri = c.Server
	r.publishChannel = make(chan ReportPayload, maxPendingReports)
	r.shutdownChannel = make(chan struct{})
	r.debug = c.Debug

	// Initialize the goroutine that actually sends the requests. We spawn
	// this even when in debug mode to prevent the payloads from
	// massively queueing. Since we by default use a queue with a maximum
	// limit, once it reaches that limit writes will block causing the
	// main goroutine to halt forever.
	go transportSend(r)
}

func (r *Reporter) PushJobstepStatus(status string, result string) {
	log.Printf("[reporter] Pushing status %s", status)
	form := make(map[string]string)
	form["status"] = status
	if len(result) > 0 {
		form["result"] = result
	}

	hostname, err := os.Hostname()
	if err == nil {
		form["node"] = hostname
	}
	r.publishChannel <- ReportPayload{"/jobsteps/" + r.jobstepID + "/", form, ""}
}

func (r *Reporter) PushCommandStatus(cID string, status string, retCode int) {
	form := make(map[string]string)
	form["status"] = status
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.publishChannel <- ReportPayload{"/commands/" + cID + "/", form, ""}
}

func (r *Reporter) PushSnapshotImageStatus(iID string, status string) {
	form := make(map[string]string)
	form["status"] = status
	r.publishChannel <- ReportPayload{"/snapshotimages/" + iID + "/", form, ""}
}

func (r *Reporter) PushLogChunk(source string, payload []byte) {
	form := make(map[string]string)
	form["source"] = source
	form["text"] = string(payload)
	if r.debug {
		log.Print(string(payload))
	}
	r.publishChannel <- ReportPayload{"/jobsteps/" + r.jobstepID + "/logappend/", form, ""}
}

func (r *Reporter) PushCommandOutput(cID string, status string, retCode int, output []byte) {
	form := make(map[string]string)
	form["status"] = status
	form["output"] = string(output)
	if retCode >= 0 {
		form["return_code"] = strconv.Itoa(retCode)
	}
	r.publishChannel <- ReportPayload{"/commands/" + cID + "/", form, ""}
}

func (r *Reporter) PublishArtifacts(cmd client.ConfigCmd, a adapter.Adapter, clientLog *client.Log) {
	if len(cmd.Artifacts) == 0 {
		clientLog.Writeln("==> Skipping artifact collection")
		return
	}

	clientLog.Writeln(fmt.Sprintf("==> Collecting artifacts matching %s", cmd.Artifacts))

	matches, err := a.CollectArtifacts(cmd.Artifacts, clientLog)
	if err != nil {
		clientLog.Writeln(fmt.Sprintf("==> ERROR: " + err.Error()))
		return
	}

	for _, artifact := range matches {
		clientLog.Writeln(fmt.Sprintf("==> Found: %s", artifact))
	}

	r.pushArtifacts(matches)
}

func (r *Reporter) pushArtifacts(artifacts []string) {
	// TODO: PushArtifacts is synchronous due to races with Adapter.Shutdown(), but
	// really what we'd want to do is just say "wait until channel empty, ok continue"
	for _, artifact := range artifacts {
		sendPayload(r, ReportPayload{"/jobsteps/" + r.jobstepID + "/artifacts/", nil, artifact})
	}
}

// Close the publish and shutdown channels, which causes the inner goroutines to
// terminate, thus cleaning up what is created by Init.
func (r *Reporter) Shutdown() {
	close(r.publishChannel)
	<-r.shutdownChannel
	close(r.shutdownChannel)
	log.Print("[reporter] Shutdown complete")
}

func init() {
	flag.IntVar(&maxPendingReports, "max_pending_reports", 64, "Backlog size")
	flag.IntVar(&numPublishRetries, "num_publish_retries", 8,
		"Number of times to retry")
	flag.IntVar(&maxPendingReports, "backoff_time_ms", 1000,
		"Time to wait between two consecutive retries")

	reporter.Register("mesos", &Reporter{})
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"github.com/uber/tchannel-go"
	"github.com/uber/tchannel-go/thrift"
	"golang.org/x/net/context"
)

// client implements TChanClient and makes outgoing Thrift calls.
type client struct {
	ch          *tchannel.Channel
	sc          *tchannel.SubChannel
	serviceName string
	opts        ClientOptions
}

// ClientOptions are options to customize the client.
type ClientOptions struct {
	// HostPort specifies a specific server to hit.
	HostPort string
}

// NewClient returns a Client that makes calls over the given tchannel to the given Hyperbahn service.
func NewClient(ch *tchannel.Channel, serviceName string, opts *ClientOptions) TChanClient {
	client := &client{
		ch:          ch,
		sc:          ch.GetSubChannel(serviceName),
		serviceName: serviceName,
	}
	if opts != nil {
		client.opts = *opts
	}
	return client
}

func (c *client) startCall(ctx context.Context, method string, callOptions *tchannel.CallOptions) (*tchannel.OutboundCall, error) {
	if c.opts.HostPort != "" {
		return c.ch.BeginCall(ctx, c.opts.HostPort, c.serviceName, method, callOptions)
	}
	return c.sc.BeginCall(ctx, method, callOptions)
}

func writeArgs(call *tchannel.OutboundCall, headers map[string]string, req RWTStruct) error {
	writer, err := call.Arg2Writer()
	if err != nil {
		return err
	}
	headers = tchannel.InjectOutboundSpan(call.Response(), headers)
	if err := thrift.WriteHeaders(writer, headers); err != nil {
		return err
	}
	if err := writer.Close(); err != nil {
		return err
	}

	writer, err = call.Arg3Writer()
	if err != nil {
		return err
	}

	if err := WriteStruct(writer, req); err != nil {
		return err
	}

	return writer.Close()
}

// readResponse reads the response struct into resp, and returns:
// (response headers, whether there was an application error, unexpected error).
func readResponse(response *tchannel.OutboundCallResponse, resp RWTStruct) (map[string]string, bool, error) {
	reader, err := response.Arg2Reader()
	if err != nil {
		return nil, false, err
	}

	headers, err := thrift.ReadHeaders(reader)
	if err != nil {
		return nil, false, err
	}

	if err := EnsureEmpty(reader, "reading response headers"); err != nil {
		return nil, false, err
	}

	if err := reader.Close(); err != nil {
		return nil, false, err
	}

	success := !response.ApplicationError()
	reader, err = response.Arg3Reader()
	if err != nil {
		return headers, success, err
	}

	if err := ReadStruct(reader, resp); err != nil {
		return headers, success, err
	}

	if err := EnsureEmpty(reader, "reading response body"); err != nil {
		return nil, false, err
	}

	return headers, success, reader.Close()
}

func (c *client) Call(ctx thrift.Context, thriftService, methodName string, req, resp RWTStruct) (bool, error) {
	var (
		headers = ctx.Headers()

		respHeaders map[string]string
		isOK        bool
	)

	err := c.ch.RunWithRetry(ctx, func(ctx context.Context, rs *tchannel.RequestState) error {
		respHeaders, isOK = nil, false

		call, err := c.startCall(ctx, thriftService+"::"+methodName, &tchannel.CallOptions{
			Format:       tchannel.Thrift,
			RequestState: rs,
		})
		if err != nil {
			return err
		}

		if err := writeArgs(call, headers, req); err != nil {
			return err
		}

		respHeaders, isOK, err = readResponse(call.Response(), resp)
		return err
	})
	if err != nil {
		return false, err
	}

	ctx.SetResponseHeaders(respHeaders)
	return isOK, nil
}
<commit_msg>Make client call only from subchannel<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package tchannel

import (
	"github.com/uber/tchannel-go"
	"github.com/uber/tchannel-go/thrift"
	"golang.org/x/net/context"
)

// client implements TChanClient and makes outgoing Thrift calls.
type Client struct {
	ch          *tchannel.Channel
	sc          *tchannel.SubChannel
	serviceName string
}

// NewClient returns a Client that makes calls over the given tchannel to the given Hyperbahn service.
func NewClient(ch *tchannel.Channel, serviceName string) TChanClient {
	client := &Client{
		ch:          ch,
		sc:          ch.GetSubChannel(serviceName),
		serviceName: serviceName,
	}
	return client
}

func writeArgs(call *tchannel.OutboundCall, headers map[string]string, req RWTStruct) error {
	writer, err := call.Arg2Writer()
	if err != nil {
		return err
	}
	headers = tchannel.InjectOutboundSpan(call.Response(), headers)
	if err := thrift.WriteHeaders(writer, headers); err != nil {
		return err
	}
	if err := writer.Close(); err != nil {
		return err
	}

	writer, err = call.Arg3Writer()
	if err != nil {
		return err
	}

	if err := WriteStruct(writer, req); err != nil {
		return err
	}

	return writer.Close()
}

// readResponse reads the response struct into resp, and returns:
// (response headers, whether there was an application error, unexpected error).
func readResponse(response *tchannel.OutboundCallResponse, resp RWTStruct) (map[string]string, bool, error) {
	reader, err := response.Arg2Reader()
	if err != nil {
		return nil, false, err
	}

	headers, err := thrift.ReadHeaders(reader)
	if err != nil {
		return nil, false, err
	}

	if err := EnsureEmpty(reader, "reading response headers"); err != nil {
		return nil, false, err
	}

	if err := reader.Close(); err != nil {
		return nil, false, err
	}

	success := !response.ApplicationError()
	reader, err = response.Arg3Reader()
	if err != nil {
		return headers, success, err
	}

	if err := ReadStruct(reader, resp); err != nil {
		return headers, success, err
	}

	if err := EnsureEmpty(reader, "reading response body"); err != nil {
		return nil, false, err
	}

	return headers, success, reader.Close()
}

func (c *Client) Call(ctx thrift.Context, thriftService, methodName string, req, resp RWTStruct) (bool, error) {
	var (
		headers = ctx.Headers()

		respHeaders map[string]string
		isOK        bool
	)

	err := c.ch.RunWithRetry(ctx, func(ctx context.Context, rs *tchannel.RequestState) error {
		respHeaders, isOK = nil, false

		call, err := c.sc.BeginCall(ctx, thriftService+"::"+methodName, &tchannel.CallOptions{
			Format:       tchannel.Thrift,
			RequestState: rs,
		})
		if err != nil {
			return err
		}

		if err := writeArgs(call, headers, req); err != nil {
			return err
		}

		respHeaders, isOK, err = readResponse(call.Response(), resp)
		return err
	})
	if err != nil {
		return false, err
	}

	ctx.SetResponseHeaders(respHeaders)
	return isOK, nil
}
<|endoftext|>


<commit_before>package components

import (
	"encoding/json"

	"github.com/Jeffail/gabs"
	"github.com/bep/gr"
	"github.com/bep/gr/attr"
	"github.com/bep/gr/el"
	"github.com/murdinc/awsm/config"
)

func ClassListBuilder(cl interface{}) *gr.Element {

	classList := cl.([]byte)

	jsonParsed, _ := gabs.ParseJSON(classList)
	classType := jsonParsed.S("classType").Data().(string)
	classes, _ := jsonParsed.S("classes").ChildrenMap()

	classListGroup := el.Div(
		el.Div(
			gr.CSS("list-group"),
		),
	)

	switch classType {

	case "alarms":
		var cType config.AlarmClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "instances":
		var cType config.InstanceClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "volumes":
		var cType config.VolumeClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "images":
		var cType config.ImageClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "keypairs":
	// TODO ?

	case "snapshots":
		var cType config.SnapshotClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "vpcs":
		var cType config.VpcClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "subnets":
		var cType config.SubnetClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "securitygroups":
		var cType config.SecurityGroupClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "addresses":
	// TODO ?

	case "launchconfigurations":
		var cType config.LaunchConfigurationClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "loadbalancers":
	// TODO ?

	case "scalingpolicies":
		var cType config.ScalingPolicyClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "simpledbdomains":
	// TODO ?

	default:
		println("Class Type not found in ClassListBuilder switch:")
		println(classType)
	}

	return classListGroup
}

func buildClassButton(name string, keys []string, values []string, classListGroup *gr.Element) {

	button := el.Button(
		attr.Type("button"),
		gr.CSS("list-group-item"),
		el.Header5(
			gr.CSS("list-group-item-heading"),
			gr.Text(name),
		),
	)

	description := el.Div(
		gr.CSS("list-group-item-text"),
	)

	left := el.DescriptionList(gr.CSS("dl-horizontal"))
	right := el.DescriptionList(gr.CSS("dl-horizontal"))

	keyLen := len(keys)

	for i, key := range keys {
		if i < (keyLen / 2) {
			el.DefinitionTerm(gr.Text(key)).Modify(left)
			el.Description(gr.Text(values[i])).Modify(left)
		} else {
			el.DefinitionTerm(gr.Text(key)).Modify(right)
			el.Description(gr.Text(values[i])).Modify(right)
		}
	}

	el.Div(
		gr.CSS("row"),
		el.Div(gr.CSS("col-sm-6"),
			left),
		el.Div(gr.CSS("col-sm-6"),
			right),
	).Modify(description)

	description.Modify(button)
	button.Modify(classListGroup)

}
<commit_msg>adds switch to build load balance class output<commit_after>package components

import (
	"encoding/json"

	"github.com/Jeffail/gabs"
	"github.com/bep/gr"
	"github.com/bep/gr/attr"
	"github.com/bep/gr/el"
	"github.com/murdinc/awsm/config"
)

func ClassListBuilder(cl interface{}) *gr.Element {

	classList := cl.([]byte)

	jsonParsed, _ := gabs.ParseJSON(classList)
	classType := jsonParsed.S("classType").Data().(string)
	classes, _ := jsonParsed.S("classes").ChildrenMap()

	classListGroup := el.Div(
		el.Div(
			gr.CSS("list-group"),
		),
	)

	switch classType {

	case "alarms":
		var cType config.AlarmClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "instances":
		var cType config.InstanceClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "volumes":
		var cType config.VolumeClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "images":
		var cType config.ImageClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "keypairs":
	// TODO ?

	case "snapshots":
		var cType config.SnapshotClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "vpcs":
		var cType config.VpcClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "subnets":
		var cType config.SubnetClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "securitygroups":
		var cType config.SecurityGroupClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "addresses":
	// TODO ?

	case "launchconfigurations":
		var cType config.LaunchConfigurationClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "loadbalancers":
		var cType config.LoadBalancerClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	case "scalingpolicies":
		var cType config.ScalingPolicyClass
		for className, class := range classes {
			cJson := class.Bytes()
			json.Unmarshal(cJson, &cType)
			keys, values := config.ExtractAwsmList(cType)
			buildClassButton(className, keys, values, classListGroup)
		}

	//case "simpledbdomains":
	// TODO ?

	default:
		println("Class Type not found in ClassListBuilder switch:")
		println(classType)
	}

	return classListGroup
}

func buildClassButton(name string, keys []string, values []string, classListGroup *gr.Element) {

	button := el.Button(
		attr.Type("button"),
		gr.CSS("list-group-item"),
		el.Header5(
			gr.CSS("list-group-item-heading"),
			gr.Text(name),
		),
	)

	description := el.Div(
		gr.CSS("list-group-item-text"),
	)

	left := el.DescriptionList(gr.CSS("dl-horizontal"))
	right := el.DescriptionList(gr.CSS("dl-horizontal"))

	keyLen := len(keys)

	for i, key := range keys {
		if i < (keyLen / 2) {
			el.DefinitionTerm(gr.Text(key)).Modify(left)
			el.Description(gr.Text(values[i])).Modify(left)
		} else {
			el.DefinitionTerm(gr.Text(key)).Modify(right)
			el.Description(gr.Text(values[i])).Modify(right)
		}
	}

	el.Div(
		gr.CSS("row"),
		el.Div(gr.CSS("col-sm-6"),
			left),
		el.Div(gr.CSS("col-sm-6"),
			right),
	).Modify(description)

	description.Modify(button)
	button.Modify(classListGroup)

}
<|endoftext|>


<commit_before>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"

	"github.com/golang/glog"
	"github.com/google/cadvisor/api"
	"github.com/google/cadvisor/container/docker"
	"github.com/google/cadvisor/container/raw"
	"github.com/google/cadvisor/healthz"
	"github.com/google/cadvisor/info"
	"github.com/google/cadvisor/manager"
	"github.com/google/cadvisor/pages"
	"github.com/google/cadvisor/pages/static"
)

var argPort = flag.Int("port", 8080, "port to listen")
var maxProcs = flag.Int("max_procs", 0, "max number of CPUs that can be used simultaneously. Less than 1 for default (number of cores).")

var argDbDriver = flag.String("storage_driver", "", "storage driver to use. Empty means none. Options are: <empty> (default), bigquery, and influxdb")

func main() {
	defer glog.Flush()
	flag.Parse()

	setMaxProcs()

	storageDriver, err := NewStorageDriver(*argDbDriver)
	if err != nil {
		glog.Fatalf("Failed to connect to database: %s", err)
	}

	containerManager, err := manager.New(storageDriver)
	if err != nil {
		glog.Fatalf("Failed to create a Container Manager: %s", err)
	}

	// Register Docker.
	if err := docker.Register(containerManager); err != nil {
		glog.Errorf("Docker registration failed: %v.", err)
	}

	// Register the raw driver.
	if err := raw.Register(containerManager); err != nil {
		glog.Fatalf("raw registration failed: %v.", err)
	}

	// Basic health handler.
	if err := healthz.RegisterHandler(); err != nil {
		glog.Fatalf("failed to register healthz handler: %s", err)
	}

	// Handler for static content.
	http.HandleFunc(static.StaticResource, func(w http.ResponseWriter, r *http.Request) {
		err := static.HandleRequest(w, r.URL)
		if err != nil {
			fmt.Fprintf(w, "%s", err)
		}
	})

	// Register API handler.
	if err := api.RegisterHandlers(containerManager); err != nil {
		glog.Fatalf("failed to register API handlers: %s", err)
	}

	// Redirect / to containers page.
	http.Handle("/", http.RedirectHandler(pages.ContainersPage, http.StatusTemporaryRedirect))

	// Register the handler for the containers page.
	http.HandleFunc(pages.ContainersPage, func(w http.ResponseWriter, r *http.Request) {
		err := pages.ServerContainersPage(containerManager, w, r.URL)
		if err != nil {
			fmt.Fprintf(w, "%s", err)
		}
	})

	// Start the manager.
	if err := containerManager.Start(); err != nil {
		glog.Fatalf("Failed to start container manager: %v", err)
	}

	// Install signal handler.
	installSignalHandler(containerManager)

	glog.Infof("Starting cAdvisor version: %q on port %d", info.VERSION, *argPort)

	addr := fmt.Sprintf(":%v", *argPort)
	glog.Fatal(http.ListenAndServe(addr, nil))
}

func setMaxProcs() {
	// TODO(vmarmol): Consider limiting if we have a CPU mask in effect.
	// Allow as many threads as we have cores unless the user specified a value.
	var numProcs int
	if *maxProcs < 1 {
		numProcs = runtime.NumCPU()
	} else {
		numProcs = *maxProcs
	}
	runtime.GOMAXPROCS(numProcs)

	// Check if the setting was successful.
	actualNumProcs := runtime.GOMAXPROCS(0)
	if actualNumProcs != numProcs {
		glog.Warningf("Specified max procs of %v but using %v", numProcs, actualNumProcs)
	}
}

func installSignalHandler(containerManager manager.Manager) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill)

	// Block until a signal is received.
	go func() {
		sig := <-c
		if err := containerManager.Stop(); err != nil {
			glog.Errorf("Failed to stop container manager: %v", err)
		}
		glog.Infof("Exiting given signal: %v", sig)
		os.Exit(0)
	}()
}
<commit_msg>Also handle SIGTERM.<commit_after>// Copyright 2014 Google Inc. All Rights Reserved.
//
// 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 main

import (
	"flag"
	"fmt"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/signal"
	"runtime"
	"syscall"

	"github.com/golang/glog"
	"github.com/google/cadvisor/api"
	"github.com/google/cadvisor/container/docker"
	"github.com/google/cadvisor/container/raw"
	"github.com/google/cadvisor/healthz"
	"github.com/google/cadvisor/info"
	"github.com/google/cadvisor/manager"
	"github.com/google/cadvisor/pages"
	"github.com/google/cadvisor/pages/static"
)

var argPort = flag.Int("port", 8080, "port to listen")
var maxProcs = flag.Int("max_procs", 0, "max number of CPUs that can be used simultaneously. Less than 1 for default (number of cores).")

var argDbDriver = flag.String("storage_driver", "", "storage driver to use. Empty means none. Options are: <empty> (default), bigquery, and influxdb")

func main() {
	defer glog.Flush()
	flag.Parse()

	setMaxProcs()

	storageDriver, err := NewStorageDriver(*argDbDriver)
	if err != nil {
		glog.Fatalf("Failed to connect to database: %s", err)
	}

	containerManager, err := manager.New(storageDriver)
	if err != nil {
		glog.Fatalf("Failed to create a Container Manager: %s", err)
	}

	// Register Docker.
	if err := docker.Register(containerManager); err != nil {
		glog.Errorf("Docker registration failed: %v.", err)
	}

	// Register the raw driver.
	if err := raw.Register(containerManager); err != nil {
		glog.Fatalf("raw registration failed: %v.", err)
	}

	// Basic health handler.
	if err := healthz.RegisterHandler(); err != nil {
		glog.Fatalf("failed to register healthz handler: %s", err)
	}

	// Handler for static content.
	http.HandleFunc(static.StaticResource, func(w http.ResponseWriter, r *http.Request) {
		err := static.HandleRequest(w, r.URL)
		if err != nil {
			fmt.Fprintf(w, "%s", err)
		}
	})

	// Register API handler.
	if err := api.RegisterHandlers(containerManager); err != nil {
		glog.Fatalf("failed to register API handlers: %s", err)
	}

	// Redirect / to containers page.
	http.Handle("/", http.RedirectHandler(pages.ContainersPage, http.StatusTemporaryRedirect))

	// Register the handler for the containers page.
	http.HandleFunc(pages.ContainersPage, func(w http.ResponseWriter, r *http.Request) {
		err := pages.ServerContainersPage(containerManager, w, r.URL)
		if err != nil {
			fmt.Fprintf(w, "%s", err)
		}
	})

	// Start the manager.
	if err := containerManager.Start(); err != nil {
		glog.Fatalf("Failed to start container manager: %v", err)
	}

	// Install signal handler.
	installSignalHandler(containerManager)

	glog.Infof("Starting cAdvisor version: %q on port %d", info.VERSION, *argPort)

	addr := fmt.Sprintf(":%v", *argPort)
	glog.Fatal(http.ListenAndServe(addr, nil))
}

func setMaxProcs() {
	// TODO(vmarmol): Consider limiting if we have a CPU mask in effect.
	// Allow as many threads as we have cores unless the user specified a value.
	var numProcs int
	if *maxProcs < 1 {
		numProcs = runtime.NumCPU()
	} else {
		numProcs = *maxProcs
	}
	runtime.GOMAXPROCS(numProcs)

	// Check if the setting was successful.
	actualNumProcs := runtime.GOMAXPROCS(0)
	if actualNumProcs != numProcs {
		glog.Warningf("Specified max procs of %v but using %v", numProcs, actualNumProcs)
	}
}

func installSignalHandler(containerManager manager.Manager) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, os.Kill, syscall.SIGTERM)

	// Block until a signal is received.
	go func() {
		sig := <-c
		if err := containerManager.Stop(); err != nil {
			glog.Errorf("Failed to stop container manager: %v", err)
		}
		glog.Infof("Exiting given signal: %v", sig)
		os.Exit(0)
	}()
}
<|endoftext|>


<commit_before>/* Copyright 2012 Marc-Antoine Ruel. 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 main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"testing"
)

type mockCasTable struct {
	entries  map[string][]byte
	needFsck bool
	t        *testing.T
	log      *log.Logger
}

func (a *ApplicationMock) MakeCasTable(rootDir string) (CasTable, error) {
	return makeCasTable(rootDir)
	if a.cas == nil {
		a.cas = &mockCasTable{make(map[string][]byte), false, a.T, a.log}
	}
	return a.cas, nil
}

func (m *mockCasTable) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.log.Printf("mockCasTable.ServeHTTP(%s)", r.URL.Path)
	w.Write(m.entries[r.URL.Path[1:]])
}

func (m *mockCasTable) Enumerate() <-chan CasEntry {
	// First make a copy of the keys.
	keys := make([]string, len(m.entries))
	i := 0
	for k, _ := range m.entries {
		keys[i] = k
		i++
	}
	m.log.Printf("mockCasTable.Enumerate() %d", len(keys))
	c := make(chan CasEntry)
	go func() {
		for _, k := range keys {
			c <- CasEntry{Item: k}
		}
		close(c)
	}()
	return c
}

func (m *mockCasTable) AddEntry(source io.Reader, item string) error {
	m.log.Printf("mockCasTable.AddEntry(%s)", item)
	data, err := ioutil.ReadAll(source)
	if err == nil {
		m.entries[item] = data
	}
	return err
}

func (m *mockCasTable) Open(item string) (ReadSeekCloser, error) {
	m.log.Printf("mockCasTable.Open(%s)", item)
	data, ok := m.entries[item]
	if !ok {
		return nil, fmt.Errorf("Missing: %s", item)
	}
	return Buffer{bytes.NewReader(data)}, nil
}

func (m *mockCasTable) Remove(item string) error {
	m.log.Printf("mockCasTable.Remove(%s)", item)
	delete(m.entries, item)
	return nil
}

func (m *mockCasTable) SetFsckBit() {
	m.log.Printf("mockCasTable.SetFsckBit()")
	m.needFsck = true
}

func (m *mockCasTable) GetFsckBit() bool {
	m.log.Printf("mockCasTable.GetFsckBit() %t", m.needFsck)
	return m.needFsck
}

func (m *mockCasTable) ClearFsckBit() {
	m.log.Printf("mockCasTable.ClearFsckBit()")
	m.needFsck = false
}

// Adds noop Close() to a bytes.Reader.
type Buffer struct {
	*bytes.Reader
}

func (b Buffer) Close() error {
	return nil
}

func TestPrefixSpace(t *testing.T) {
	t.Parallel()
	type S struct {
		i int
		s string
	}
	checks := map[int]S{
		0: S{0, ""},
		1: S{16, "f"},
		2: S{256, "ff"},
		3: S{4096, "fff"},
		4: S{65536, "ffff"},
	}
	for prefixLength, s := range checks {
		x := prefixSpace(uint(prefixLength))
		if x != s.i {
			t.Fatalf("%d: %d != %d", prefixLength, x, s.i)
		}
		if x != 0 {
			res := fmt.Sprintf("%0*x", prefixLength, x-1)
			if res != s.s {
				t.Fatalf("%d: %s != %s", prefixLength, res, s.s)
			}
		}
	}
}

func TestCasTableImpl(t *testing.T) {
	t.Parallel()
	tempData, err := makeTempDir("cas")
	if err != nil {
		t.Fatalf("Failed to create tempdir", err)
	}
	defer removeTempDir(tempData)

	cas, err := makeCasTable(tempData)
	if err != nil {
		t.Fatal(err)
	}
	testCasTableImpl(t, cas)
}

func TestCasTableNode(t *testing.T) {
	t.Parallel()

	log := getLog(false)
	cas := &mockCasTable{make(map[string][]byte), false, t, log}
	testCasTableImpl(t, cas)
}

func testCasTableImpl(t *testing.T, cas CasTable) {
	for _ = range cas.Enumerate() {
		t.Fatal("Found unexpected value")
	}

	file1, err := AddBytes(cas, []byte("content1"))
	if err != nil {
		t.Fatal(err)
	}

	count := 0
	for v := range cas.Enumerate() {
		if v.Item != file1 {
			t.Fatal("Found unexpected value")
		}
		count++
	}
	if count != 1 {
		t.Fatalf("Found %d items", count)
	}

	if cas.GetFsckBit() {
		t.Fatal("Unexpected")
	}
	cas.SetFsckBit()
	if !cas.GetFsckBit() {
		t.Fatal("Unexpected")
	}
	cas.ClearFsckBit()
	if cas.GetFsckBit() {
		t.Fatal("Unexpected")
	}
}
<commit_msg>Improve CasTable testing and its mock<commit_after>/* Copyright 2012 Marc-Antoine Ruel. 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 main

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"testing"
)

type mockCasTable struct {
	entries  map[string][]byte
	needFsck bool
	t        *testing.T
	log      *log.Logger
}

func (a *ApplicationMock) MakeCasTable(rootDir string) (CasTable, error) {
	return makeCasTable(rootDir)
	if a.cas == nil {
		a.cas = &mockCasTable{make(map[string][]byte), false, a.T, a.log}
	}
	return a.cas, nil
}

func (m *mockCasTable) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	m.log.Printf("mockCasTable.ServeHTTP(%s)", r.URL.Path)
	w.Write(m.entries[r.URL.Path[1:]])
}

func (m *mockCasTable) Enumerate() <-chan CasEntry {
	// First make a copy of the keys.
	keys := make([]string, len(m.entries))
	i := 0
	for k, _ := range m.entries {
		keys[i] = k
		i++
	}
	m.log.Printf("mockCasTable.Enumerate() %d", len(keys))
	c := make(chan CasEntry)
	go func() {
		for _, k := range keys {
			c <- CasEntry{Item: k}
		}
		close(c)
	}()
	return c
}

func (m *mockCasTable) AddEntry(source io.Reader, item string) error {
	m.log.Printf("mockCasTable.AddEntry(%s)", item)
	if _, ok := m.entries[item]; ok {
		return os.ErrExist
	}
	data, err := ioutil.ReadAll(source)
	if err == nil {
		m.entries[item] = data
	}
	return err
}

func (m *mockCasTable) Open(item string) (ReadSeekCloser, error) {
	m.log.Printf("mockCasTable.Open(%s)", item)
	data, ok := m.entries[item]
	if !ok {
		return nil, fmt.Errorf("Missing: %s", item)
	}
	return Buffer{bytes.NewReader(data)}, nil
}

func (m *mockCasTable) Remove(item string) error {
	m.log.Printf("mockCasTable.Remove(%s)", item)
	if _, ok := m.entries[item]; !ok {
		return os.ErrNotExist
	}
	delete(m.entries, item)
	return nil
}

func (m *mockCasTable) SetFsckBit() {
	m.log.Printf("mockCasTable.SetFsckBit()")
	m.needFsck = true
}

func (m *mockCasTable) GetFsckBit() bool {
	m.log.Printf("mockCasTable.GetFsckBit() %t", m.needFsck)
	return m.needFsck
}

func (m *mockCasTable) ClearFsckBit() {
	m.log.Printf("mockCasTable.ClearFsckBit()")
	m.needFsck = false
}

// Adds noop Close() to a bytes.Reader.
type Buffer struct {
	*bytes.Reader
}

func (b Buffer) Close() error {
	return nil
}

func TestPrefixSpace(t *testing.T) {
	t.Parallel()
	type S struct {
		i int
		s string
	}
	checks := map[int]S{
		0: S{0, ""},
		1: S{16, "f"},
		2: S{256, "ff"},
		3: S{4096, "fff"},
		4: S{65536, "ffff"},
	}
	for prefixLength, s := range checks {
		x := prefixSpace(uint(prefixLength))
		if x != s.i {
			t.Fatalf("%d: %d != %d", prefixLength, x, s.i)
		}
		if x != 0 {
			res := fmt.Sprintf("%0*x", prefixLength, x-1)
			if res != s.s {
				t.Fatalf("%d: %s != %s", prefixLength, res, s.s)
			}
		}
	}
}

func TestCasTableImpl(t *testing.T) {
	t.Parallel()
	tempData, err := makeTempDir("cas")
	if err != nil {
		t.Fatalf("Failed to create tempdir", err)
	}
	defer removeTempDir(tempData)

	cas, err := makeCasTable(tempData)
	if err != nil {
		t.Fatal(err)
	}
	testCasTableImpl(t, cas)
}

func TestCasTableNode(t *testing.T) {
	t.Parallel()

	log := getLog(false)
	cas := &mockCasTable{make(map[string][]byte), false, t, log}
	testCasTableImpl(t, cas)
}

func testCasTableImpl(t *testing.T, cas CasTable) {
	for _ = range cas.Enumerate() {
		t.Fatal("Found unexpected value")
	}

	file1, err := AddBytes(cas, []byte("content1"))
	if err != nil {
		t.Fatal(err)
	}

	count := 0
	for v := range cas.Enumerate() {
		if v.Item != file1 {
			t.Fatal("Found unexpected value")
		}
		count++
	}
	if count != 1 {
		t.Fatalf("Found %d items", count)
	}

	// Add the same content.
	file2, err := AddBytes(cas, []byte("content1"))
	if !os.IsExist(err) {
		t.Fatal(err)
	}
	if file1 != file2 {
		t.Fatal("Hash mismatch")
	}
	count = 0
	for v := range cas.Enumerate() {
		if v.Item != file1 {
			t.Fatal("Found unexpected value")
		}
		count++
	}
	if count != 1 {
		t.Fatalf("Found %d items", count)
	}

	f, err := cas.Open(file1)
	if err != nil {
		t.Fatal(err)
	}
	data, err := ioutil.ReadAll(f)
	f.Close()
	if err != nil {
		t.Fatal(err)
	}
	if string(data) != "content1" {
		t.Fatal(string(data))
	}
	_, err = cas.Open("0")
	if err == nil {
		t.Fatal("Unexpected success")
	}
	if err := cas.Remove(file1); err != nil {
		t.Fatal(err)
	}
	if err := cas.Remove(file1); err == nil {
		t.Fatal("Unexpected success")
	}

	// Test fsck bit.
	if cas.GetFsckBit() {
		t.Fatal("Unexpected")
	}
	cas.SetFsckBit()
	if !cas.GetFsckBit() {
		t.Fatal("Unexpected")
	}
	cas.ClearFsckBit()
	if cas.GetFsckBit() {
		t.Fatal("Unexpected")
	}

}
<|endoftext|>