<commit_before>// +build libvirt

package libvirt

import (
	"strings"

	libvirt "github.com/libvirt/libvirt-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	"github.com/openshift/installer/pkg/destroy"
	"github.com/openshift/installer/pkg/types"
)

// filterFunc allows filtering based on names.
// returns true, when the name should be handled.
type filterFunc func(name string) bool

// ClusterIDPrefixFilter returns true for names
// that are prefixed with clusterid.
// `clusterid` cannot be empty.
var ClusterIDPrefixFilter = func(clusterid string) filterFunc {
	if clusterid == "" {
		panic("clusterid cannot be empty")
	}
	return func(name string) bool {
		return strings.HasPrefix(name, clusterid)
	}
}

// AlwaysTrueFilter returns true for all
// names except `default`.
var AlwaysTrueFilter = func() filterFunc {
	return func(name string) bool {
		return name != "default"
	}
}

// deleteFunc is the interface a function needs to implement to be delete resources.
type deleteFunc func(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error

// ClusterUninstaller holds the various options for the cluster we want to delete.
type ClusterUninstaller struct {
	LibvirtURI string
	Filter     filterFunc
	Logger     logrus.FieldLogger
}

// Run is the entrypoint to start the uninstall process.
func (o *ClusterUninstaller) Run() error {
	conn, err := libvirt.NewConnect(o.LibvirtURI)
	if err != nil {
		return errors.Wrap(err, "failed to connect to Libvirt daemon")
	}

	for _, del := range []deleteFunc{
		deleteDomains,
		deleteNetwork,
		deleteVolumes,
	} {
		err = del(conn, o.Filter, o.Logger)
		if err != nil {
			return err
		}
	}

	return nil
}

// deleteDomains calls deleteDomainsSinglePass until it finds no
// matching domains.  This guards against the machine-API launching
// additional nodes after the initial list call.  We continue deleting
// domains until we either hit an error or we have a list call with no
// matching domains.
func deleteDomains(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error {
	logger.Debug("Deleting libvirt domains")
	var err error
	nothingToDelete := false
	for !nothingToDelete {
		nothingToDelete, err = deleteDomainsSinglePass(conn, filter, logger)
		if err != nil {
			return err
		}
	}
	return nil
}

func deleteDomainsSinglePass(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) (nothingToDelete bool, err error) {
	domains, err := conn.ListAllDomains(0)
	if err != nil {
		return false, errors.Wrap(err, "list domains")
	}

	nothingToDelete = true
	for _, domain := range domains {
		defer domain.Free()
		dName, err := domain.GetName()
		if err != nil {
			return false, errors.Wrap(err, "get domain name")
		}
		if !filter(dName) {
			continue
		}

		nothingToDelete = false
		dState, _, err := domain.GetState()
		if err != nil {
			return false, errors.Wrapf(err, "get domain state %d", dName)
		}

		if dState != libvirt.DOMAIN_SHUTOFF && dState != libvirt.DOMAIN_SHUTDOWN {
			if err := domain.Destroy(); err != nil {
				return false, errors.Wrapf(err, "destroy domain %q", dName)
			}
		}
		if err := domain.Undefine(); err != nil {
			return false, errors.Wrapf(err, "undefine domain %q", dName)
		}
		logger.WithField("domain", dName).Info("Deleted domain")
	}

	return nothingToDelete, nil
}

func deleteVolumes(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error {
	logger.Debug("Deleting libvirt volumes")

	pools, err := conn.ListStoragePools()
	if err != nil {
		return errors.Wrap(err, "list storage pools")
	}

	tpool := "default"
	for _, pname := range pools {
		// pool name that returns true from filter, override default.
		if filter(pname) {
			tpool = pname
		}
	}
	pool, err := conn.LookupStoragePoolByName(tpool)
	if err != nil {
		return errors.Wrapf(err, "get storage pool %q", tpool)
	}
	defer pool.Free()

	switch tpool {
	case "default":
		// delete all vols that return true from filter.
		vols, err := pool.ListAllStorageVolumes(0)
		if err != nil {
			return errors.Wrapf(err, "list volumes in %q", tpool)
		}

		for _, vol := range vols {
			defer vol.Free()
			vName, err := vol.GetName()
			if err != nil {
				return errors.Wrapf(err, "get volume names in %q", tpool)
			}
			if !filter(vName) {
				continue
			}
			if err := vol.Delete(0); err != nil {
				return errors.Wrapf(err, "delete volume %q from %q", vName, tpool)
			}
			logger.WithField("volume", vName).Info("Deleted volume")
		}
	default:
		// blow away entire pool.
		if err := pool.Destroy(); err != nil {
			return errors.Wrapf(err, "destroy pool %q", tpool)
		}

		if err := pool.Undefine(); err != nil {
			return errors.Wrapf(err, "undefine pool %q", tpool)
		}
		logger.WithField("pool", tpool).Info("Deleted pool")
	}

	return nil
}

func deleteNetwork(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error {
	logger.Debug("Deleting libvirt network")

	networks, err := conn.ListNetworks()
	if err != nil {
		return errors.Wrap(err, "list networks")
	}

	for _, nName := range networks {
		if !filter(nName) {
			continue
		}
		network, err := conn.LookupNetworkByName(nName)
		if err != nil {
			return errors.Wrapf(err, "get network %q", nName)
		}
		defer network.Free()

		if err := network.Destroy(); err != nil {
			return errors.Wrapf(err, "destroy network %q", nName)
		}

		if err := network.Undefine(); err != nil {
			return errors.Wrapf(err, "undefine network %q", nName)
		}
		logger.WithField("network", nName).Info("Deleted network")
	}
	return nil
}

// New returns libvirt Uninstaller from ClusterMetadata.
func New(logger logrus.FieldLogger, metadata *types.ClusterMetadata) (destroy.Destroyer, error) {
	return &ClusterUninstaller{
		LibvirtURI: metadata.ClusterPlatformMetadata.Libvirt.URI,
		Filter:     ClusterIDPrefixFilter(metadata.ClusterID),
		Logger:     logger,
	}, nil
}
<commit_msg>libvirt: Pass correct InfraID prefix to destroy<commit_after>// +build libvirt

package libvirt

import (
	"strings"

	libvirt "github.com/libvirt/libvirt-go"
	"github.com/pkg/errors"
	"github.com/sirupsen/logrus"

	"github.com/openshift/installer/pkg/destroy"
	"github.com/openshift/installer/pkg/types"
)

// filterFunc allows filtering based on names.
// returns true, when the name should be handled.
type filterFunc func(name string) bool

// ClusterIDPrefixFilter returns true for names
// that are prefixed with clusterid.
// `clusterid` cannot be empty.
var ClusterIDPrefixFilter = func(clusterid string) filterFunc {
	if clusterid == "" {
		panic("clusterid cannot be empty")
	}
	return func(name string) bool {
		return strings.HasPrefix(name, clusterid)
	}
}

// AlwaysTrueFilter returns true for all
// names except `default`.
var AlwaysTrueFilter = func() filterFunc {
	return func(name string) bool {
		return name != "default"
	}
}

// deleteFunc is the interface a function needs to implement to be delete resources.
type deleteFunc func(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error

// ClusterUninstaller holds the various options for the cluster we want to delete.
type ClusterUninstaller struct {
	LibvirtURI string
	Filter     filterFunc
	Logger     logrus.FieldLogger
}

// Run is the entrypoint to start the uninstall process.
func (o *ClusterUninstaller) Run() error {
	conn, err := libvirt.NewConnect(o.LibvirtURI)
	if err != nil {
		return errors.Wrap(err, "failed to connect to Libvirt daemon")
	}

	for _, del := range []deleteFunc{
		deleteDomains,
		deleteNetwork,
		deleteVolumes,
	} {
		err = del(conn, o.Filter, o.Logger)
		if err != nil {
			return err
		}
	}

	return nil
}

// deleteDomains calls deleteDomainsSinglePass until it finds no
// matching domains.  This guards against the machine-API launching
// additional nodes after the initial list call.  We continue deleting
// domains until we either hit an error or we have a list call with no
// matching domains.
func deleteDomains(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error {
	logger.Debug("Deleting libvirt domains")
	var err error
	nothingToDelete := false
	for !nothingToDelete {
		nothingToDelete, err = deleteDomainsSinglePass(conn, filter, logger)
		if err != nil {
			return err
		}
	}
	return nil
}

func deleteDomainsSinglePass(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) (nothingToDelete bool, err error) {
	domains, err := conn.ListAllDomains(0)
	if err != nil {
		return false, errors.Wrap(err, "list domains")
	}

	nothingToDelete = true
	for _, domain := range domains {
		defer domain.Free()
		dName, err := domain.GetName()
		if err != nil {
			return false, errors.Wrap(err, "get domain name")
		}
		if !filter(dName) {
			continue
		}

		nothingToDelete = false
		dState, _, err := domain.GetState()
		if err != nil {
			return false, errors.Wrapf(err, "get domain state %d", dName)
		}

		if dState != libvirt.DOMAIN_SHUTOFF && dState != libvirt.DOMAIN_SHUTDOWN {
			if err := domain.Destroy(); err != nil {
				return false, errors.Wrapf(err, "destroy domain %q", dName)
			}
		}
		if err := domain.Undefine(); err != nil {
			return false, errors.Wrapf(err, "undefine domain %q", dName)
		}
		logger.WithField("domain", dName).Info("Deleted domain")
	}

	return nothingToDelete, nil
}

func deleteVolumes(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error {
	logger.Debug("Deleting libvirt volumes")

	pools, err := conn.ListStoragePools()
	if err != nil {
		return errors.Wrap(err, "list storage pools")
	}

	tpool := "default"
	for _, pname := range pools {
		// pool name that returns true from filter, override default.
		if filter(pname) {
			tpool = pname
		}
	}
	pool, err := conn.LookupStoragePoolByName(tpool)
	if err != nil {
		return errors.Wrapf(err, "get storage pool %q", tpool)
	}
	defer pool.Free()

	switch tpool {
	case "default":
		// delete all vols that return true from filter.
		vols, err := pool.ListAllStorageVolumes(0)
		if err != nil {
			return errors.Wrapf(err, "list volumes in %q", tpool)
		}

		for _, vol := range vols {
			defer vol.Free()
			vName, err := vol.GetName()
			if err != nil {
				return errors.Wrapf(err, "get volume names in %q", tpool)
			}
			if !filter(vName) {
				continue
			}
			if err := vol.Delete(0); err != nil {
				return errors.Wrapf(err, "delete volume %q from %q", vName, tpool)
			}
			logger.WithField("volume", vName).Info("Deleted volume")
		}
	default:
		// blow away entire pool.
		if err := pool.Destroy(); err != nil {
			return errors.Wrapf(err, "destroy pool %q", tpool)
		}

		if err := pool.Undefine(); err != nil {
			return errors.Wrapf(err, "undefine pool %q", tpool)
		}
		logger.WithField("pool", tpool).Info("Deleted pool")
	}

	return nil
}

func deleteNetwork(conn *libvirt.Connect, filter filterFunc, logger logrus.FieldLogger) error {
	logger.Debug("Deleting libvirt network")

	networks, err := conn.ListNetworks()
	if err != nil {
		return errors.Wrap(err, "list networks")
	}

	for _, nName := range networks {
		if !filter(nName) {
			continue
		}
		network, err := conn.LookupNetworkByName(nName)
		if err != nil {
			return errors.Wrapf(err, "get network %q", nName)
		}
		defer network.Free()

		if err := network.Destroy(); err != nil {
			return errors.Wrapf(err, "destroy network %q", nName)
		}

		if err := network.Undefine(); err != nil {
			return errors.Wrapf(err, "undefine network %q", nName)
		}
		logger.WithField("network", nName).Info("Deleted network")
	}
	return nil
}

// New returns libvirt Uninstaller from ClusterMetadata.
func New(logger logrus.FieldLogger, metadata *types.ClusterMetadata) (destroy.Destroyer, error) {
	return &ClusterUninstaller{
		LibvirtURI: metadata.ClusterPlatformMetadata.Libvirt.URI,
		Filter:     ClusterIDPrefixFilter(metadata.InfraID),
		Logger:     logger,
	}, nil
}
<|endoftext|>


<commit_before>package manipulators

import (
	"testing"

	"github.com/MJKWoolnough/minecraft"
)

func TestAreaDimensions(t *testing.T) {
	l, _ := minecraft.NewLevel(minecraft.NewMemPath())
	defer l.Close()
	tests := []struct {
		x1, y1, z1, x2, y2, z2, w, h, d int32
	}{
		{0, 0, 0, 0, 0, 0, 1, 1, 1},
		{-1, -1, -1, -1, -1, -1, 1, 1, 1},
		{-1, -1, -1, 0, 0, 0, 2, 2, 2},
		{64, 3, -5, 70, 1, -10, 7, 3, 6},
	}

	for n, test := range tests {
		a := NewArea(test.x1, test.y1, test.z1, test.x2, test.y2, test.z2, l)
		if a.Width() != test.w {
			t.Errorf("test %d: expecting width of %d, got %d", n+1, test.w, a.Width())
		}
		if a.Height() != test.h {
			t.Errorf("test %d: expecting height of %d, got %d", n+1, test.h, a.Height())
		}
		if a.Depth() != test.d {
			t.Errorf("test %d: expecting depth of %d, got %d", n+1, test.d, a.Depth())
		}
	}
}
<commit_msg>Added more simple tests<commit_after>package manipulators

import (
	"testing"

	"github.com/MJKWoolnough/minecraft"
)

func TestAreaDimensions(t *testing.T) {
	t.Parallel()
	l, _ := minecraft.NewLevel(minecraft.NewMemPath())
	defer l.Close()
	tests := []struct {
		x1, y1, z1, x2, y2, z2, w, h, d int32
	}{
		{0, 0, 0, 0, 0, 0, 1, 1, 1},
		{-1, -1, -1, -1, -1, -1, 1, 1, 1},
		{-1, -1, -1, 0, 0, 0, 2, 2, 2},
		{64, 3, -5, 70, 1, -10, 7, 3, 6},
	}

	for n, test := range tests {
		a := NewArea(test.x1, test.y1, test.z1, test.x2, test.y2, test.z2, l)
		if a.Width() != test.w {
			t.Errorf("test %d: expecting width of %d, got %d", n+1, test.w, a.Width())
		}
		if a.Height() != test.h {
			t.Errorf("test %d: expecting height of %d, got %d", n+1, test.h, a.Height())
		}
		if a.Depth() != test.d {
			t.Errorf("test %d: expecting depth of %d, got %d", n+1, test.d, a.Depth())
		}
	}
}

func TestAreaGetSet(t *testing.T) {
	t.Parallel()
	l, _ := minecraft.NewLevel(minecraft.NewMemPath())
	a1 := NewArea(3, 4, 5, 6, 7, 8, l)
	a2 := NewArea(-3, 12, -5, -6, 17, -8, l)
	tests := []struct {
		x1, y1, z1, x2, y2, z2, x3, y3, z3 int32
		b                                  minecraft.Block
	}{
		{0, 0, 0, 3, 4, 5, -6, 12, -8, minecraft.Block{ID: 1}},
		{1, 1, 1, 4, 5, 6, -5, 13, -7, minecraft.Block{ID: 2}},
	}

	for n, test := range tests {
		a1.Set(test.x1, test.y1, test.z1, test.b)
		a2.Set(test.x1, test.y1, test.z1, test.b)
		b, _ := a1.Get(test.x1, test.y1, test.z1)
		if !b.EqualBlock(test.b) {
			t.Errorf("test %d-1: incorrect block gotten", n+1)
		}
		b, _ = a2.Get(test.x1, test.y1, test.z1)
		if !b.EqualBlock(test.b) {
			t.Errorf("test %d-2: incorrect block gotten", n+1)
		}
		b, _ = l.GetBlock(test.x2, test.y2, test.z2)
		if !b.EqualBlock(test.b) {
			t.Errorf("test %d-3: incorrect block gotten", n+1)
		}
		b, _ = l.GetBlock(test.x3, test.y3, test.z3)
		if !b.EqualBlock(test.b) {
			t.Errorf("test %d-4: incorrect block gotten", n+1)
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// VPA collects CPU and memory usage measurements from all containers running in
// the cluster and aggregates them in memory in structures called
// AggregateContainerState.
// During aggregation the usage samples are grouped together by the key called
// AggregateStateKey and stored in structures such as histograms of CPU and
// memory usage, that are parts of the AggregateContainerState.
//
// The AggregateStateKey consists of the container name, the namespace and the
// set of labels on the pod the container belongs to. In other words, whenever
// two samples come from containers with the same name, in the same namespace
// and with the same pod labels, they end up in the same histogram.
//
// Recall that VPA produces one recommendation for all containers with a given
// name and namespace, having pod labels that match a given selector. Therefore
// for each VPA object and container name the recommender has to take all
// matching AggregateContainerStates and further aggregate them together, in
// order to obtain the final aggregation that is the input to the recommender
// function.

package model

import (
	"fmt"
	"math"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	vpa_types "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/apis/autoscaling.k8s.io/v1"
	"k8s.io/autoscaler/vertical-pod-autoscaler/pkg/recommender/util"
)

// ContainerNameToAggregateStateMap maps a container name to AggregateContainerState
// that aggregates state of containers with that name.
type ContainerNameToAggregateStateMap map[string]*AggregateContainerState

const (
	// SupportedCheckpointVersion is the tag of the supported version of serialized checkpoints.
	// Version id should be incremented on every non incompatible change, i.e. if the new
	// version of the recommender binary can't initialize from the old checkpoint format or the
	// previous version of the recommender binary can't initialize from the new checkpoint format.
	SupportedCheckpointVersion = "v3"
)

var (
	defaultControlledResources = []ResourceName{ResourceCPU, ResourceMemory}
)

// ContainerStateAggregator is an interface for objects that consume and
// aggregate container usage samples.
type ContainerStateAggregator interface {
	// AddSample aggregates a single usage sample.
	AddSample(sample *ContainerUsageSample)
	// SubtractSample removes a single usage sample. The subtracted sample
	// should be equal to some sample that was aggregated with AddSample()
	// in the past.
	SubtractSample(sample *ContainerUsageSample)
	// GetLastRecommendation returns last recommendation calculated for this
	// aggregator.
	GetLastRecommendation() corev1.ResourceList
	// NeedsRecommendation returns true if this aggregator should have
	// a recommendation calculated.
	NeedsRecommendation() bool
	// GetUpdateMode returns the update mode of VPA controlling this aggregator,
	// nil if aggregator is not autoscaled.
	GetUpdateMode() *vpa_types.UpdateMode
}

// AggregateContainerState holds input signals aggregated from a set of containers.
// It can be used as an input to compute the recommendation.
// The CPU and memory distributions use decaying histograms by default
// (see NewAggregateContainerState()).
// Implements ContainerStateAggregator interface.
type AggregateContainerState struct {
	// AggregateCPUUsage is a distribution of all CPU samples.
	AggregateCPUUsage util.Histogram
	// AggregateMemoryPeaks is a distribution of memory peaks from all containers:
	// each container should add one peak per memory aggregation interval (e.g. once every 24h).
	AggregateMemoryPeaks util.Histogram
	// Note: first/last sample timestamps as well as the sample count are based only on CPU samples.
	FirstSampleStart  time.Time
	LastSampleStart   time.Time
	TotalSamplesCount int
	CreationTime      time.Time

	// Following fields are needed to correctly report quality metrics
	// for VPA. When we record a new sample in an AggregateContainerState
	// we want to know if it needs recommendation, if the recommendation
	// is present and if the automatic updates are on (are we able to
	// apply the recommendation to the pods).
	LastRecommendation  corev1.ResourceList
	IsUnderVPA          bool
	UpdateMode          *vpa_types.UpdateMode
	ScalingMode         *vpa_types.ContainerScalingMode
	ControlledResources *[]ResourceName
}

// GetLastRecommendation returns last recorded recommendation.
func (a *AggregateContainerState) GetLastRecommendation() corev1.ResourceList {
	return a.LastRecommendation
}

// NeedsRecommendation returns true if the state should have recommendation calculated.
func (a *AggregateContainerState) NeedsRecommendation() bool {
	return a.IsUnderVPA && a.ScalingMode != nil && *a.ScalingMode != vpa_types.ContainerScalingModeOff
}

// GetUpdateMode returns the update mode of VPA controlling this aggregator,
// nil if aggregator is not autoscaled.
func (a *AggregateContainerState) GetUpdateMode() *vpa_types.UpdateMode {
	return a.UpdateMode
}

// GetScalingMode returns the container scaling mode of the container
// represented byt his aggregator, nil if aggregator is not autoscaled.
func (a *AggregateContainerState) GetScalingMode() *vpa_types.ContainerScalingMode {
	return a.ScalingMode
}

// GetControlledResources returns the list of resources controlled by VPA controlling this aggregator.
// Returns default if not set.
func (a *AggregateContainerState) GetControlledResources() []ResourceName {
	if a.ControlledResources != nil {
		return *a.ControlledResources
	}
	return defaultControlledResources
}

// MarkNotAutoscaled registers that this container state is not controlled by
// a VPA object.
func (a *AggregateContainerState) MarkNotAutoscaled() {
	a.IsUnderVPA = false
	a.LastRecommendation = nil
	a.UpdateMode = nil
	a.ScalingMode = nil
	a.ControlledResources = nil
}

// MergeContainerState merges two AggregateContainerStates.
func (a *AggregateContainerState) MergeContainerState(other *AggregateContainerState) {
	a.AggregateCPUUsage.Merge(other.AggregateCPUUsage)
	a.AggregateMemoryPeaks.Merge(other.AggregateMemoryPeaks)

	if !other.FirstSampleStart.IsZero() && other.FirstSampleStart.Before(a.FirstSampleStart) {
		a.FirstSampleStart = other.FirstSampleStart
	}
	if other.LastSampleStart.After(a.LastSampleStart) {
		a.LastSampleStart = other.LastSampleStart
	}
	a.TotalSamplesCount += other.TotalSamplesCount
}

// NewAggregateContainerState returns a new, empty AggregateContainerState.
func NewAggregateContainerState() *AggregateContainerState {
	return &AggregateContainerState{
		AggregateCPUUsage:    util.NewDecayingHistogram(CPUHistogramOptions, CPUHistogramDecayHalfLife),
		AggregateMemoryPeaks: util.NewDecayingHistogram(MemoryHistogramOptions, MemoryHistogramDecayHalfLife),
		CreationTime:         time.Now(),
	}
}

// AddSample aggregates a single usage sample.
func (a *AggregateContainerState) AddSample(sample *ContainerUsageSample) {
	switch sample.Resource {
	case ResourceCPU:
		a.addCPUSample(sample)
	case ResourceMemory:
		a.AggregateMemoryPeaks.AddSample(BytesFromMemoryAmount(sample.Usage), 1.0, sample.MeasureStart)
	default:
		panic(fmt.Sprintf("AddSample doesn't support resource '%s'", sample.Resource))
	}
}

// SubtractSample removes a single usage sample from an aggregation.
// The subtracted sample should be equal to some sample that was aggregated with
// AddSample() in the past.
// Only memory samples can be subtracted at the moment. Support for CPU could be
// added if necessary.
func (a *AggregateContainerState) SubtractSample(sample *ContainerUsageSample) {
	switch sample.Resource {
	case ResourceMemory:
		a.AggregateMemoryPeaks.SubtractSample(BytesFromMemoryAmount(sample.Usage), 1.0, sample.MeasureStart)
	default:
		panic(fmt.Sprintf("SubtractSample doesn't support resource '%s'", sample.Resource))
	}
}

func (a *AggregateContainerState) addCPUSample(sample *ContainerUsageSample) {
	cpuUsageCores := CoresFromCPUAmount(sample.Usage)
	cpuRequestCores := CoresFromCPUAmount(sample.Request)
	// Samples are added with the weight equal to the current request. This means that
	// whenever the request is increased, the history accumulated so far effectively decays,
	// which helps react quickly to CPU starvation.
	a.AggregateCPUUsage.AddSample(
		cpuUsageCores, math.Max(cpuRequestCores, minSampleWeight), sample.MeasureStart)
	if sample.MeasureStart.After(a.LastSampleStart) {
		a.LastSampleStart = sample.MeasureStart
	}
	if a.FirstSampleStart.IsZero() || sample.MeasureStart.Before(a.FirstSampleStart) {
		a.FirstSampleStart = sample.MeasureStart
	}
	a.TotalSamplesCount++
}

// SaveToCheckpoint serializes AggregateContainerState as VerticalPodAutoscalerCheckpointStatus.
// The serialization may result in loss of precission of the histograms.
func (a *AggregateContainerState) SaveToCheckpoint() (*vpa_types.VerticalPodAutoscalerCheckpointStatus, error) {
	memory, err := a.AggregateMemoryPeaks.SaveToChekpoint()
	if err != nil {
		return nil, err
	}
	cpu, err := a.AggregateCPUUsage.SaveToChekpoint()
	if err != nil {
		return nil, err
	}
	return &vpa_types.VerticalPodAutoscalerCheckpointStatus{
		FirstSampleStart:  metav1.NewTime(a.FirstSampleStart),
		LastSampleStart:   metav1.NewTime(a.LastSampleStart),
		TotalSamplesCount: a.TotalSamplesCount,
		MemoryHistogram:   *memory,
		CPUHistogram:      *cpu,
		Version:           SupportedCheckpointVersion,
	}, nil
}

// LoadFromCheckpoint deserializes data from VerticalPodAutoscalerCheckpointStatus
// into the AggregateContainerState.
func (a *AggregateContainerState) LoadFromCheckpoint(checkpoint *vpa_types.VerticalPodAutoscalerCheckpointStatus) error {
	if checkpoint.Version != SupportedCheckpointVersion {
		return fmt.Errorf("unsuported checkpoint version %s", checkpoint.Version)
	}
	a.TotalSamplesCount = checkpoint.TotalSamplesCount
	a.FirstSampleStart = checkpoint.FirstSampleStart.Time
	a.LastSampleStart = checkpoint.LastSampleStart.Time
	err := a.AggregateMemoryPeaks.LoadFromCheckpoint(&checkpoint.MemoryHistogram)
	if err != nil {
		return err
	}
	err = a.AggregateCPUUsage.LoadFromCheckpoint(&checkpoint.CPUHistogram)
	if err != nil {
		return err
	}
	return nil
}

func (a *AggregateContainerState) isExpired(now time.Time) bool {
	if a.isEmpty() {
		return now.Sub(a.CreationTime) >= MemoryAggregationWindowLength
	}
	return now.Sub(a.LastSampleStart) >= MemoryAggregationWindowLength
}

func (a *AggregateContainerState) isEmpty() bool {
	return a.TotalSamplesCount == 0
}

// UpdateFromPolicy updates container state scaling mode and controlled resources based on resource
// policy of the VPA object.
func (a *AggregateContainerState) UpdateFromPolicy(resourcePolicy *vpa_types.ContainerResourcePolicy) {
	// ContainerScalingModeAuto is the default scaling mode
	scalingModeAuto := vpa_types.ContainerScalingModeAuto
	a.ScalingMode = &scalingModeAuto
	if resourcePolicy != nil && resourcePolicy.Mode != nil {
		a.ScalingMode = resourcePolicy.Mode
	}
	a.ControlledResources = &defaultControlledResources
	if resourcePolicy != nil && resourcePolicy.ControlledResources != nil {
		a.ControlledResources = resourceNamesApiToModel(*resourcePolicy.ControlledResources)
	}
}

// AggregateStateByContainerName takes a set of AggregateContainerStates and merge them
// grouping by the container name. The result is a map from the container name to the aggregation
// from all input containers with the given name.
func AggregateStateByContainerName(aggregateContainerStateMap aggregateContainerStatesMap) ContainerNameToAggregateStateMap {
	containerNameToAggregateStateMap := make(ContainerNameToAggregateStateMap)
	for aggregationKey, aggregation := range aggregateContainerStateMap {
		containerName := aggregationKey.ContainerName()
		aggregateContainerState, isInitialized := containerNameToAggregateStateMap[containerName]
		if !isInitialized {
			aggregateContainerState = NewAggregateContainerState()
			containerNameToAggregateStateMap[containerName] = aggregateContainerState
		}
		aggregateContainerState.MergeContainerState(aggregation)
	}
	return containerNameToAggregateStateMap
}

// ContainerStateAggregatorProxy is a wrapper for ContainerStateAggregator
// that creates CnontainerStateAgregator for container if it is no longer
// present in the cluster state.
type ContainerStateAggregatorProxy struct {
	containerID ContainerID
	cluster     *ClusterState
}

// NewContainerStateAggregatorProxy creates a ContainerStateAggregatorProxy
// pointing to the cluster state.
func NewContainerStateAggregatorProxy(cluster *ClusterState, containerID ContainerID) ContainerStateAggregator {
	return &ContainerStateAggregatorProxy{containerID, cluster}
}

// AddSample adds a container sample to the aggregator.
func (p *ContainerStateAggregatorProxy) AddSample(sample *ContainerUsageSample) {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	aggregator.AddSample(sample)
}

// SubtractSample subtracts a container sample from the aggregator.
func (p *ContainerStateAggregatorProxy) SubtractSample(sample *ContainerUsageSample) {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	aggregator.SubtractSample(sample)
}

// GetLastRecommendation returns last recorded recommendation.
func (p *ContainerStateAggregatorProxy) GetLastRecommendation() corev1.ResourceList {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.GetLastRecommendation()
}

// NeedsRecommendation returns true if the aggregator should have recommendation calculated.
func (p *ContainerStateAggregatorProxy) NeedsRecommendation() bool {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.NeedsRecommendation()
}

// GetUpdateMode returns update mode of VPA controlling the aggregator.
func (p *ContainerStateAggregatorProxy) GetUpdateMode() *vpa_types.UpdateMode {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.GetUpdateMode()
}

// GetScalingMode returns scaling mode of container represented by the aggregator.
func (p *ContainerStateAggregatorProxy) GetScalingMode() *vpa_types.ContainerScalingMode {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.GetScalingMode()
}
<commit_msg>fixing container state merge<commit_after>/*
Copyright 2018 The Kubernetes Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// VPA collects CPU and memory usage measurements from all containers running in
// the cluster and aggregates them in memory in structures called
// AggregateContainerState.
// During aggregation the usage samples are grouped together by the key called
// AggregateStateKey and stored in structures such as histograms of CPU and
// memory usage, that are parts of the AggregateContainerState.
//
// The AggregateStateKey consists of the container name, the namespace and the
// set of labels on the pod the container belongs to. In other words, whenever
// two samples come from containers with the same name, in the same namespace
// and with the same pod labels, they end up in the same histogram.
//
// Recall that VPA produces one recommendation for all containers with a given
// name and namespace, having pod labels that match a given selector. Therefore
// for each VPA object and container name the recommender has to take all
// matching AggregateContainerStates and further aggregate them together, in
// order to obtain the final aggregation that is the input to the recommender
// function.

package model

import (
	"fmt"
	"math"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	vpa_types "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/apis/autoscaling.k8s.io/v1"
	"k8s.io/autoscaler/vertical-pod-autoscaler/pkg/recommender/util"
)

// ContainerNameToAggregateStateMap maps a container name to AggregateContainerState
// that aggregates state of containers with that name.
type ContainerNameToAggregateStateMap map[string]*AggregateContainerState

const (
	// SupportedCheckpointVersion is the tag of the supported version of serialized checkpoints.
	// Version id should be incremented on every non incompatible change, i.e. if the new
	// version of the recommender binary can't initialize from the old checkpoint format or the
	// previous version of the recommender binary can't initialize from the new checkpoint format.
	SupportedCheckpointVersion = "v3"
)

var (
	defaultControlledResources = []ResourceName{ResourceCPU, ResourceMemory}
)

// ContainerStateAggregator is an interface for objects that consume and
// aggregate container usage samples.
type ContainerStateAggregator interface {
	// AddSample aggregates a single usage sample.
	AddSample(sample *ContainerUsageSample)
	// SubtractSample removes a single usage sample. The subtracted sample
	// should be equal to some sample that was aggregated with AddSample()
	// in the past.
	SubtractSample(sample *ContainerUsageSample)
	// GetLastRecommendation returns last recommendation calculated for this
	// aggregator.
	GetLastRecommendation() corev1.ResourceList
	// NeedsRecommendation returns true if this aggregator should have
	// a recommendation calculated.
	NeedsRecommendation() bool
	// GetUpdateMode returns the update mode of VPA controlling this aggregator,
	// nil if aggregator is not autoscaled.
	GetUpdateMode() *vpa_types.UpdateMode
}

// AggregateContainerState holds input signals aggregated from a set of containers.
// It can be used as an input to compute the recommendation.
// The CPU and memory distributions use decaying histograms by default
// (see NewAggregateContainerState()).
// Implements ContainerStateAggregator interface.
type AggregateContainerState struct {
	// AggregateCPUUsage is a distribution of all CPU samples.
	AggregateCPUUsage util.Histogram
	// AggregateMemoryPeaks is a distribution of memory peaks from all containers:
	// each container should add one peak per memory aggregation interval (e.g. once every 24h).
	AggregateMemoryPeaks util.Histogram
	// Note: first/last sample timestamps as well as the sample count are based only on CPU samples.
	FirstSampleStart  time.Time
	LastSampleStart   time.Time
	TotalSamplesCount int
	CreationTime      time.Time

	// Following fields are needed to correctly report quality metrics
	// for VPA. When we record a new sample in an AggregateContainerState
	// we want to know if it needs recommendation, if the recommendation
	// is present and if the automatic updates are on (are we able to
	// apply the recommendation to the pods).
	LastRecommendation  corev1.ResourceList
	IsUnderVPA          bool
	UpdateMode          *vpa_types.UpdateMode
	ScalingMode         *vpa_types.ContainerScalingMode
	ControlledResources *[]ResourceName
}

// GetLastRecommendation returns last recorded recommendation.
func (a *AggregateContainerState) GetLastRecommendation() corev1.ResourceList {
	return a.LastRecommendation
}

// NeedsRecommendation returns true if the state should have recommendation calculated.
func (a *AggregateContainerState) NeedsRecommendation() bool {
	return a.IsUnderVPA && a.ScalingMode != nil && *a.ScalingMode != vpa_types.ContainerScalingModeOff
}

// GetUpdateMode returns the update mode of VPA controlling this aggregator,
// nil if aggregator is not autoscaled.
func (a *AggregateContainerState) GetUpdateMode() *vpa_types.UpdateMode {
	return a.UpdateMode
}

// GetScalingMode returns the container scaling mode of the container
// represented byt his aggregator, nil if aggregator is not autoscaled.
func (a *AggregateContainerState) GetScalingMode() *vpa_types.ContainerScalingMode {
	return a.ScalingMode
}

// GetControlledResources returns the list of resources controlled by VPA controlling this aggregator.
// Returns default if not set.
func (a *AggregateContainerState) GetControlledResources() []ResourceName {
	if a.ControlledResources != nil {
		return *a.ControlledResources
	}
	return defaultControlledResources
}

// MarkNotAutoscaled registers that this container state is not controlled by
// a VPA object.
func (a *AggregateContainerState) MarkNotAutoscaled() {
	a.IsUnderVPA = false
	a.LastRecommendation = nil
	a.UpdateMode = nil
	a.ScalingMode = nil
	a.ControlledResources = nil
}

// MergeContainerState merges two AggregateContainerStates.
func (a *AggregateContainerState) MergeContainerState(other *AggregateContainerState) {
	a.AggregateCPUUsage.Merge(other.AggregateCPUUsage)
	a.AggregateMemoryPeaks.Merge(other.AggregateMemoryPeaks)

	if a.FirstSampleStart.IsZero() ||
		(!other.FirstSampleStart.IsZero() && other.FirstSampleStart.Before(a.FirstSampleStart)) {
		a.FirstSampleStart = other.FirstSampleStart
	}
	if other.LastSampleStart.After(a.LastSampleStart) {
		a.LastSampleStart = other.LastSampleStart
	}
	a.TotalSamplesCount += other.TotalSamplesCount
}

// NewAggregateContainerState returns a new, empty AggregateContainerState.
func NewAggregateContainerState() *AggregateContainerState {
	return &AggregateContainerState{
		AggregateCPUUsage:    util.NewDecayingHistogram(CPUHistogramOptions, CPUHistogramDecayHalfLife),
		AggregateMemoryPeaks: util.NewDecayingHistogram(MemoryHistogramOptions, MemoryHistogramDecayHalfLife),
		CreationTime:         time.Now(),
	}
}

// AddSample aggregates a single usage sample.
func (a *AggregateContainerState) AddSample(sample *ContainerUsageSample) {
	switch sample.Resource {
	case ResourceCPU:
		a.addCPUSample(sample)
	case ResourceMemory:
		a.AggregateMemoryPeaks.AddSample(BytesFromMemoryAmount(sample.Usage), 1.0, sample.MeasureStart)
	default:
		panic(fmt.Sprintf("AddSample doesn't support resource '%s'", sample.Resource))
	}
}

// SubtractSample removes a single usage sample from an aggregation.
// The subtracted sample should be equal to some sample that was aggregated with
// AddSample() in the past.
// Only memory samples can be subtracted at the moment. Support for CPU could be
// added if necessary.
func (a *AggregateContainerState) SubtractSample(sample *ContainerUsageSample) {
	switch sample.Resource {
	case ResourceMemory:
		a.AggregateMemoryPeaks.SubtractSample(BytesFromMemoryAmount(sample.Usage), 1.0, sample.MeasureStart)
	default:
		panic(fmt.Sprintf("SubtractSample doesn't support resource '%s'", sample.Resource))
	}
}

func (a *AggregateContainerState) addCPUSample(sample *ContainerUsageSample) {
	cpuUsageCores := CoresFromCPUAmount(sample.Usage)
	cpuRequestCores := CoresFromCPUAmount(sample.Request)
	// Samples are added with the weight equal to the current request. This means that
	// whenever the request is increased, the history accumulated so far effectively decays,
	// which helps react quickly to CPU starvation.
	a.AggregateCPUUsage.AddSample(
		cpuUsageCores, math.Max(cpuRequestCores, minSampleWeight), sample.MeasureStart)
	if sample.MeasureStart.After(a.LastSampleStart) {
		a.LastSampleStart = sample.MeasureStart
	}
	if a.FirstSampleStart.IsZero() || sample.MeasureStart.Before(a.FirstSampleStart) {
		a.FirstSampleStart = sample.MeasureStart
	}
	a.TotalSamplesCount++
}

// SaveToCheckpoint serializes AggregateContainerState as VerticalPodAutoscalerCheckpointStatus.
// The serialization may result in loss of precission of the histograms.
func (a *AggregateContainerState) SaveToCheckpoint() (*vpa_types.VerticalPodAutoscalerCheckpointStatus, error) {
	memory, err := a.AggregateMemoryPeaks.SaveToChekpoint()
	if err != nil {
		return nil, err
	}
	cpu, err := a.AggregateCPUUsage.SaveToChekpoint()
	if err != nil {
		return nil, err
	}
	return &vpa_types.VerticalPodAutoscalerCheckpointStatus{
		FirstSampleStart:  metav1.NewTime(a.FirstSampleStart),
		LastSampleStart:   metav1.NewTime(a.LastSampleStart),
		TotalSamplesCount: a.TotalSamplesCount,
		MemoryHistogram:   *memory,
		CPUHistogram:      *cpu,
		Version:           SupportedCheckpointVersion,
	}, nil
}

// LoadFromCheckpoint deserializes data from VerticalPodAutoscalerCheckpointStatus
// into the AggregateContainerState.
func (a *AggregateContainerState) LoadFromCheckpoint(checkpoint *vpa_types.VerticalPodAutoscalerCheckpointStatus) error {
	if checkpoint.Version != SupportedCheckpointVersion {
		return fmt.Errorf("unsuported checkpoint version %s", checkpoint.Version)
	}
	a.TotalSamplesCount = checkpoint.TotalSamplesCount
	a.FirstSampleStart = checkpoint.FirstSampleStart.Time
	a.LastSampleStart = checkpoint.LastSampleStart.Time
	err := a.AggregateMemoryPeaks.LoadFromCheckpoint(&checkpoint.MemoryHistogram)
	if err != nil {
		return err
	}
	err = a.AggregateCPUUsage.LoadFromCheckpoint(&checkpoint.CPUHistogram)
	if err != nil {
		return err
	}
	return nil
}

func (a *AggregateContainerState) isExpired(now time.Time) bool {
	if a.isEmpty() {
		return now.Sub(a.CreationTime) >= MemoryAggregationWindowLength
	}
	return now.Sub(a.LastSampleStart) >= MemoryAggregationWindowLength
}

func (a *AggregateContainerState) isEmpty() bool {
	return a.TotalSamplesCount == 0
}

// UpdateFromPolicy updates container state scaling mode and controlled resources based on resource
// policy of the VPA object.
func (a *AggregateContainerState) UpdateFromPolicy(resourcePolicy *vpa_types.ContainerResourcePolicy) {
	// ContainerScalingModeAuto is the default scaling mode
	scalingModeAuto := vpa_types.ContainerScalingModeAuto
	a.ScalingMode = &scalingModeAuto
	if resourcePolicy != nil && resourcePolicy.Mode != nil {
		a.ScalingMode = resourcePolicy.Mode
	}
	a.ControlledResources = &defaultControlledResources
	if resourcePolicy != nil && resourcePolicy.ControlledResources != nil {
		a.ControlledResources = resourceNamesApiToModel(*resourcePolicy.ControlledResources)
	}
}

// AggregateStateByContainerName takes a set of AggregateContainerStates and merge them
// grouping by the container name. The result is a map from the container name to the aggregation
// from all input containers with the given name.
func AggregateStateByContainerName(aggregateContainerStateMap aggregateContainerStatesMap) ContainerNameToAggregateStateMap {
	containerNameToAggregateStateMap := make(ContainerNameToAggregateStateMap)
	for aggregationKey, aggregation := range aggregateContainerStateMap {
		containerName := aggregationKey.ContainerName()
		aggregateContainerState, isInitialized := containerNameToAggregateStateMap[containerName]
		if !isInitialized {
			aggregateContainerState = NewAggregateContainerState()
			containerNameToAggregateStateMap[containerName] = aggregateContainerState
		}
		aggregateContainerState.MergeContainerState(aggregation)
	}
	return containerNameToAggregateStateMap
}

// ContainerStateAggregatorProxy is a wrapper for ContainerStateAggregator
// that creates CnontainerStateAgregator for container if it is no longer
// present in the cluster state.
type ContainerStateAggregatorProxy struct {
	containerID ContainerID
	cluster     *ClusterState
}

// NewContainerStateAggregatorProxy creates a ContainerStateAggregatorProxy
// pointing to the cluster state.
func NewContainerStateAggregatorProxy(cluster *ClusterState, containerID ContainerID) ContainerStateAggregator {
	return &ContainerStateAggregatorProxy{containerID, cluster}
}

// AddSample adds a container sample to the aggregator.
func (p *ContainerStateAggregatorProxy) AddSample(sample *ContainerUsageSample) {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	aggregator.AddSample(sample)
}

// SubtractSample subtracts a container sample from the aggregator.
func (p *ContainerStateAggregatorProxy) SubtractSample(sample *ContainerUsageSample) {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	aggregator.SubtractSample(sample)
}

// GetLastRecommendation returns last recorded recommendation.
func (p *ContainerStateAggregatorProxy) GetLastRecommendation() corev1.ResourceList {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.GetLastRecommendation()
}

// NeedsRecommendation returns true if the aggregator should have recommendation calculated.
func (p *ContainerStateAggregatorProxy) NeedsRecommendation() bool {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.NeedsRecommendation()
}

// GetUpdateMode returns update mode of VPA controlling the aggregator.
func (p *ContainerStateAggregatorProxy) GetUpdateMode() *vpa_types.UpdateMode {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.GetUpdateMode()
}

// GetScalingMode returns scaling mode of container represented by the aggregator.
func (p *ContainerStateAggregatorProxy) GetScalingMode() *vpa_types.ContainerScalingMode {
	aggregator := p.cluster.findOrCreateAggregateContainerState(p.containerID)
	return aggregator.GetScalingMode()
}
<|endoftext|>


<commit_before>package lightstep_test

import (
	"context"
	"fmt"
	"reflect"

	"github.com/lightstep/lightstep-tracer-common/golang/gogo/collectorpb"
	cpbfakes "github.com/lightstep/lightstep-tracer-common/golang/gogo/collectorpb/collectorpbfakes"
	. "github.com/lightstep/lightstep-tracer-go"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/types"
	"github.com/opentracing/opentracing-go"
)

func closeTestTracer(tracer opentracing.Tracer) {
	complete := make(chan struct{})
	go func() {
		Close(context.Background(), tracer)
		close(complete)
	}()
	Eventually(complete).Should(BeClosed())
}

func startNSpans(n int, tracer opentracing.Tracer) {
	for i := 0; i < n; i++ {
		tracer.StartSpan(string(i)).Finish()
	}
}

type haveKeyValuesMatcher []*collectorpb.KeyValue

func HaveKeyValues(keyValues ...*collectorpb.KeyValue) types.GomegaMatcher {
	return haveKeyValuesMatcher(keyValues)
}

func (matcher haveKeyValuesMatcher) Match(actual interface{}) (bool, error) {
	switch v := actual.(type) {
	case []*collectorpb.KeyValue:
		return matcher.MatchProtos(v)
	case *collectorpb.Log:
		return matcher.MatchProtos(v.GetFields())
	default:
		return false, fmt.Errorf("HaveKeyValues matcher expects either a []*KeyValue or a *Log/*LogRecord")
	}
}

func (matcher haveKeyValuesMatcher) MatchProtos(actualKeyValues []*collectorpb.KeyValue) (bool, error) {
	expectedKeyValues := []*collectorpb.KeyValue(matcher)
	if len(expectedKeyValues) != len(actualKeyValues) {
		return false, nil
	}

	for i := range actualKeyValues {
		if !reflect.DeepEqual(actualKeyValues[i], expectedKeyValues[i]) {
			return false, nil
		}
	}

	return true, nil
}

func (matcher haveKeyValuesMatcher) FailureMessage(actual interface{}) string {
	return fmt.Sprintf("Expected '%v' to have key values '%v'", actual, matcher)
}

func (matcher haveKeyValuesMatcher) NegatedFailureMessage(actual interface{}) string {
	return fmt.Sprintf("Expected '%v' to not have key values '%v'", actual, matcher)
}

func KeyValue(key string, value interface{}, storeAsJson ...bool) *collectorpb.KeyValue {
	tag := &collectorpb.KeyValue{Key: key}
	switch typedValue := value.(type) {
	case int:
		tag.Value = &collectorpb.KeyValue_IntValue{IntValue: int64(typedValue)}
	case string:
		if len(storeAsJson) > 0 && storeAsJson[0] {
			tag.Value = &collectorpb.KeyValue_JsonValue{JsonValue: typedValue}
		} else {
			tag.Value = &collectorpb.KeyValue_StringValue{StringValue: typedValue}
		}
	case bool:
		tag.Value = &collectorpb.KeyValue_BoolValue{BoolValue: typedValue}
	case float32:
		tag.Value = &collectorpb.KeyValue_DoubleValue{DoubleValue: float64(typedValue)}
	case float64:
		tag.Value = &collectorpb.KeyValue_DoubleValue{DoubleValue: typedValue}
	}
	return tag
}

//////////////////
// GRPC HELPERS //
//////////////////

func getReportedGRPCSpans(fakeClient *cpbfakes.FakeCollectorServiceClient) []*collectorpb.Span {
	callCount := fakeClient.ReportCallCount()
	spans := make([]*collectorpb.Span, 0)
	for i := 0; i < callCount; i++ {
		_, report, _ := fakeClient.ReportArgsForCall(i)
		spans = append(spans, report.GetSpans()...)
	}
	return spans
}

type dummyConnection struct{}

func (*dummyConnection) Close() error { return nil }

func fakeGrpcConnection(fakeClient *cpbfakes.FakeCollectorServiceClient) ConnectorFactory {
	return func() (interface{}, Connection, error) {
		return fakeClient, new(dummyConnection), nil
	}
}
<commit_msg>Convert id to string using strconv.Itoa<commit_after>package lightstep_test

import (
	"context"
	"fmt"
	"reflect"
    "strconv"

	"github.com/lightstep/lightstep-tracer-common/golang/gogo/collectorpb"
	cpbfakes "github.com/lightstep/lightstep-tracer-common/golang/gogo/collectorpb/collectorpbfakes"
	. "github.com/lightstep/lightstep-tracer-go"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/types"
	"github.com/opentracing/opentracing-go"
)

func closeTestTracer(tracer opentracing.Tracer) {
	complete := make(chan struct{})
	go func() {
		Close(context.Background(), tracer)
		close(complete)
	}()
	Eventually(complete).Should(BeClosed())
}

func startNSpans(n int, tracer opentracing.Tracer) {
	for i := 0; i < n; i++ {
		tracer.StartSpan(strconv.Itoa(i)).Finish()
	}
}

type haveKeyValuesMatcher []*collectorpb.KeyValue

func HaveKeyValues(keyValues ...*collectorpb.KeyValue) types.GomegaMatcher {
	return haveKeyValuesMatcher(keyValues)
}

func (matcher haveKeyValuesMatcher) Match(actual interface{}) (bool, error) {
	switch v := actual.(type) {
	case []*collectorpb.KeyValue:
		return matcher.MatchProtos(v)
	case *collectorpb.Log:
		return matcher.MatchProtos(v.GetFields())
	default:
		return false, fmt.Errorf("HaveKeyValues matcher expects either a []*KeyValue or a *Log/*LogRecord")
	}
}

func (matcher haveKeyValuesMatcher) MatchProtos(actualKeyValues []*collectorpb.KeyValue) (bool, error) {
	expectedKeyValues := []*collectorpb.KeyValue(matcher)
	if len(expectedKeyValues) != len(actualKeyValues) {
		return false, nil
	}

	for i := range actualKeyValues {
		if !reflect.DeepEqual(actualKeyValues[i], expectedKeyValues[i]) {
			return false, nil
		}
	}

	return true, nil
}

func (matcher haveKeyValuesMatcher) FailureMessage(actual interface{}) string {
	return fmt.Sprintf("Expected '%v' to have key values '%v'", actual, matcher)
}

func (matcher haveKeyValuesMatcher) NegatedFailureMessage(actual interface{}) string {
	return fmt.Sprintf("Expected '%v' to not have key values '%v'", actual, matcher)
}

func KeyValue(key string, value interface{}, storeAsJson ...bool) *collectorpb.KeyValue {
	tag := &collectorpb.KeyValue{Key: key}
	switch typedValue := value.(type) {
	case int:
		tag.Value = &collectorpb.KeyValue_IntValue{IntValue: int64(typedValue)}
	case string:
		if len(storeAsJson) > 0 && storeAsJson[0] {
			tag.Value = &collectorpb.KeyValue_JsonValue{JsonValue: typedValue}
		} else {
			tag.Value = &collectorpb.KeyValue_StringValue{StringValue: typedValue}
		}
	case bool:
		tag.Value = &collectorpb.KeyValue_BoolValue{BoolValue: typedValue}
	case float32:
		tag.Value = &collectorpb.KeyValue_DoubleValue{DoubleValue: float64(typedValue)}
	case float64:
		tag.Value = &collectorpb.KeyValue_DoubleValue{DoubleValue: typedValue}
	}
	return tag
}

//////////////////
// GRPC HELPERS //
//////////////////

func getReportedGRPCSpans(fakeClient *cpbfakes.FakeCollectorServiceClient) []*collectorpb.Span {
	callCount := fakeClient.ReportCallCount()
	spans := make([]*collectorpb.Span, 0)
	for i := 0; i < callCount; i++ {
		_, report, _ := fakeClient.ReportArgsForCall(i)
		spans = append(spans, report.GetSpans()...)
	}
	return spans
}

type dummyConnection struct{}

func (*dummyConnection) Close() error { return nil }

func fakeGrpcConnection(fakeClient *cpbfakes.FakeCollectorServiceClient) ConnectorFactory {
	return func() (interface{}, Connection, error) {
		return fakeClient, new(dummyConnection), nil
	}
}
<|endoftext|>


<commit_before>package ip2region

import (
	"fmt"
	"strings"

	"github.com/admpub/ip2region/binding/golang/ip2region"
	syncOnce "github.com/admpub/once"
	"github.com/webx-top/echo"
)

var (
	region   *ip2region.Ip2Region
	dictFile string
	once     syncOnce.Once
)

func init() {
	dictFile = echo.Wd() + echo.FilePathSeparator + `data` + echo.FilePathSeparator + `ip2region` + echo.FilePathSeparator + `ip2region.db`
}

func SetDictFile(f string) {
	dictFile = f
}

func SetInstance(newInstance *ip2region.Ip2Region) {
	if region == nil {
		region = newInstance
	} else {
		oldRegion := *region
		*region = *newInstance
		oldRegion.Close()
	}
}

func Initialize() (err error) {
	if region == nil {
		region, err = ip2region.New(dictFile)
	}
	return
}

func IsInitialized() bool {
	return region != nil
}

func IPInfo(ip string) (info ip2region.IpInfo, err error) {
	if len(ip) == 0 {
		return
	}
	once.Do(func() {
		err = Initialize()
	})
	if err != nil {
		return
	}
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf(`%v`, e)
		}
	}()
	info, err = region.MemorySearch(ip)
	return
}

func Stringify(info ip2region.IpInfo) string {
	var (
		formats []string
		args    []interface{}
	)
	if len(info.Country) > 0 && info.Country != `0` {
		formats = append(formats, `"国家":%q`)
		args = append(args, info.Country)
	}
	if len(info.Region) > 0 && info.Region != `0` {
		formats = append(formats, `"地区":%q`)
		args = append(args, info.Region)
	}
	if len(info.Province) > 0 && info.Province != `0` {
		formats = append(formats, `"省份":%q`)
		args = append(args, info.Province)
	}
	if len(info.City) > 0 && info.City != `0` {
		formats = append(formats, `"城市":%q`)
		args = append(args, info.City)
	}
	if len(info.ISP) > 0 && info.ISP != `0` {
		formats = append(formats, `"线路":%q`)
		args = append(args, info.ISP)
	}
	return fmt.Sprintf(`{`+strings.Join(formats, `,`)+`}`, args...)
}
<commit_msg>update<commit_after>package ip2region

import (
	"fmt"
	"strings"

	"github.com/admpub/ip2region/binding/golang/ip2region"
	syncOnce "github.com/admpub/once"
	"github.com/webx-top/echo"
)

var (
	region   *ip2region.Ip2Region
	dictFile string
	once     syncOnce.Once
)

func init() {
	dictFile = echo.Wd() + echo.FilePathSeparator + `data` + echo.FilePathSeparator + `ip2region` + echo.FilePathSeparator + `ip2region.db`
}

func SetDictFile(f string) {
	dictFile = f
	once.Reset()
}

func SetInstance(newInstance *ip2region.Ip2Region) {
	if region == nil {
		region = newInstance
	} else {
		oldRegion := *region
		*region = *newInstance
		oldRegion.Close()
	}
}

func initialize() (err error) {
	if region != nil {
		region.Close()
	}
	region, err = ip2region.New(dictFile)
	return
}

func IsInitialized() bool {
	return region != nil
}

func IPInfo(ip string) (info ip2region.IpInfo, err error) {
	if len(ip) == 0 {
		return
	}
	once.Do(func() {
		err = initialize()
	})
	if err != nil {
		return
	}
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf(`%v`, e)
		}
	}()
	info, err = region.MemorySearch(ip)
	return
}

func Stringify(info ip2region.IpInfo) string {
	var (
		formats []string
		args    []interface{}
	)
	if len(info.Country) > 0 && info.Country != `0` {
		formats = append(formats, `"国家":%q`)
		args = append(args, info.Country)
	}
	if len(info.Region) > 0 && info.Region != `0` {
		formats = append(formats, `"地区":%q`)
		args = append(args, info.Region)
	}
	if len(info.Province) > 0 && info.Province != `0` {
		formats = append(formats, `"省份":%q`)
		args = append(args, info.Province)
	}
	if len(info.City) > 0 && info.City != `0` {
		formats = append(formats, `"城市":%q`)
		args = append(args, info.City)
	}
	if len(info.ISP) > 0 && info.ISP != `0` {
		formats = append(formats, `"线路":%q`)
		args = append(args, info.ISP)
	}
	return fmt.Sprintf(`{`+strings.Join(formats, `,`)+`}`, args...)
}
<|endoftext|>


<commit_before>package mail

import (
	"crypto/tls"
	"github.com/hackform/governor"
	"github.com/labstack/echo"
	"github.com/sirupsen/logrus"
	"gopkg.in/gomail.v2"
	"net/http"
	"time"
)

type (
	// Mail is a service wrapper around a mailer instance
	Mail interface {
		governor.Service
		Send(to, subject, body string) *governor.Error
	}

	goMail struct {
		host        string
		port        int
		username    string
		password    string
		insecure    bool
		bufferSize  int
		workerSize  int
		fromAddress string
		msgc        chan *gomail.Message
	}
)

const (
	moduleID = "mail"
)

// New creates a new mailer service
func New(c governor.Config, l *logrus.Logger) Mail {
	v := c.Conf()
	rconf := v.GetStringMapString("mail")

	l.Info("initialized mail service")

	return &goMail{
		host:        rconf["host"],
		port:        v.GetInt("mail.port"),
		username:    rconf["username"],
		password:    rconf["password"],
		insecure:    v.GetBool("mail.insecure"),
		bufferSize:  v.GetInt("mail.buffer_size"),
		workerSize:  v.GetInt("mail.worker_size"),
		fromAddress: rconf["from_address"],
		msgc:        make(chan *gomail.Message),
	}
}

func (m *goMail) dialer() *gomail.Dialer {
	d := gomail.NewDialer(m.host, m.port, m.username, m.password)

	if m.insecure {
		d.TLSConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	return d
}

func (m *goMail) mailWorker(l *logrus.Logger, ch <-chan *gomail.Message) {
	d := m.dialer()

	var s gomail.SendCloser
	var err error
	open := false
	for {
		select {
		case m, ok := <-ch:
			if !ok {
				return
			}
			if !open {
				if s, err = d.Dial(); err != nil {
					panic(err)
				}
				open = true
			}
			if err := gomail.Send(s, m); err != nil {
				l.Error(err)
			}
		// Close the connection to the SMTP server if no email was sent in
		// the last 30 seconds.
		case <-time.After(30 * time.Second):
			if open {
				if err := s.Close(); err != nil {
					panic(err)
				}
				open = false
			}
		}
	}
}

// Mount is a place to mount routes to satisfy the Service interface
func (m *goMail) Mount(conf governor.Config, r *echo.Group, l *logrus.Logger) error {
	l.Info("mounted mail service")
	return nil
}

// Health is a health check for the service
func (m *goMail) Health() *governor.Error {
	return nil
}

// Setup is run on service setup
func (m *goMail) Setup(conf governor.Config, l *logrus.Logger, rsetup governor.ReqSetupPost) *governor.Error {
	return nil
}

const (
	moduleIDSend = moduleID + ".send"
)

// Send creates and sends a new message
func (m *goMail) Send(to, subject, body string) *governor.Error {
	msg := gomail.NewMessage()
	msg.SetHeader("From", m.fromAddress)
	msg.SetHeader("To", to)
	msg.SetHeader("Subject", subject)
	msg.SetBody("text/html", body)

	//if err := m.mailer.DialAndSend(msg); err != nil {
	//	return governor.NewError(moduleIDSend, err.Error(), 0, http.StatusInternalServerError)
	//}

	return nil
}
<commit_msg>mail service now using workers<commit_after>package mail

import (
	"crypto/tls"
	"github.com/hackform/governor"
	"github.com/labstack/echo"
	"github.com/sirupsen/logrus"
	"gopkg.in/gomail.v2"
	"net/http"
	"time"
)

type (
	// Mail is a service wrapper around a mailer instance
	Mail interface {
		governor.Service
		Send(to, subject, body string) *governor.Error
	}

	goMail struct {
		host        string
		port        int
		username    string
		password    string
		insecure    bool
		bufferSize  int
		workerSize  int
		fromAddress string
		msgc        chan *gomail.Message
	}
)

const (
	moduleID = "mail"
)

// New creates a new mailer service
func New(c governor.Config, l *logrus.Logger) Mail {
	v := c.Conf()
	rconf := v.GetStringMapString("mail")

	l.Info("initialized mail service")

	gm := &goMail{
		host:        rconf["host"],
		port:        v.GetInt("mail.port"),
		username:    rconf["username"],
		password:    rconf["password"],
		insecure:    v.GetBool("mail.insecure"),
		bufferSize:  v.GetInt("mail.buffer_size"),
		workerSize:  v.GetInt("mail.worker_size"),
		fromAddress: rconf["from_address"],
		msgc:        make(chan *gomail.Message, v.GetInt("mail.buffer_size")),
	}

	gm.startWorkers(l)

	return gm
}

func (m *goMail) dialer() *gomail.Dialer {
	d := gomail.NewDialer(m.host, m.port, m.username, m.password)

	if m.insecure {
		d.TLSConfig = &tls.Config{
			InsecureSkipVerify: true,
		}
	}

	return d
}

func (m *goMail) mailWorker(l *logrus.Logger) {
	d := m.dialer()
	var sender gomail.SendCloser

	for {
		select {
		case m, ok := <-m.msgc:
			if !ok {
				return
			}
			if sender == nil {
				if s, err := d.Dial(); err == nil {
					sender = s
				} else {
					l.Error(err)
				}
			}
			if err := gomail.Send(sender, m); err != nil {
				l.Error(err)
			}

		case <-time.After(30 * time.Second):
			if sender != nil {
				if err := sender.Close(); err != nil {
					l.Error(err)
				}
				sender = nil
			}
		}
	}
}

func (m *goMail) startWorkers(l *logrus.Logger) {
	for i := 0; i < m.workerSize; i++ {
		go m.mailWorker(l)
	}
}

const (
	moduleIDenqueue = moduleID + ".enqueue"
)

func (m *goMail) enqueue(msg *gomail.Message) *governor.Error {
	select {
	case m.msgc <- msg:
	default:
		return governor.NewError(moduleIDenqueue, "email service experiencing load", 0, http.StatusInternalServerError)
	}

	return nil
}

// Mount is a place to mount routes to satisfy the Service interface
func (m *goMail) Mount(conf governor.Config, r *echo.Group, l *logrus.Logger) error {
	l.Info("mounted mail service")
	return nil
}

// Health is a health check for the service
func (m *goMail) Health() *governor.Error {
	return nil
}

// Setup is run on service setup
func (m *goMail) Setup(conf governor.Config, l *logrus.Logger, rsetup governor.ReqSetupPost) *governor.Error {
	return nil
}

const (
	moduleIDSend = moduleID + ".Send"
)

// Send creates and enqueues a new message to be sent
func (m *goMail) Send(to, subject, body string) *governor.Error {
	msg := gomail.NewMessage()
	msg.SetHeader("From", m.fromAddress)
	msg.SetHeader("To", to)
	msg.SetHeader("Subject", subject)
	msg.SetBody("text/html", body)

	return m.enqueue(msg)
}
<|endoftext|>


<commit_before>package session

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"strings"

	"github.com/stretchr/testify/assert"
)

var (
	uuid   = "2e075a73-98d3-4980-b0c7-ba06fbd2cc36"
	expiry = time.Date(2009, 11, 10, 23, 00, 00, 00, time.UTC)
)

type testService struct{}

func (s *testService) CreateSession(string, string) (*Session, error) {
	panic("implement me")
}

func (s *testService) FindSession(id string) (*Session, error) {
	if id == uuid {
		return &Session{
			ID:     uuid,
			Expiry: expiry,
			User: User{
				ID:       "ab2dfdfc-0603-4752-ad7f-0e57256feaa8",
				Username: "foobar",
			},
		}, nil
	}
	return nil, errors.New("session not found")
}

func (s *testService) ClearSessions() (int64, error) {
	panic("implement me")
}

func testHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusTeapot)
	user := GetSessionUser(r)
	w.Write([]byte(fmt.Sprintf(`{"id":"%s","username":"%s"}`, user.ID, user.Username)))
}

func TestAuthorized(t *testing.T) {
	s := &testService{}

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Cookie", fmt.Sprintf("%s=%s", CookieName, uuid))

	w := httptest.NewRecorder()
	h := Authorized(s)(http.HandlerFunc(testHandler))

	h.ServeHTTP(w, req)

	expected := `{"id":"ab2dfdfc-0603-4752-ad7f-0e57256feaa8","username":"foobar"}`
	assert.Equal(t, http.StatusTeapot, w.Code)
	assert.Equal(t, expected, w.Body.String())
}

func TestAuthorizedNoCookie(t *testing.T) {
	s := &testService{}

	req := httptest.NewRequest(http.MethodGet, "/", nil)

	w := httptest.NewRecorder()
	h := Authorized(s)(http.HandlerFunc(testHandler))

	h.ServeHTTP(w, req)

	expected := `{"errors":[{"title":"Unauthorized","detail":"Your Cookie is not valid","status":"401"}]}`
	assert.Equal(t, http.StatusUnauthorized, w.Code)
	assert.Equal(t, expected, strings.TrimSpace(w.Body.String()))
}

func TestAuthorizedInvalidCookie(t *testing.T) {
	s := &testService{}

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Cookie", fmt.Sprintf("%s=%s", CookieName, "blablabla"))

	w := httptest.NewRecorder()
	h := Authorized(s)(http.HandlerFunc(testHandler))

	h.ServeHTTP(w, req)

	expected := `{"errors":[{"title":"Unauthorized","detail":"Your Cookie is not valid","status":"401"}]}`
	assert.Equal(t, http.StatusUnauthorized, w.Code)
	assert.Equal(t, expected, strings.TrimSpace(w.Body.String()))
}

func TestGetSessionUser(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	ctx := context.WithValue(req.Context(), CookieUserID, "id")
	ctx = context.WithValue(ctx, CookieUserUsername, "username")
	req = req.WithContext(ctx)

	user := GetSessionUser(req)
	assert.Equal(t, "id", user.ID)
	assert.Equal(t, "username", user.Username)
}
<commit_msg>Fix session/http_test.go to use request context directly<commit_after>package session

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"strings"

	"github.com/stretchr/testify/assert"
)

var (
	uuid   = "2e075a73-98d3-4980-b0c7-ba06fbd2cc36"
	expiry = time.Date(2009, 11, 10, 23, 00, 00, 00, time.UTC)
)

type testService struct{}

func (s *testService) CreateSession(string, string) (*Session, error) {
	panic("implement me")
}

func (s *testService) FindSession(id string) (*Session, error) {
	if id == uuid {
		return &Session{
			ID:     uuid,
			Expiry: expiry,
			User: User{
				ID:       "ab2dfdfc-0603-4752-ad7f-0e57256feaa8",
				Username: "foobar",
			},
		}, nil
	}
	return nil, errors.New("session not found")
}

func (s *testService) ClearSessions() (int64, error) {
	panic("implement me")
}

func testHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusTeapot)
	user := GetSessionUser(r.Context())
	w.Write([]byte(fmt.Sprintf(`{"id":"%s","username":"%s"}`, user.ID, user.Username)))
}

func TestAuthorized(t *testing.T) {
	s := &testService{}

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Cookie", fmt.Sprintf("%s=%s", CookieName, uuid))

	w := httptest.NewRecorder()
	h := Authorized(s)(http.HandlerFunc(testHandler))

	h.ServeHTTP(w, req)

	expected := `{"id":"ab2dfdfc-0603-4752-ad7f-0e57256feaa8","username":"foobar"}`
	assert.Equal(t, http.StatusTeapot, w.Code)
	assert.Equal(t, expected, w.Body.String())
}

func TestAuthorizedNoCookie(t *testing.T) {
	s := &testService{}

	req := httptest.NewRequest(http.MethodGet, "/", nil)

	w := httptest.NewRecorder()
	h := Authorized(s)(http.HandlerFunc(testHandler))

	h.ServeHTTP(w, req)

	expected := `{"errors":[{"title":"Unauthorized","detail":"Your Cookie is not valid","status":"401"}]}`
	assert.Equal(t, http.StatusUnauthorized, w.Code)
	assert.Equal(t, expected, strings.TrimSpace(w.Body.String()))
}

func TestAuthorizedInvalidCookie(t *testing.T) {
	s := &testService{}

	req := httptest.NewRequest(http.MethodGet, "/", nil)
	req.Header.Set("Cookie", fmt.Sprintf("%s=%s", CookieName, "blablabla"))

	w := httptest.NewRecorder()
	h := Authorized(s)(http.HandlerFunc(testHandler))

	h.ServeHTTP(w, req)

	expected := `{"errors":[{"title":"Unauthorized","detail":"Your Cookie is not valid","status":"401"}]}`
	assert.Equal(t, http.StatusUnauthorized, w.Code)
	assert.Equal(t, expected, strings.TrimSpace(w.Body.String()))
}

func TestGetSessionUser(t *testing.T) {
	req := httptest.NewRequest(http.MethodGet, "/", nil)
	ctx := context.WithValue(req.Context(), CookieUserID, "id")
	ctx = context.WithValue(ctx, CookieUserUsername, "username")
	req = req.WithContext(ctx)

	user := GetSessionUser(req.Context())
	assert.Equal(t, "id", user.ID)
	assert.Equal(t, "username", user.Username)
}
<|endoftext|>


<commit_before>// Copyright 2019 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 smt

import (
	"bytes"
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/google/trillian/merkle/maphasher"
	"github.com/google/trillian/storage/tree"
	"github.com/google/trillian/testonly"
	"golang.org/x/sync/errgroup"
)

const treeID = int64(0)

var (
	hasher = maphasher.Default
	b64    = testonly.MustDecodeBase64
)

func TestWriterSplit(t *testing.T) {
	ids := []tree.NodeID2{
		tree.NewNodeID2("\x01\x00\x00\x00", 32),
		tree.NewNodeID2("\x00\x00\x00\x00", 32),
		tree.NewNodeID2("\x02\x00\x00\x00", 32),
		tree.NewNodeID2("\x03\x00\x00\x00", 32),
		tree.NewNodeID2("\x02\x00\x01\x00", 32),
		tree.NewNodeID2("\x03\x00\x00\x00", 32),
	}
	// Generate some node updates based on IDs.
	upd := make([]NodeUpdate, len(ids))
	for i, id := range ids {
		upd[i] = NodeUpdate{ID: id, Hash: []byte(fmt.Sprintf("%32d", i))}
	}

	for _, tc := range []struct {
		desc  string
		split uint
		upd   []NodeUpdate
		want  [][]NodeUpdate
		err   bool
	}{
		{desc: "dup", upd: upd, err: true},
		{desc: "wrong-len", upd: []NodeUpdate{{ID: tree.NewNodeID2("ab", 10)}}, err: true},
		{desc: "ok-24", split: 24, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1]}, {upd[0]}, {upd[2]}, {upd[4]}, {upd[3]}}},
		{desc: "ok-21", split: 21, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1]}, {upd[0]}, {upd[2], upd[4]}, {upd[3]}}},
		{desc: "ok-16", split: 16, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1]}, {upd[0]}, {upd[2], upd[4]}, {upd[3]}}},
		{desc: "ok-0", split: 0, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1], upd[0], upd[2], upd[4], upd[3]}}},
	} {
		t.Run(tc.desc, func(t *testing.T) {
			upd := make([]NodeUpdate, len(tc.upd))
			copy(upd, tc.upd) // Avoid shuffling effects.

			w := NewWriter(treeID, hasher, 32, tc.split)
			shards, err := w.Split(upd)
			if !reflect.DeepEqual(shards, tc.want) {
				t.Error("shards mismatch")
			}
			if got, want := err != nil, tc.err; got != want {
				t.Errorf("got err: %v, want %v", err, want)
			}
		})
	}
}

func TestWriterWrite(t *testing.T) {
	ctx := context.Background()
	upd := []NodeUpdate{genUpd("key1", "value1"), genUpd("key2", "value2"), genUpd("key3", "value3")}
	for _, tc := range []struct {
		desc     string
		split    uint
		acc      *testAccessor
		upd      []NodeUpdate
		wantRoot []byte
		wantErr  string
	}{
		// Taken from SparseMerkleTreeWriter tests.
		{
			desc:     "single-leaf",
			upd:      []NodeUpdate{upd[0]},
			wantRoot: b64("PPI818D5CiUQQMZulH58LikjxeOFWw2FbnGM0AdVHWA="),
		},
		{
			desc:     "multi-leaf",
			upd:      []NodeUpdate{upd[0], upd[1], upd[2]},
			wantRoot: b64("Ms8A+VeDImofprfgq7Hoqh9cw+YrD/P/qibTmCm5JvQ="),
		},

		{desc: "empty", wantErr: "nothing to write"},
		{desc: "unaligned", upd: []NodeUpdate{{ID: tree.NewNodeID2("ab", 10)}}, wantErr: "unexpected depth"},
		{desc: "dup", upd: []NodeUpdate{upd[0], upd[0]}, wantErr: "duplicate ID"},
		{desc: "2-shards", split: 128, upd: []NodeUpdate{upd[0], upd[1]}, wantErr: "writing across"},
		{desc: "get-err", acc: &testAccessor{get: errors.New("nope")}, upd: []NodeUpdate{upd[0]}, wantErr: "nope"},
		{desc: "set-err", acc: &testAccessor{set: errors.New("nope")}, upd: []NodeUpdate{upd[0]}, wantErr: "nope"},
	} {
		t.Run(tc.desc, func(t *testing.T) {
			w := NewWriter(treeID, hasher, 256, tc.split)
			acc := tc.acc
			if acc == nil {
				acc = &testAccessor{}
			}
			rootUpd, err := w.Write(ctx, tc.upd, acc)
			gotErr := ""
			if err != nil {
				gotErr = err.Error()
			}
			if got, want := gotErr, tc.wantErr; !strings.Contains(got, want) {
				t.Errorf("Write: want err containing %q, got %v", want, err)
			}
			if got, want := rootUpd.Hash, tc.wantRoot; !bytes.Equal(got, want) {
				t.Errorf("Write: got root %x, want %x", got, want)
			}
		})
	}
}

func TestWriterBigBatch(t *testing.T) {
	testWriterBigBatch(t)
}

func BenchmarkWriterBigBatch(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testWriterBigBatch(b)
	}
}

func testWriterBigBatch(t testing.TB) {
	if testing.Short() {
		t.Skip("BigBatch test is not short")
	}
	ctx := context.Background()

	const batchSize = 1024
	const numBatches = 4
	upd := make([]NodeUpdate, 0, batchSize*numBatches)
	for x := 0; x < numBatches; x++ {
		for y := 0; y < batchSize; y++ {
			u := genUpd(fmt.Sprintf("key-%d-%d", x, y), fmt.Sprintf("value-%d-%d", x, y))
			upd = append(upd, u)
		}
	}

	w := NewWriter(treeID, hasher, 256, 8)
	rootUpd := update(ctx, t, w, &testAccessor{}, upd)

	// Calculated using Python code from the original Revocation Transparency
	// doc: https://www.links.org/files/RevocationTransparency.pdf.
	want := b64("Av30xkERsepT6F/AgbZX3sp91TUmV1TKaXE6QPFfUZA=")
	if got := rootUpd.Hash; !bytes.Equal(got, want) {
		t.Errorf("root mismatch: got %x, want %x", got, want)
	}
}

func TestWriterBigBatchMultipleWrites(t *testing.T) {
	if testing.Short() {
		t.Skip("BigBatch test is not short")
	}
	ctx := context.Background()

	const batchSize = 1024
	const numBatches = 4
	roots := [numBatches][]byte{
		b64("7R5uvGy5MJ2Y8xrQr4/mnn3aPw39vYscghmg9KBJaKc="),
		b64("VTrPStz/chupeOjzAYFIHGfhiMT8yN+v589jxWZO1F0="),
		b64("nRvRV/NfC06rXGI5cKeTieyyp/69bHoMcVDs0AtZzus="),
		b64("Av30xkERsepT6F/AgbZX3sp91TUmV1TKaXE6QPFfUZA="),
	}

	w := NewWriter(treeID, hasher, 256, 8)
	acc := &testAccessor{h: make(map[tree.NodeID2][]byte), save: true}

	for i := 0; i < numBatches; i++ {
		upd := make([]NodeUpdate, 0, batchSize)
		for j := 0; j < batchSize; j++ {
			u := genUpd(fmt.Sprintf("key-%d-%d", i, j), fmt.Sprintf("value-%d-%d", i, j))
			upd = append(upd, u)
		}
		rootUpd := update(ctx, t, w, acc, upd)
		if got, want := rootUpd.Hash, roots[i]; !bytes.Equal(got, want) {
			t.Errorf("%d: root mismatch: got %x, want %x", i, got, want)
		}
	}
}

func update(ctx context.Context, t testing.TB, w *Writer, acc NodeBatchAccessor, upd []NodeUpdate) NodeUpdate {
	shards, err := w.Split(upd)
	if err != nil {
		t.Fatalf("Split: %v", err)
	}

	var mu sync.Mutex
	splitUpd := make([]NodeUpdate, 0, 256)

	eg, _ := errgroup.WithContext(ctx)
	for _, upd := range shards {
		upd := upd
		eg.Go(func() error {
			rootUpd, err := w.Write(ctx, upd, acc)
			if err != nil {
				return err
			}
			mu.Lock()
			defer mu.Unlock()
			splitUpd = append(splitUpd, rootUpd)
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		t.Fatalf("Wait: %v", err)
	}

	rootUpd, err := w.Write(ctx, splitUpd, acc)
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	return rootUpd
}

// genUpd returns a NodeUpdate for the given key and value. The returned node
// ID is 256-bit map key based on SHA256 of the given key string.
func genUpd(key, value string) NodeUpdate {
	key256 := sha256.Sum256([]byte(key))
	hash := hasher.HashLeaf(treeID, key256[:], []byte(value))
	return NodeUpdate{ID: tree.NewNodeID2(string(key256[:]), 256), Hash: hash}
}

// testAccessor implements NodeBatchAccessor for testing purposes.
type testAccessor struct {
	mu   sync.RWMutex // Guards the h map.
	h    map[tree.NodeID2][]byte
	save bool  // Persist node updates in this accessor.
	get  error // The error returned by Get.
	set  error // The error returned by Set.
}

func (t *testAccessor) Get(ctx context.Context, ids []tree.NodeID2) (map[tree.NodeID2][]byte, error) {
	if err := t.get; err != nil {
		return nil, err
	} else if !t.save {
		return nil, nil
	}
	t.mu.RLock()
	defer t.mu.RUnlock()
	h := make(map[tree.NodeID2][]byte, len(ids))
	for _, id := range ids {
		if hash, ok := t.h[id]; ok {
			h[id] = hash
		}
	}
	return h, nil
}

func (t *testAccessor) Set(ctx context.Context, upd []NodeUpdate) error {
	if err := t.set; err != nil {
		return err
	} else if !t.save {
		return nil
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, u := range upd {
		t.h[u.ID] = u.Hash
	}
	return nil
}
<commit_msg>Clarify error message in tests (#1909)<commit_after>// Copyright 2019 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 smt

import (
	"bytes"
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"reflect"
	"strings"
	"sync"
	"testing"

	"github.com/google/trillian/merkle/maphasher"
	"github.com/google/trillian/storage/tree"
	"github.com/google/trillian/testonly"
	"golang.org/x/sync/errgroup"
)

const treeID = int64(0)

var (
	hasher = maphasher.Default
	b64    = testonly.MustDecodeBase64
)

func TestWriterSplit(t *testing.T) {
	ids := []tree.NodeID2{
		tree.NewNodeID2("\x01\x00\x00\x00", 32),
		tree.NewNodeID2("\x00\x00\x00\x00", 32),
		tree.NewNodeID2("\x02\x00\x00\x00", 32),
		tree.NewNodeID2("\x03\x00\x00\x00", 32),
		tree.NewNodeID2("\x02\x00\x01\x00", 32),
		tree.NewNodeID2("\x03\x00\x00\x00", 32),
	}
	// Generate some node updates based on IDs.
	upd := make([]NodeUpdate, len(ids))
	for i, id := range ids {
		upd[i] = NodeUpdate{ID: id, Hash: []byte(fmt.Sprintf("%32d", i))}
	}

	for _, tc := range []struct {
		desc  string
		split uint
		upd   []NodeUpdate
		want  [][]NodeUpdate
		err   bool
	}{
		{desc: "dup", upd: upd, err: true},
		{desc: "wrong-len", upd: []NodeUpdate{{ID: tree.NewNodeID2("ab", 10)}}, err: true},
		{desc: "ok-24", split: 24, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1]}, {upd[0]}, {upd[2]}, {upd[4]}, {upd[3]}}},
		{desc: "ok-21", split: 21, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1]}, {upd[0]}, {upd[2], upd[4]}, {upd[3]}}},
		{desc: "ok-16", split: 16, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1]}, {upd[0]}, {upd[2], upd[4]}, {upd[3]}}},
		{desc: "ok-0", split: 0, upd: upd[:5],
			want: [][]NodeUpdate{{upd[1], upd[0], upd[2], upd[4], upd[3]}}},
	} {
		t.Run(tc.desc, func(t *testing.T) {
			upd := make([]NodeUpdate, len(tc.upd))
			copy(upd, tc.upd) // Avoid shuffling effects.

			w := NewWriter(treeID, hasher, 32, tc.split)
			shards, err := w.Split(upd)
			if !reflect.DeepEqual(shards, tc.want) {
				t.Error("shards mismatch")
			}
			if got, want := err != nil, tc.err; got != want {
				t.Errorf("got err: %v, want %v", err, want)
			}
		})
	}
}

func TestWriterWrite(t *testing.T) {
	ctx := context.Background()
	upd := []NodeUpdate{genUpd("key1", "value1"), genUpd("key2", "value2"), genUpd("key3", "value3")}
	for _, tc := range []struct {
		desc     string
		split    uint
		acc      *testAccessor
		upd      []NodeUpdate
		wantRoot []byte
		wantErr  string
	}{
		// Taken from SparseMerkleTreeWriter tests.
		{
			desc:     "single-leaf",
			upd:      []NodeUpdate{upd[0]},
			wantRoot: b64("PPI818D5CiUQQMZulH58LikjxeOFWw2FbnGM0AdVHWA="),
		},
		{
			desc:     "multi-leaf",
			upd:      []NodeUpdate{upd[0], upd[1], upd[2]},
			wantRoot: b64("Ms8A+VeDImofprfgq7Hoqh9cw+YrD/P/qibTmCm5JvQ="),
		},

		{desc: "empty", wantErr: "nothing to write"},
		{desc: "unaligned", upd: []NodeUpdate{{ID: tree.NewNodeID2("ab", 10)}}, wantErr: "unexpected depth"},
		{desc: "dup", upd: []NodeUpdate{upd[0], upd[0]}, wantErr: "duplicate ID"},
		{desc: "2-shards", split: 128, upd: []NodeUpdate{upd[0], upd[1]}, wantErr: "writing across"},
		{desc: "get-err", acc: &testAccessor{get: errors.New("fail")}, upd: []NodeUpdate{upd[0]}, wantErr: "fail"},
		{desc: "set-err", acc: &testAccessor{set: errors.New("fail")}, upd: []NodeUpdate{upd[0]}, wantErr: "fail"},
	} {
		t.Run(tc.desc, func(t *testing.T) {
			w := NewWriter(treeID, hasher, 256, tc.split)
			acc := tc.acc
			if acc == nil {
				acc = &testAccessor{}
			}
			rootUpd, err := w.Write(ctx, tc.upd, acc)
			gotErr := ""
			if err != nil {
				gotErr = err.Error()
			}
			if got, want := gotErr, tc.wantErr; !strings.Contains(got, want) {
				t.Errorf("Write: want err containing %q, got %v", want, err)
			}
			if got, want := rootUpd.Hash, tc.wantRoot; !bytes.Equal(got, want) {
				t.Errorf("Write: got root %x, want %x", got, want)
			}
		})
	}
}

func TestWriterBigBatch(t *testing.T) {
	testWriterBigBatch(t)
}

func BenchmarkWriterBigBatch(b *testing.B) {
	for i := 0; i < b.N; i++ {
		testWriterBigBatch(b)
	}
}

func testWriterBigBatch(t testing.TB) {
	if testing.Short() {
		t.Skip("BigBatch test is not short")
	}
	ctx := context.Background()

	const batchSize = 1024
	const numBatches = 4
	upd := make([]NodeUpdate, 0, batchSize*numBatches)
	for x := 0; x < numBatches; x++ {
		for y := 0; y < batchSize; y++ {
			u := genUpd(fmt.Sprintf("key-%d-%d", x, y), fmt.Sprintf("value-%d-%d", x, y))
			upd = append(upd, u)
		}
	}

	w := NewWriter(treeID, hasher, 256, 8)
	rootUpd := update(ctx, t, w, &testAccessor{}, upd)

	// Calculated using Python code from the original Revocation Transparency
	// doc: https://www.links.org/files/RevocationTransparency.pdf.
	want := b64("Av30xkERsepT6F/AgbZX3sp91TUmV1TKaXE6QPFfUZA=")
	if got := rootUpd.Hash; !bytes.Equal(got, want) {
		t.Errorf("root mismatch: got %x, want %x", got, want)
	}
}

func TestWriterBigBatchMultipleWrites(t *testing.T) {
	if testing.Short() {
		t.Skip("BigBatch test is not short")
	}
	ctx := context.Background()

	const batchSize = 1024
	const numBatches = 4
	roots := [numBatches][]byte{
		b64("7R5uvGy5MJ2Y8xrQr4/mnn3aPw39vYscghmg9KBJaKc="),
		b64("VTrPStz/chupeOjzAYFIHGfhiMT8yN+v589jxWZO1F0="),
		b64("nRvRV/NfC06rXGI5cKeTieyyp/69bHoMcVDs0AtZzus="),
		b64("Av30xkERsepT6F/AgbZX3sp91TUmV1TKaXE6QPFfUZA="),
	}

	w := NewWriter(treeID, hasher, 256, 8)
	acc := &testAccessor{h: make(map[tree.NodeID2][]byte), save: true}

	for i := 0; i < numBatches; i++ {
		upd := make([]NodeUpdate, 0, batchSize)
		for j := 0; j < batchSize; j++ {
			u := genUpd(fmt.Sprintf("key-%d-%d", i, j), fmt.Sprintf("value-%d-%d", i, j))
			upd = append(upd, u)
		}
		rootUpd := update(ctx, t, w, acc, upd)
		if got, want := rootUpd.Hash, roots[i]; !bytes.Equal(got, want) {
			t.Errorf("%d: root mismatch: got %x, want %x", i, got, want)
		}
	}
}

func update(ctx context.Context, t testing.TB, w *Writer, acc NodeBatchAccessor, upd []NodeUpdate) NodeUpdate {
	shards, err := w.Split(upd)
	if err != nil {
		t.Fatalf("Split: %v", err)
	}

	var mu sync.Mutex
	splitUpd := make([]NodeUpdate, 0, 256)

	eg, _ := errgroup.WithContext(ctx)
	for _, upd := range shards {
		upd := upd
		eg.Go(func() error {
			rootUpd, err := w.Write(ctx, upd, acc)
			if err != nil {
				return err
			}
			mu.Lock()
			defer mu.Unlock()
			splitUpd = append(splitUpd, rootUpd)
			return nil
		})
	}
	if err := eg.Wait(); err != nil {
		t.Fatalf("Wait: %v", err)
	}

	rootUpd, err := w.Write(ctx, splitUpd, acc)
	if err != nil {
		t.Fatalf("Write: %v", err)
	}
	return rootUpd
}

// genUpd returns a NodeUpdate for the given key and value. The returned node
// ID is 256-bit map key based on SHA256 of the given key string.
func genUpd(key, value string) NodeUpdate {
	key256 := sha256.Sum256([]byte(key))
	hash := hasher.HashLeaf(treeID, key256[:], []byte(value))
	return NodeUpdate{ID: tree.NewNodeID2(string(key256[:]), 256), Hash: hash}
}

// testAccessor implements NodeBatchAccessor for testing purposes.
type testAccessor struct {
	mu   sync.RWMutex // Guards the h map.
	h    map[tree.NodeID2][]byte
	save bool  // Persist node updates in this accessor.
	get  error // The error returned by Get.
	set  error // The error returned by Set.
}

func (t *testAccessor) Get(ctx context.Context, ids []tree.NodeID2) (map[tree.NodeID2][]byte, error) {
	if err := t.get; err != nil {
		return nil, err
	} else if !t.save {
		return nil, nil
	}
	t.mu.RLock()
	defer t.mu.RUnlock()
	h := make(map[tree.NodeID2][]byte, len(ids))
	for _, id := range ids {
		if hash, ok := t.h[id]; ok {
			h[id] = hash
		}
	}
	return h, nil
}

func (t *testAccessor) Set(ctx context.Context, upd []NodeUpdate) error {
	if err := t.set; err != nil {
		return err
	} else if !t.save {
		return nil
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	for _, u := range upd {
		t.h[u.ID] = u.Hash
	}
	return nil
}
<|endoftext|>