<commit_before>// Copyright 2016-2018 Authors of Cilium
//
// 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 endpointmanager

import (
	"context"
	"fmt"
	"sync"

	"github.com/cilium/cilium/pkg/endpoint"
	endpointid "github.com/cilium/cilium/pkg/endpoint/id"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/metrics"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/policy"

	"github.com/prometheus/client_golang/prometheus"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "endpoint-manager")

	// mutex protects endpoints and endpointsAux
	mutex lock.RWMutex

	// endpoints is the global list of endpoints indexed by ID. mutex must
	// be held to read and write.
	endpoints    = map[uint16]*endpoint.Endpoint{}
	endpointsAux = map[string]*endpoint.Endpoint{}
)

func init() {
	// EndpointCount is a function used to collect this metric. We cannot
	// increment/decrement a gauge since we invoke Remove gratuitiously and that
	// would result in negative counts.
	// It must be thread-safe.
	metrics.EndpointCount = prometheus.NewGaugeFunc(prometheus.GaugeOpts{
		Namespace: metrics.Namespace,
		Name:      "endpoint_count",
		Help:      "Number of endpoints managed by this agent",
	},
		func() float64 { return float64(len(GetEndpoints())) },
	)
	metrics.MustRegister(metrics.EndpointCount)
}

// Insert inserts the endpoint into the global maps.
func Insert(ep *endpoint.Endpoint) {
	// No need to check liveness as an endpoint can only be deleted via the
	// API after it has been inserted into the manager.
	ep.UnconditionalRLock()
	mutex.Lock()

	endpoints[ep.ID] = ep
	updateReferences(ep)

	mutex.Unlock()
	ep.RUnlock()

	ep.RunK8sCiliumEndpointSync()
}

// Lookup looks up the endpoint by prefix id
func Lookup(id string) (*endpoint.Endpoint, error) {
	mutex.RLock()
	defer mutex.RUnlock()

	prefix, eid, err := endpointid.ParseID(id)
	if err != nil {
		return nil, err
	}

	switch prefix {
	case endpointid.CiliumLocalIdPrefix:
		n, err := endpointid.ParseCiliumID(id)
		if err != nil {
			return nil, err
		}
		return lookupCiliumID(uint16(n)), nil

	case endpointid.CiliumGlobalIdPrefix:
		return nil, fmt.Errorf("Unsupported id format for now")

	case endpointid.ContainerIdPrefix:
		return lookupDockerID(eid), nil

	case endpointid.DockerEndpointPrefix:
		return lookupDockerEndpoint(eid), nil

	case endpointid.ContainerNamePrefix:
		return lookupDockerContainerName(eid), nil

	case endpointid.PodNamePrefix:
		return lookupPodNameLocked(eid), nil

	case endpointid.IPv4Prefix:
		return lookupIPv4(eid), nil

	default:
		return nil, fmt.Errorf("Unknown endpoint prefix %s", prefix)
	}
}

// LookupCiliumID looks up endpoint by endpoint ID
func LookupCiliumID(id uint16) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupCiliumID(id)
	mutex.RUnlock()
	return ep
}

// LookupDockerID looks up endpoint by Docker ID
func LookupDockerID(id string) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupDockerID(id)
	mutex.RUnlock()
	return ep
}

// LookupIPv4 looks up endpoint by IPv4 address
func LookupIPv4(ipv4 string) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupIPv4(ipv4)
	mutex.RUnlock()
	return ep
}

// LookupPodName looks up endpoint by namespace + pod name
func LookupPodName(name string) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupPodNameLocked(name)
	mutex.RUnlock()
	return ep
}

// UpdateReferences makes an endpoint available by all possible reference
// fields as available for this endpoint (containerID, IPv4 address, ...)
// Must be called with ep.Mutex.RLock held.
func UpdateReferences(ep *endpoint.Endpoint) {
	mutex.Lock()
	defer mutex.Unlock()
	updateReferences(ep)
}

// Remove removes the endpoint from the global maps.
// Must be called with ep.Mutex.RLock held.
func Remove(ep *endpoint.Endpoint) {
	mutex.Lock()
	defer mutex.Unlock()
	delete(endpoints, ep.ID)

	if ep.ContainerID != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.ContainerIdPrefix, ep.ContainerID))
	}

	if ep.DockerEndpointID != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.DockerEndpointPrefix, ep.DockerEndpointID))
	}

	if ep.IPv4.String() != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.IPv4Prefix, ep.IPv4.String()))
	}

	if ep.ContainerName != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.ContainerNamePrefix, ep.ContainerName))
	}

	if podName := ep.GetK8sNamespaceAndPodNameLocked(); podName != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.PodNamePrefix, podName))
	}
}

// RemoveAll removes all endpoints from the global maps.
func RemoveAll() {
	mutex.Lock()
	defer mutex.Unlock()
	endpoints = map[uint16]*endpoint.Endpoint{}
	endpointsAux = map[string]*endpoint.Endpoint{}
}

// lookupCiliumID looks up endpoint by endpoint ID
func lookupCiliumID(id uint16) *endpoint.Endpoint {
	if ep, ok := endpoints[id]; ok {
		return ep
	}
	return nil
}

func lookupDockerEndpoint(id string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.DockerEndpointPrefix, id)]; ok {
		return ep
	}
	return nil
}

func lookupPodNameLocked(name string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.PodNamePrefix, name)]; ok {
		return ep
	}
	return nil
}

func lookupDockerContainerName(name string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.ContainerNamePrefix, name)]; ok {
		return ep
	}
	return nil
}

func lookupIPv4(ipv4 string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.IPv4Prefix, ipv4)]; ok {
		return ep
	}
	return nil
}

func lookupDockerID(id string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.ContainerIdPrefix, id)]; ok {
		return ep
	}
	return nil
}

func linkContainerID(ep *endpoint.Endpoint) {
	endpointsAux[endpointid.NewID(endpointid.ContainerIdPrefix, ep.ContainerID)] = ep
}

// UpdateReferences updates the mappings of various values to their corresponding
// endpoints, such as ContainerID, Docker Container Name, Pod Name, etc.
func updateReferences(ep *endpoint.Endpoint) {
	if ep.ContainerID != "" {
		linkContainerID(ep)
	}

	if ep.DockerEndpointID != "" {
		endpointsAux[endpointid.NewID(endpointid.DockerEndpointPrefix, ep.DockerEndpointID)] = ep
	}

	if ep.IPv4.String() != "" {
		endpointsAux[endpointid.NewID(endpointid.IPv4Prefix, ep.IPv4.String())] = ep
	}

	if ep.ContainerName != "" {
		endpointsAux[endpointid.NewID(endpointid.ContainerNamePrefix, ep.ContainerName)] = ep
	}

	if podName := ep.GetK8sNamespaceAndPodNameLocked(); podName != "" {
		endpointsAux[endpointid.NewID(endpointid.PodNamePrefix, podName)] = ep
	}
}

// RegenerateAllEndpoints calls a SetStateLocked for each endpoint and
// regenerates if state transaction is valid. During this process, the endpoint
// list is locked and cannot be modified.
// Returns a waiting group that can be used to know when all the endpoints are
// regenerated.
func RegenerateAllEndpoints(owner endpoint.Owner, regenContext *endpoint.RegenerationContext) *sync.WaitGroup {
	var wg sync.WaitGroup

	eps := GetEndpoints()
	wg.Add(len(eps))

	log.Infof("regenerating all endpoints due to %s", regenContext.Reason)
	for _, ep := range eps {
		go func(ep *endpoint.Endpoint, wg *sync.WaitGroup) {
			if err := ep.LockAlive(); err != nil {
				log.WithError(err).Warn("Error regenerating endpoint for event %s", regenContext.Reason)
				ep.LogStatus(endpoint.Policy, endpoint.Failure, "Error while handling policy updates for endpoint: "+err.Error())
			} else {
				regen := ep.SetStateLocked(endpoint.StateWaitingToRegenerate, fmt.Sprintf("Triggering endpoint regeneration due to %s", regenContext.Reason))
				ep.Unlock()
				if regen {
					// Regenerate logs status according to the build success/failure
					<-ep.Regenerate(owner, regenContext)
				}
			}
			wg.Done()
		}(ep, &wg)
	}

	return &wg
}

// HasGlobalCT returns true if the endpoints have a global CT, false otherwise.
func HasGlobalCT() bool {
	eps := GetEndpoints()
	for _, e := range eps {
		if !e.Options.IsEnabled(option.ConntrackLocal) {
			return true
		}
	}
	return false
}

// GetEndpoints returns a slice of all endpoints present in endpoint manager.
func GetEndpoints() []*endpoint.Endpoint {
	mutex.RLock()
	eps := make([]*endpoint.Endpoint, 0, len(endpoints))
	for _, ep := range endpoints {
		eps = append(eps, ep)
	}
	mutex.RUnlock()
	return eps
}

// AddEndpoint takes the prepared endpoint object and starts managing it.
func AddEndpoint(owner endpoint.Owner, ep *endpoint.Endpoint, reason string) error {
	alwaysEnforce := policy.GetPolicyEnabled() == option.AlwaysEnforce
	ep.SetIngressPolicyEnabled(alwaysEnforce)
	ep.SetEgressPolicyEnabled(alwaysEnforce)

	if err := ep.CreateDirectory(); err != nil {
		return err
	}

	// Regenerate immediately if ready or waiting for identity
	if err := ep.LockAlive(); err != nil {
		return err
	}
	build := false
	state := ep.GetStateLocked()

	// We can only trigger regeneration of endpoints if the endpoint is in a
	// state where it can regenerate. See endpoint.SetStateLocked().
	if state == endpoint.StateReady {
		ep.SetStateLocked(endpoint.StateWaitingToRegenerate, reason)
		build = true
	}
	ep.Unlock()

	if build {
		if err := ep.RegenerateWait(owner, reason); err != nil {
			ep.RemoveDirectory()
			return err
		}
	}

	Insert(ep)
	ep.InsertEvent()

	return nil
}

// WaitForEndpointsAtPolicyRev waits for all endpoints which existed at the time
// this function is called to be at a given policy revision.
// New endpoints appearing while waiting are ignored.
func WaitForEndpointsAtPolicyRev(ctx context.Context, rev uint64) error {
	eps := GetEndpoints()
	for i := range eps {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-eps[i].WaitForPolicyRevision(ctx, rev):
			if ctx.Err() != nil {
				return ctx.Err()
			}
		}
	}
	return nil
}
<commit_msg>endpoint: Correctly remove directory on endpoint add failure<commit_after>// Copyright 2016-2018 Authors of Cilium
//
// 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 endpointmanager

import (
	"context"
	"fmt"
	"sync"

	"github.com/cilium/cilium/pkg/endpoint"
	endpointid "github.com/cilium/cilium/pkg/endpoint/id"
	"github.com/cilium/cilium/pkg/lock"
	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/metrics"
	"github.com/cilium/cilium/pkg/option"
	"github.com/cilium/cilium/pkg/policy"

	"github.com/prometheus/client_golang/prometheus"
)

var (
	log = logging.DefaultLogger.WithField(logfields.LogSubsys, "endpoint-manager")

	// mutex protects endpoints and endpointsAux
	mutex lock.RWMutex

	// endpoints is the global list of endpoints indexed by ID. mutex must
	// be held to read and write.
	endpoints    = map[uint16]*endpoint.Endpoint{}
	endpointsAux = map[string]*endpoint.Endpoint{}
)

func init() {
	// EndpointCount is a function used to collect this metric. We cannot
	// increment/decrement a gauge since we invoke Remove gratuitiously and that
	// would result in negative counts.
	// It must be thread-safe.
	metrics.EndpointCount = prometheus.NewGaugeFunc(prometheus.GaugeOpts{
		Namespace: metrics.Namespace,
		Name:      "endpoint_count",
		Help:      "Number of endpoints managed by this agent",
	},
		func() float64 { return float64(len(GetEndpoints())) },
	)
	metrics.MustRegister(metrics.EndpointCount)
}

// Insert inserts the endpoint into the global maps.
func Insert(ep *endpoint.Endpoint) {
	// No need to check liveness as an endpoint can only be deleted via the
	// API after it has been inserted into the manager.
	ep.UnconditionalRLock()
	mutex.Lock()

	endpoints[ep.ID] = ep
	updateReferences(ep)

	mutex.Unlock()
	ep.RUnlock()

	ep.RunK8sCiliumEndpointSync()
}

// Lookup looks up the endpoint by prefix id
func Lookup(id string) (*endpoint.Endpoint, error) {
	mutex.RLock()
	defer mutex.RUnlock()

	prefix, eid, err := endpointid.ParseID(id)
	if err != nil {
		return nil, err
	}

	switch prefix {
	case endpointid.CiliumLocalIdPrefix:
		n, err := endpointid.ParseCiliumID(id)
		if err != nil {
			return nil, err
		}
		return lookupCiliumID(uint16(n)), nil

	case endpointid.CiliumGlobalIdPrefix:
		return nil, fmt.Errorf("Unsupported id format for now")

	case endpointid.ContainerIdPrefix:
		return lookupDockerID(eid), nil

	case endpointid.DockerEndpointPrefix:
		return lookupDockerEndpoint(eid), nil

	case endpointid.ContainerNamePrefix:
		return lookupDockerContainerName(eid), nil

	case endpointid.PodNamePrefix:
		return lookupPodNameLocked(eid), nil

	case endpointid.IPv4Prefix:
		return lookupIPv4(eid), nil

	default:
		return nil, fmt.Errorf("Unknown endpoint prefix %s", prefix)
	}
}

// LookupCiliumID looks up endpoint by endpoint ID
func LookupCiliumID(id uint16) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupCiliumID(id)
	mutex.RUnlock()
	return ep
}

// LookupDockerID looks up endpoint by Docker ID
func LookupDockerID(id string) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupDockerID(id)
	mutex.RUnlock()
	return ep
}

// LookupIPv4 looks up endpoint by IPv4 address
func LookupIPv4(ipv4 string) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupIPv4(ipv4)
	mutex.RUnlock()
	return ep
}

// LookupPodName looks up endpoint by namespace + pod name
func LookupPodName(name string) *endpoint.Endpoint {
	mutex.RLock()
	ep := lookupPodNameLocked(name)
	mutex.RUnlock()
	return ep
}

// UpdateReferences makes an endpoint available by all possible reference
// fields as available for this endpoint (containerID, IPv4 address, ...)
// Must be called with ep.Mutex.RLock held.
func UpdateReferences(ep *endpoint.Endpoint) {
	mutex.Lock()
	defer mutex.Unlock()
	updateReferences(ep)
}

// Remove removes the endpoint from the global maps.
// Must be called with ep.Mutex.RLock held.
func Remove(ep *endpoint.Endpoint) {
	mutex.Lock()
	defer mutex.Unlock()
	delete(endpoints, ep.ID)

	if ep.ContainerID != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.ContainerIdPrefix, ep.ContainerID))
	}

	if ep.DockerEndpointID != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.DockerEndpointPrefix, ep.DockerEndpointID))
	}

	if ep.IPv4.String() != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.IPv4Prefix, ep.IPv4.String()))
	}

	if ep.ContainerName != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.ContainerNamePrefix, ep.ContainerName))
	}

	if podName := ep.GetK8sNamespaceAndPodNameLocked(); podName != "" {
		delete(endpointsAux, endpointid.NewID(endpointid.PodNamePrefix, podName))
	}
}

// RemoveAll removes all endpoints from the global maps.
func RemoveAll() {
	mutex.Lock()
	defer mutex.Unlock()
	endpoints = map[uint16]*endpoint.Endpoint{}
	endpointsAux = map[string]*endpoint.Endpoint{}
}

// lookupCiliumID looks up endpoint by endpoint ID
func lookupCiliumID(id uint16) *endpoint.Endpoint {
	if ep, ok := endpoints[id]; ok {
		return ep
	}
	return nil
}

func lookupDockerEndpoint(id string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.DockerEndpointPrefix, id)]; ok {
		return ep
	}
	return nil
}

func lookupPodNameLocked(name string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.PodNamePrefix, name)]; ok {
		return ep
	}
	return nil
}

func lookupDockerContainerName(name string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.ContainerNamePrefix, name)]; ok {
		return ep
	}
	return nil
}

func lookupIPv4(ipv4 string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.IPv4Prefix, ipv4)]; ok {
		return ep
	}
	return nil
}

func lookupDockerID(id string) *endpoint.Endpoint {
	if ep, ok := endpointsAux[endpointid.NewID(endpointid.ContainerIdPrefix, id)]; ok {
		return ep
	}
	return nil
}

func linkContainerID(ep *endpoint.Endpoint) {
	endpointsAux[endpointid.NewID(endpointid.ContainerIdPrefix, ep.ContainerID)] = ep
}

// UpdateReferences updates the mappings of various values to their corresponding
// endpoints, such as ContainerID, Docker Container Name, Pod Name, etc.
func updateReferences(ep *endpoint.Endpoint) {
	if ep.ContainerID != "" {
		linkContainerID(ep)
	}

	if ep.DockerEndpointID != "" {
		endpointsAux[endpointid.NewID(endpointid.DockerEndpointPrefix, ep.DockerEndpointID)] = ep
	}

	if ep.IPv4.String() != "" {
		endpointsAux[endpointid.NewID(endpointid.IPv4Prefix, ep.IPv4.String())] = ep
	}

	if ep.ContainerName != "" {
		endpointsAux[endpointid.NewID(endpointid.ContainerNamePrefix, ep.ContainerName)] = ep
	}

	if podName := ep.GetK8sNamespaceAndPodNameLocked(); podName != "" {
		endpointsAux[endpointid.NewID(endpointid.PodNamePrefix, podName)] = ep
	}
}

// RegenerateAllEndpoints calls a SetStateLocked for each endpoint and
// regenerates if state transaction is valid. During this process, the endpoint
// list is locked and cannot be modified.
// Returns a waiting group that can be used to know when all the endpoints are
// regenerated.
func RegenerateAllEndpoints(owner endpoint.Owner, regenContext *endpoint.RegenerationContext) *sync.WaitGroup {
	var wg sync.WaitGroup

	eps := GetEndpoints()
	wg.Add(len(eps))

	log.Infof("regenerating all endpoints due to %s", regenContext.Reason)
	for _, ep := range eps {
		go func(ep *endpoint.Endpoint, wg *sync.WaitGroup) {
			if err := ep.LockAlive(); err != nil {
				log.WithError(err).Warn("Error regenerating endpoint for event %s", regenContext.Reason)
				ep.LogStatus(endpoint.Policy, endpoint.Failure, "Error while handling policy updates for endpoint: "+err.Error())
			} else {
				regen := ep.SetStateLocked(endpoint.StateWaitingToRegenerate, fmt.Sprintf("Triggering endpoint regeneration due to %s", regenContext.Reason))
				ep.Unlock()
				if regen {
					// Regenerate logs status according to the build success/failure
					<-ep.Regenerate(owner, regenContext)
				}
			}
			wg.Done()
		}(ep, &wg)
	}

	return &wg
}

// HasGlobalCT returns true if the endpoints have a global CT, false otherwise.
func HasGlobalCT() bool {
	eps := GetEndpoints()
	for _, e := range eps {
		if !e.Options.IsEnabled(option.ConntrackLocal) {
			return true
		}
	}
	return false
}

// GetEndpoints returns a slice of all endpoints present in endpoint manager.
func GetEndpoints() []*endpoint.Endpoint {
	mutex.RLock()
	eps := make([]*endpoint.Endpoint, 0, len(endpoints))
	for _, ep := range endpoints {
		eps = append(eps, ep)
	}
	mutex.RUnlock()
	return eps
}

// AddEndpoint takes the prepared endpoint object and starts managing it.
func AddEndpoint(owner endpoint.Owner, ep *endpoint.Endpoint, reason string) (err error) {
	alwaysEnforce := policy.GetPolicyEnabled() == option.AlwaysEnforce
	ep.SetIngressPolicyEnabled(alwaysEnforce)
	ep.SetEgressPolicyEnabled(alwaysEnforce)

	if err := ep.CreateDirectory(); err != nil {
		return err
	}
	defer func() {
		if err != nil {
			ep.RemoveDirectory()
		}
	}()

	// Regenerate immediately if ready or waiting for identity
	if err := ep.LockAlive(); err != nil {
		return err
	}
	build := false
	state := ep.GetStateLocked()

	// We can only trigger regeneration of endpoints if the endpoint is in a
	// state where it can regenerate. See endpoint.SetStateLocked().
	if state == endpoint.StateReady {
		ep.SetStateLocked(endpoint.StateWaitingToRegenerate, reason)
		build = true
	}
	ep.Unlock()

	if build {
		if err := ep.RegenerateWait(owner, reason); err != nil {
			return err
		}
	}

	Insert(ep)
	ep.InsertEvent()

	return nil
}

// WaitForEndpointsAtPolicyRev waits for all endpoints which existed at the time
// this function is called to be at a given policy revision.
// New endpoints appearing while waiting are ignored.
func WaitForEndpointsAtPolicyRev(ctx context.Context, rev uint64) error {
	eps := GetEndpoints()
	for i := range eps {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-eps[i].WaitForPolicyRevision(ctx, rev):
			if ctx.Err() != nil {
				return ctx.Err()
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2020 The Knative Authors

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

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

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

package server

import (
	"context"
	"fmt"
	"net"

	envoyv2 "github.com/envoyproxy/go-control-plane/envoy/api/v2"
	core "github.com/envoyproxy/go-control-plane/envoy/api/v2/core"
	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v2"
	cache "github.com/envoyproxy/go-control-plane/pkg/cache/v2"
	xds "github.com/envoyproxy/go-control-plane/pkg/server/v2"
	"google.golang.org/grpc"
)

const (
	grpcMaxConcurrentStreams = 1000000
)

type XdsServer struct {
	managementPort uint
	ctx            context.Context
	server         xds.Server
	snapshotCache  cache.SnapshotCache
}

// hasher returns node ID as an ID
type hasher struct {
}

func (h hasher) ID(node *core.Node) string {
	if node == nil {
		return "unknown"
	}
	return node.Id
}

func NewXdsServer(managementPort uint, callbacks xds.Callbacks) *XdsServer {
	ctx := context.Background()
	snapshotCache := cache.NewSnapshotCache(true, hasher{}, nil)
	srv := xds.NewServer(ctx, snapshotCache, callbacks)

	return &XdsServer{
		managementPort: managementPort,
		ctx:            ctx,
		server:         srv,
		snapshotCache:  snapshotCache,
	}
}

// RunManagementServer starts an xDS server at the given Port.
func (envoyXdsServer *XdsServer) RunManagementServer() error {
	port := envoyXdsServer.managementPort
	server := envoyXdsServer.server

	grpcServer := grpc.NewServer(grpc.MaxConcurrentStreams(grpcMaxConcurrentStreams))
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return fmt.Errorf("failed to listen: %w", err)
	}

	// register services
	discovery.RegisterAggregatedDiscoveryServiceServer(grpcServer, server)
	envoyv2.RegisterClusterDiscoveryServiceServer(grpcServer, server)
	envoyv2.RegisterListenerDiscoveryServiceServer(grpcServer, server)
	envoyv2.RegisterRouteDiscoveryServiceServer(grpcServer, server)

	errCh := make(chan error)
	go func() {
		if err = grpcServer.Serve(lis); err != nil {
			errCh <- err
		}
	}()

	select {
	case <-envoyXdsServer.ctx.Done():
		grpcServer.GracefulStop()
		return nil
	case err := <-errCh:
		return fmt.Errorf("failed to serve: %w", err)
	}
}

func (envoyXdsServer *XdsServer) SetSnapshot(nodeID string, snapshot cache.Snapshot) error {
	return envoyXdsServer.snapshotCache.SetSnapshot(nodeID, snapshot)
}
<commit_msg>Use IDHasher from envoy's library (#364)<commit_after>/*
Copyright 2020 The Knative Authors

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

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

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

package server

import (
	"context"
	"fmt"
	"net"

	envoyv2 "github.com/envoyproxy/go-control-plane/envoy/api/v2"
	discovery "github.com/envoyproxy/go-control-plane/envoy/service/discovery/v2"
	cache "github.com/envoyproxy/go-control-plane/pkg/cache/v2"
	xds "github.com/envoyproxy/go-control-plane/pkg/server/v2"
	"google.golang.org/grpc"
)

const (
	grpcMaxConcurrentStreams = 1000000
)

type XdsServer struct {
	managementPort uint
	ctx            context.Context
	server         xds.Server
	snapshotCache  cache.SnapshotCache
}

func NewXdsServer(managementPort uint, callbacks xds.Callbacks) *XdsServer {
	ctx := context.Background()
	snapshotCache := cache.NewSnapshotCache(true, cache.IDHash{}, nil)
	srv := xds.NewServer(ctx, snapshotCache, callbacks)

	return &XdsServer{
		managementPort: managementPort,
		ctx:            ctx,
		server:         srv,
		snapshotCache:  snapshotCache,
	}
}

// RunManagementServer starts an xDS server at the given Port.
func (envoyXdsServer *XdsServer) RunManagementServer() error {
	port := envoyXdsServer.managementPort
	server := envoyXdsServer.server

	grpcServer := grpc.NewServer(grpc.MaxConcurrentStreams(grpcMaxConcurrentStreams))
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
	if err != nil {
		return fmt.Errorf("failed to listen: %w", err)
	}

	// register services
	discovery.RegisterAggregatedDiscoveryServiceServer(grpcServer, server)
	envoyv2.RegisterClusterDiscoveryServiceServer(grpcServer, server)
	envoyv2.RegisterListenerDiscoveryServiceServer(grpcServer, server)
	envoyv2.RegisterRouteDiscoveryServiceServer(grpcServer, server)

	errCh := make(chan error)
	go func() {
		if err = grpcServer.Serve(lis); err != nil {
			errCh <- err
		}
	}()

	select {
	case <-envoyXdsServer.ctx.Done():
		grpcServer.GracefulStop()
		return nil
	case err := <-errCh:
		return fmt.Errorf("failed to serve: %w", err)
	}
}

func (envoyXdsServer *XdsServer) SetSnapshot(nodeID string, snapshot cache.Snapshot) error {
	return envoyXdsServer.snapshotCache.SetSnapshot(nodeID, snapshot)
}
<|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 registry

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/apiserver"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/cloudprovider"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

// ServiceRegistryStorage adapts a service registry into apiserver's RESTStorage model.
type ServiceRegistryStorage struct {
	registry ServiceRegistry
	cloud    cloudprovider.Interface
	machines MinionRegistry
}

// MakeServiceRegistryStorage makes a new ServiceRegistryStorage.
func MakeServiceRegistryStorage(registry ServiceRegistry, cloud cloudprovider.Interface, machines MinionRegistry) apiserver.RESTStorage {
	return &ServiceRegistryStorage{
		registry: registry,
		cloud:    cloud,
		machines: machines,
	}
}

func makeLinkVariables(service api.Service, machine string) []api.EnvVar {
	prefix := strings.ToUpper(service.ID)
	var port string
	if service.ContainerPort.Kind == util.IntstrString {
		port = service.ContainerPort.StrVal
	} else {
		port = strconv.Itoa(service.ContainerPort.IntVal)
	}
	portPrefix := prefix + "_PORT_" + strings.ToUpper(strings.Replace(port, "-", "_", -1)) + "_TCP"
	return []api.EnvVar{
		{
			Name:  prefix + "_PORT",
			Value: fmt.Sprintf("tcp://%s:%d", machine, service.Port),
		},
		{
			Name:  portPrefix,
			Value: fmt.Sprintf("tcp://%s:%d", machine, service.Port),
		},
		{
			Name:  portPrefix + "_PROTO",
			Value: "tcp",
		},
		{
			Name:  portPrefix + "_PORT",
			Value: strconv.Itoa(service.Port),
		},
		{
			Name:  portPrefix + "_ADDR",
			Value: machine,
		},
	}
}

// GetServiceEnvironmentVariables populates a list of environment variables that are use
// in the container environment to get access to services.
func GetServiceEnvironmentVariables(registry ServiceRegistry, machine string) ([]api.EnvVar, error) {
	var result []api.EnvVar
	services, err := registry.ListServices()
	if err != nil {
		return result, err
	}
	for _, service := range services.Items {
		name := strings.ToUpper(service.ID) + "_SERVICE_PORT"
		value := strconv.Itoa(service.Port)
		result = append(result, api.EnvVar{Name: name, Value: value})
		result = append(result, makeLinkVariables(service, machine)...)
	}
	result = append(result, api.EnvVar{Name: "SERVICE_HOST", Value: machine})
	return result, nil
}

func (sr *ServiceRegistryStorage) List(selector labels.Selector) (interface{}, error) {
	list, err := sr.registry.ListServices()
	if err != nil {
		return nil, err
	}
	var filtered []api.Service
	for _, service := range list.Items {
		if selector.Matches(labels.Set(service.Labels)) {
			filtered = append(filtered, service)
		}
	}
	list.Items = filtered
	return list, err
}

func (sr *ServiceRegistryStorage) Get(id string) (interface{}, error) {
	service, err := sr.registry.GetService(id)
	if err != nil {
		return nil, err
	}
	return service, err
}

func (sr *ServiceRegistryStorage) Delete(id string) (<-chan interface{}, error) {
	service, err := sr.registry.GetService(id)
	if err != nil {
		return nil, err
	}
	return apiserver.MakeAsync(func() (interface{}, error) {
		if service.CreateExternalLoadBalancer {
			var balancer cloudprovider.TCPLoadBalancer
			var ok bool
			if sr.cloud != nil {
				balancer, ok = sr.cloud.TCPLoadBalancer()
			}
			if ok && balancer != nil {
				err = balancer.DeleteTCPLoadBalancer(id, "us-central1")
				if err != nil {
					return nil, err
				}
			}
		}
		return api.Status{Status: api.StatusSuccess}, sr.registry.DeleteService(id)
	}), nil
}

func (sr *ServiceRegistryStorage) Extract(body []byte) (interface{}, error) {
	var svc api.Service
	err := api.DecodeInto(body, &svc)
	return svc, err
}

func (sr *ServiceRegistryStorage) Create(obj interface{}) (<-chan interface{}, error) {
	srv := obj.(api.Service)
	if errs := api.ValidateService(&srv); len(errs) > 0 {
		return nil, fmt.Errorf("Validation errors: %v", errs)
	}
	return apiserver.MakeAsync(func() (interface{}, error) {
		// TODO: Consider moving this to a rectification loop, so that we make/remove external load balancers
		// correctly no matter what http operations happen.
		if srv.CreateExternalLoadBalancer {
			var balancer cloudprovider.TCPLoadBalancer
			var zones cloudprovider.Zones
			var ok bool
			if sr.cloud != nil {
				balancer, ok = sr.cloud.TCPLoadBalancer()
				if ok {
					zones, ok = sr.cloud.Zones()
				}
			}
			if ok && balancer != nil && zones != nil {
				hosts, err := sr.machines.List()
				if err != nil {
					return nil, err
				}
				zone, err := zones.GetZone()
				if err != nil {
					return nil, err
				}
				err = balancer.CreateTCPLoadBalancer(srv.ID, zone, srv.Port, hosts)
				if err != nil {
					return nil, err
				}
			} else {
				return nil, fmt.Errorf("requested an external service, but no cloud provider supplied.")
			}
		}
		// TODO actually wait for the object to be fully created here.
		err := sr.registry.CreateService(srv)
		if err != nil {
			return nil, err
		}
		return sr.registry.GetService(srv.ID)
	}), nil
}

func (sr *ServiceRegistryStorage) Update(obj interface{}) (<-chan interface{}, error) {
	srv := obj.(api.Service)
	if srv.ID == "" {
		return nil, fmt.Errorf("ID should not be empty: %#v", srv)
	}
	if errs := api.ValidateService(&srv); len(errs) > 0 {
		return nil, fmt.Errorf("Validation errors: %v", errs)
	}
	return apiserver.MakeAsync(func() (interface{}, error) {
		// TODO: check to see if external load balancer status changed
		err := sr.registry.UpdateService(srv)
		if err != nil {
			return nil, err
		}
		return sr.registry.GetService(srv.ID)
	}), nil
}
<commit_msg>Reduce nesting on ServiceRegistryStorage.Create<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 registry

import (
	"fmt"
	"strconv"
	"strings"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/apiserver"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/cloudprovider"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
)

// ServiceRegistryStorage adapts a service registry into apiserver's RESTStorage model.
type ServiceRegistryStorage struct {
	registry ServiceRegistry
	cloud    cloudprovider.Interface
	machines MinionRegistry
}

// MakeServiceRegistryStorage makes a new ServiceRegistryStorage.
func MakeServiceRegistryStorage(registry ServiceRegistry, cloud cloudprovider.Interface, machines MinionRegistry) apiserver.RESTStorage {
	return &ServiceRegistryStorage{
		registry: registry,
		cloud:    cloud,
		machines: machines,
	}
}

func makeLinkVariables(service api.Service, machine string) []api.EnvVar {
	prefix := strings.ToUpper(service.ID)
	var port string
	if service.ContainerPort.Kind == util.IntstrString {
		port = service.ContainerPort.StrVal
	} else {
		port = strconv.Itoa(service.ContainerPort.IntVal)
	}
	portPrefix := prefix + "_PORT_" + strings.ToUpper(strings.Replace(port, "-", "_", -1)) + "_TCP"
	return []api.EnvVar{
		{
			Name:  prefix + "_PORT",
			Value: fmt.Sprintf("tcp://%s:%d", machine, service.Port),
		},
		{
			Name:  portPrefix,
			Value: fmt.Sprintf("tcp://%s:%d", machine, service.Port),
		},
		{
			Name:  portPrefix + "_PROTO",
			Value: "tcp",
		},
		{
			Name:  portPrefix + "_PORT",
			Value: strconv.Itoa(service.Port),
		},
		{
			Name:  portPrefix + "_ADDR",
			Value: machine,
		},
	}
}

// GetServiceEnvironmentVariables populates a list of environment variables that are use
// in the container environment to get access to services.
func GetServiceEnvironmentVariables(registry ServiceRegistry, machine string) ([]api.EnvVar, error) {
	var result []api.EnvVar
	services, err := registry.ListServices()
	if err != nil {
		return result, err
	}
	for _, service := range services.Items {
		name := strings.ToUpper(service.ID) + "_SERVICE_PORT"
		value := strconv.Itoa(service.Port)
		result = append(result, api.EnvVar{Name: name, Value: value})
		result = append(result, makeLinkVariables(service, machine)...)
	}
	result = append(result, api.EnvVar{Name: "SERVICE_HOST", Value: machine})
	return result, nil
}

func (sr *ServiceRegistryStorage) List(selector labels.Selector) (interface{}, error) {
	list, err := sr.registry.ListServices()
	if err != nil {
		return nil, err
	}
	var filtered []api.Service
	for _, service := range list.Items {
		if selector.Matches(labels.Set(service.Labels)) {
			filtered = append(filtered, service)
		}
	}
	list.Items = filtered
	return list, err
}

func (sr *ServiceRegistryStorage) Get(id string) (interface{}, error) {
	service, err := sr.registry.GetService(id)
	if err != nil {
		return nil, err
	}
	return service, err
}

func (sr *ServiceRegistryStorage) Delete(id string) (<-chan interface{}, error) {
	service, err := sr.registry.GetService(id)
	if err != nil {
		return nil, err
	}
	return apiserver.MakeAsync(func() (interface{}, error) {
		if service.CreateExternalLoadBalancer {
			var balancer cloudprovider.TCPLoadBalancer
			var ok bool
			if sr.cloud != nil {
				balancer, ok = sr.cloud.TCPLoadBalancer()
			}
			if ok && balancer != nil {
				err = balancer.DeleteTCPLoadBalancer(id, "us-central1")
				if err != nil {
					return nil, err
				}
			}
		}
		return api.Status{Status: api.StatusSuccess}, sr.registry.DeleteService(id)
	}), nil
}

func (sr *ServiceRegistryStorage) Extract(body []byte) (interface{}, error) {
	var svc api.Service
	err := api.DecodeInto(body, &svc)
	return svc, err
}

func (sr *ServiceRegistryStorage) Create(obj interface{}) (<-chan interface{}, error) {
	srv := obj.(api.Service)
	if errs := api.ValidateService(&srv); len(errs) > 0 {
		return nil, fmt.Errorf("Validation errors: %v", errs)
	}
	return apiserver.MakeAsync(func() (interface{}, error) {
		// TODO: Consider moving this to a rectification loop, so that we make/remove external load balancers
		// correctly no matter what http operations happen.
		if srv.CreateExternalLoadBalancer {
			if sr.cloud == nil {
				return nil, fmt.Errorf("requested an external service, but no cloud provider supplied.")
			}
			balancer, ok := sr.cloud.TCPLoadBalancer()
			if !ok {
				return nil, fmt.Errorf("The cloud provider does not support external TCP load balancers.")
			}
			zones, ok := sr.cloud.Zones()
			if !ok {
				return nil, fmt.Errorf("The cloud provider does not support zone enumeration.")
			}
			hosts, err := sr.machines.List()
			if err != nil {
				return nil, err
			}
			zone, err := zones.GetZone()
			if err != nil {
				return nil, err
			}
			err = balancer.CreateTCPLoadBalancer(srv.ID, zone, srv.Port, hosts)
			if err != nil {
				return nil, err
			}
		}
		// TODO actually wait for the object to be fully created here.
		err := sr.registry.CreateService(srv)
		if err != nil {
			return nil, err
		}
		return sr.registry.GetService(srv.ID)
	}), nil
}

func (sr *ServiceRegistryStorage) Update(obj interface{}) (<-chan interface{}, error) {
	srv := obj.(api.Service)
	if srv.ID == "" {
		return nil, fmt.Errorf("ID should not be empty: %#v", srv)
	}
	if errs := api.ValidateService(&srv); len(errs) > 0 {
		return nil, fmt.Errorf("Validation errors: %v", errs)
	}
	return apiserver.MakeAsync(func() (interface{}, error) {
		// TODO: check to see if external load balancer status changed
		err := sr.registry.UpdateService(srv)
		if err != nil {
			return nil, err
		}
		return sr.registry.GetService(srv.ID)
	}), nil
}
<|endoftext|>


<commit_before>// +build unittest

package storkctl

import (
	"strings"
	"testing"

	storkv1 "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	"github.com/portworx/sched-ops/k8s"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestGetMigrationsNoMigration(t *testing.T) {
	cmdArgs := []string{"migrations"}

	var migrationList storkv1.MigrationList
	expected := "No resources found.\n"
	testCommon(t, newGetCommand, cmdArgs, &migrationList, expected, false)
}

func createMigrationAndVerify(t *testing.T, name string, clusterpair string, namespaces []string) {
	cmdArgs := []string{"migrations", "-c", clusterpair, "--namespaces", strings.Join(namespaces, ","), name}

	var migrationList storkv1.MigrationList
	expected := "Migration " + name + " created successfully\n"
	testCommon(t, newCreateCommand, cmdArgs, &migrationList, expected, false)

	// Make sure it was created correctly
	migration, err := k8s.Instance().GetMigration(name, "test")
	require.NoError(t, err, "Error getting migration")
	require.Equal(t, name, migration.Name, "Migration name mismatch")
	require.Equal(t, "test", migration.Namespace, "Migration namespace mismatch")
	require.Equal(t, clusterpair, migration.Spec.ClusterPair, "Migration clusterpair mismatch")
	require.Equal(t, namespaces, migration.Spec.Namespaces, "Migration namespace mismatch")
}

func TestGetMigrationsOneMigration(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationtest", "clusterpair1", []string{"namespace1"})

	expected := "NAME               CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest   clusterpair1                       0/0       0/0         \n"

	cmdArgs := []string{"migrations"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestGetMigrationsMultiple(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationtest1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "getmigrationtest2", "clusterpair2", []string{"namespace1"})

	expected := "NAME                CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest1   clusterpair1                       0/0       0/0         \n" +
		"getmigrationtest2   clusterpair2                       0/0       0/0         \n"

	cmdArgs := []string{"migrations", "getmigrationtest1", "getmigrationtest2"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)

	// Should get all migrations if no name given
	cmdArgs = []string{"migrations"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)

	expected = "NAME                CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest1   clusterpair1                       0/0       0/0         \n"
	// Should get only one migration if name given
	cmdArgs = []string{"migrations", "getmigrationtest1"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestGetMigrationsWithStatusAndProgress(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationstatustest", "clusterpair1", []string{"namespace1"})
	migration, err := k8s.Instance().GetMigration("getmigrationstatustest", "test")
	require.NoError(t, err, "Error getting migration")

	// Update the status of the migration
	migration.CreationTimestamp = metav1.Now()
	migration.Status.Stage = storkv1.MigrationStageFinal
	migration.Status.Status = storkv1.MigrationStatusSuccessful
	migration.Status.Volumes = []*storkv1.VolumeInfo{}
	migration, err = k8s.Instance().UpdateMigration(migration)

	expected := "NAME                     CLUSTERPAIR    STAGE     STATUS       VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationstatustest   clusterpair1   Final     Successful   0/0       0/0         " + toTimeString(migration.CreationTimestamp) + "\n"
	cmdArgs := []string{"migrations", "getmigrationstatustest"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestCreateMigrationsNoNamespace(t *testing.T) {
	cmdArgs := []string{"migrations", "-c", "clusterPair1", "migration1"}

	expected := "error: Need to provide atleast one namespace to migrate"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestCreateMigrationsNoClusterPair(t *testing.T) {
	cmdArgs := []string{"migrations", "migration1"}

	expected := "error: ClusterPair name needs to be provided for migration"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestCreateMigrationsNoName(t *testing.T) {
	cmdArgs := []string{"migrations"}

	expected := "error: Exactly one name needs to be provided for migration name"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestCreateMigrations(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "createmigration", "clusterpair1", []string{"namespace1"})
}

func TestDeleteMigrationsNoMigrationName(t *testing.T) {
	cmdArgs := []string{"migrations"}

	var migrationList storkv1.MigrationList
	expected := "error: At least one argument needs to be provided for migration name"
	testCommon(t, newDeleteCommand, cmdArgs, &migrationList, expected, true)
}

func TestDeleteMigrationsNoMigration(t *testing.T) {
	cmdArgs := []string{"migrations", "-c", "migration1"}

	expected := "No resources found.\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)
}

func TestDeleteMigrations(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "deletemigration", "clusterpair1", []string{"namespace1"})

	cmdArgs := []string{"migrations", "deletemigration"}
	expected := "Migration deletemigration deleted successfully\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)

	cmdArgs = []string{"migrations", "deletemigration"}
	expected = "Error from server (NotFound): migrations.stork.libopenstorage.org \"deletemigration\" not found"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, true)

	createMigrationAndVerify(t, "deletemigration1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "deletemigration2", "clusterpair2", []string{"namespace1"})

	cmdArgs = []string{"migrations", "deletemigration1", "deletemigration2"}
	expected = "Migration deletemigration1 deleted successfully\n"
	expected += "Migration deletemigration2 deleted successfully\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)

	createMigrationAndVerify(t, "deletemigration1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "deletemigration2", "clusterpair1", []string{"namespace1"})

	cmdArgs = []string{"migrations", "-c", "clusterpair1"}
	expected = "Migration deletemigration1 deleted successfully\n"
	expected += "Migration deletemigration2 deleted successfully\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)
}
<commit_msg>UTs for storkctl migration subcommand<commit_after>// +build unittest

package storkctl

import (
	"strings"
	"testing"

	storkv1 "github.com/libopenstorage/stork/pkg/apis/stork/v1alpha1"
	"github.com/portworx/sched-ops/k8s"
	"github.com/stretchr/testify/require"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func TestGetMigrationsNoMigration(t *testing.T) {
	cmdArgs := []string{"migrations"}

	var migrationList storkv1.MigrationList
	expected := "No resources found.\n"
	testCommon(t, newGetCommand, cmdArgs, &migrationList, expected, false)
}

func createMigrationAndVerify(t *testing.T, name string, clusterpair string, namespaces []string) {
	cmdArgs := []string{"migrations", "-c", clusterpair, "--namespaces", strings.Join(namespaces, ","), name}

	expected := "Migration " + name + " created successfully\n"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, false)

	// Make sure it was created correctly
	migration, err := k8s.Instance().GetMigration(name, "test")
	require.NoError(t, err, "Error getting migration")
	require.Equal(t, name, migration.Name, "Migration name mismatch")
	require.Equal(t, "test", migration.Namespace, "Migration namespace mismatch")
	require.Equal(t, clusterpair, migration.Spec.ClusterPair, "Migration clusterpair mismatch")
	require.Equal(t, namespaces, migration.Spec.Namespaces, "Migration namespace mismatch")
}

func TestGetMigrationsOneMigration(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationtest", "clusterpair1", []string{"namespace1"})

	expected := "NAME               CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest   clusterpair1                       0/0       0/0         \n"

	cmdArgs := []string{"migrations"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestGetMigrationsMultiple(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationtest1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "getmigrationtest2", "clusterpair2", []string{"namespace1"})

	expected := "NAME                CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest1   clusterpair1                       0/0       0/0         \n" +
		"getmigrationtest2   clusterpair2                       0/0       0/0         \n"

	cmdArgs := []string{"migrations", "getmigrationtest1", "getmigrationtest2"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)

	// Should get all migrations if no name given
	cmdArgs = []string{"migrations"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)

	expected = "NAME                CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest1   clusterpair1                       0/0       0/0         \n"
	// Should get only one migration if name given
	cmdArgs = []string{"migrations", "getmigrationtest1"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestGetMigrationsWithClusterPair(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationtest1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "getmigrationtest2", "clusterpair2", []string{"namespace1"})

	expected := "NAME                CLUSTERPAIR    STAGE     STATUS    VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationtest1   clusterpair1                       0/0       0/0         \n"

	cmdArgs := []string{"migrations", "-c", "clusterpair1"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestGetMigrationsWithStatusAndProgress(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "getmigrationstatustest", "clusterpair1", []string{"namespace1"})
	migration, err := k8s.Instance().GetMigration("getmigrationstatustest", "test")
	require.NoError(t, err, "Error getting migration")

	// Update the status of the migration
	migration.CreationTimestamp = metav1.Now()
	migration.Status.Stage = storkv1.MigrationStageFinal
	migration.Status.Status = storkv1.MigrationStatusSuccessful
	migration.Status.Volumes = []*storkv1.VolumeInfo{}
	migration, err = k8s.Instance().UpdateMigration(migration)

	expected := "NAME                     CLUSTERPAIR    STAGE     STATUS       VOLUMES   RESOURCES   CREATED\n" +
		"getmigrationstatustest   clusterpair1   Final     Successful   0/0       0/0         " + toTimeString(migration.CreationTimestamp) + "\n"
	cmdArgs := []string{"migrations", "getmigrationstatustest"}
	testCommon(t, newGetCommand, cmdArgs, nil, expected, false)
}

func TestCreateMigrationsNoNamespace(t *testing.T) {
	cmdArgs := []string{"migrations", "-c", "clusterPair1", "migration1"}

	expected := "error: Need to provide atleast one namespace to migrate"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestCreateMigrationsNoClusterPair(t *testing.T) {
	cmdArgs := []string{"migrations", "migration1"}

	expected := "error: ClusterPair name needs to be provided for migration"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestCreateMigrationsNoName(t *testing.T) {
	cmdArgs := []string{"migrations"}

	expected := "error: Exactly one name needs to be provided for migration name"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestCreateMigrations(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "createmigration", "clusterpair1", []string{"namespace1"})
}

func TestCreateDuplicateMigrations(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "createmigration", "clusterpair1", []string{"namespace1"})
	cmdArgs := []string{"migrations", "-c", "clusterpair1", "--namespaces", "namespace1", "createmigration"}

	expected := "Error from server (AlreadyExists): migrations.stork.libopenstorage.org \"createmigration\" already exists"
	testCommon(t, newCreateCommand, cmdArgs, nil, expected, true)
}

func TestDeleteMigrationsNoMigrationName(t *testing.T) {
	cmdArgs := []string{"migrations"}

	var migrationList storkv1.MigrationList
	expected := "error: At least one argument needs to be provided for migration name"
	testCommon(t, newDeleteCommand, cmdArgs, &migrationList, expected, true)
}

func TestDeleteMigrationsNoMigration(t *testing.T) {
	cmdArgs := []string{"migrations", "-c", "migration1"}

	expected := "No resources found.\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)
}

func TestDeleteMigrations(t *testing.T) {
	defer resetTest()
	createMigrationAndVerify(t, "deletemigration", "clusterpair1", []string{"namespace1"})

	cmdArgs := []string{"migrations", "deletemigration"}
	expected := "Migration deletemigration deleted successfully\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)

	cmdArgs = []string{"migrations", "deletemigration"}
	expected = "Error from server (NotFound): migrations.stork.libopenstorage.org \"deletemigration\" not found"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, true)

	createMigrationAndVerify(t, "deletemigration1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "deletemigration2", "clusterpair2", []string{"namespace1"})

	cmdArgs = []string{"migrations", "deletemigration1", "deletemigration2"}
	expected = "Migration deletemigration1 deleted successfully\n"
	expected += "Migration deletemigration2 deleted successfully\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)

	createMigrationAndVerify(t, "deletemigration1", "clusterpair1", []string{"namespace1"})
	createMigrationAndVerify(t, "deletemigration2", "clusterpair1", []string{"namespace1"})

	cmdArgs = []string{"migrations", "-c", "clusterpair1"}
	expected = "Migration deletemigration1 deleted successfully\n"
	expected += "Migration deletemigration2 deleted successfully\n"
	testCommon(t, newDeleteCommand, cmdArgs, nil, expected, false)
}
<|endoftext|>


<commit_before>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package router

import (
	"sync"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/engine/common"
	"github.com/ava-labs/avalanchego/snow/networking/benchlist"
	"github.com/ava-labs/avalanchego/snow/networking/timeout"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/timer"
)

func TestShutdown(t *testing.T) {
	vdrs := validators.NewSet()
	err := vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	benchlist := benchlist.NewNoBenchlist()
	tm := timeout.Manager{}
	err = tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     time.Millisecond,
			MinimumTimeout:     time.Millisecond,
			MaximumTimeout:     10 * time.Second,
			TimeoutCoefficient: 1.25,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist,
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Second, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	engine := common.EngineTest{T: t}
	engine.Default(false)

	shutdownCalled := make(chan struct{}, 1)

	engine.ContextF = snow.DefaultContextTest
	engine.ShutdownF = func() error { shutdownCalled <- struct{}{}; return nil }

	handler := &Handler{}
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	go handler.Dispatch()

	chainRouter.AddChain(handler)

	chainRouter.Shutdown()

	ticker := time.NewTicker(250 * time.Millisecond)
	select {
	case <-ticker.C:
		t.Fatalf("Handler shutdown was not called or timed out after 250ms during chainRouter shutdown")
	case <-shutdownCalled:
	}

	select {
	case <-handler.closed:
	default:
		t.Fatal("handler shutdown but never closed its closing channel")
	}
}

func TestShutdownTimesOut(t *testing.T) {
	vdrs := validators.NewSet()
	err := vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	benchlist := benchlist.NewNoBenchlist()
	tm := timeout.Manager{}
	// Ensure that the MultiPut request does not timeout
	err = tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     time.Second,
			MinimumTimeout:     500 * time.Millisecond,
			MaximumTimeout:     10 * time.Second,
			TimeoutCoefficient: 1.25,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist,
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	engine := common.EngineTest{T: t}
	engine.Default(false)

	engineFinished := make(chan struct{}, 1)

	// MultiPut blocks for two seconds
	engine.MultiPutF = func(nodeID ids.ShortID, requestID uint32, containers [][]byte) error {
		time.Sleep(2 * time.Second)
		engineFinished <- struct{}{}
		return nil
	}

	closed := new(int)

	engine.ContextF = snow.DefaultContextTest
	engine.ShutdownF = func() error { *closed++; return nil }

	handler := &Handler{}
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)

	go handler.Dispatch()

	shutdownFinished := make(chan struct{}, 1)

	go func() {
		handler.MultiPut(ids.ShortID{}, 1, nil, func() {})
		time.Sleep(50 * time.Millisecond) // Pause to ensure message gets processed

		chainRouter.Shutdown()
		shutdownFinished <- struct{}{}
	}()

	select {
	case <-engineFinished:
		t.Fatalf("Shutdown should have finished in one millisecond before timing out instead of waiting for engine to finish shutting down.")
	case <-shutdownFinished:
	}
}

// Ensure that a timeout fires if we don't get a response to a request
func TestRouterTimeout(t *testing.T) {
	// Create a timeout manager
	maxTimeout := 25 * time.Millisecond
	tm := timeout.Manager{}
	err := tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     10 * time.Millisecond,
			MinimumTimeout:     10 * time.Millisecond,
			MaximumTimeout:     maxTimeout,
			TimeoutCoefficient: 1,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist.NewNoBenchlist(),
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	// Create a router
	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	// Create an engine and handler
	engine := common.EngineTest{T: t}
	engine.Default(false)

	var (
		calledGetFailed, calledGetAncestorsFailed,
		calledQueryFailed, calledQueryFailed2,
		calledGetAcceptedFailed, calledGetAcceptedFrontierFailed bool

		wg = sync.WaitGroup{}
	)

	engine.GetFailedF = func(nodeID ids.ShortID, requestID uint32) error { wg.Done(); calledGetFailed = true; return nil }
	engine.GetAncestorsFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAncestorsFailed = true
		return nil
	}
	engine.QueryFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		if !calledQueryFailed {
			calledQueryFailed = true
			return nil
		}
		calledQueryFailed2 = true
		return nil
	}
	engine.GetAcceptedFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAcceptedFailed = true
		return nil
	}
	engine.GetAcceptedFrontierFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAcceptedFrontierFailed = true
		return nil
	}

	engine.ContextF = snow.DefaultContextTest

	handler := &Handler{}
	vdrs := validators.NewSet()
	err = vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)
	go handler.Dispatch()

	// Register requests for each request type
	msgs := []constants.MsgType{
		constants.GetMsg,
		constants.GetAncestorsMsg,
		constants.PullQueryMsg,
		constants.PushQueryMsg,
		constants.GetAcceptedMsg,
		constants.GetAcceptedFrontierMsg,
	}

	wg.Add(len(msgs))

	for i, msg := range msgs {
		chainRouter.RegisterRequest(ids.GenerateTestShortID(), handler.ctx.ChainID, uint32(i), msg)
	}

	wg.Wait()
	chainRouter.lock.Lock()
	defer chainRouter.lock.Unlock()
	assert.True(t, calledGetFailed && calledGetAncestorsFailed && calledQueryFailed2 && calledGetAcceptedFailed && calledGetAcceptedFrontierFailed)
}

func TestRouterClearTimeouts(t *testing.T) {
	// Create a timeout manager
	tm := timeout.Manager{}
	err := tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     3 * time.Second,
			MinimumTimeout:     3 * time.Second,
			MaximumTimeout:     5 * time.Minute,
			TimeoutCoefficient: 1,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist.NewNoBenchlist(),
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	// Create a router
	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	// Create an engine and handler
	engine := common.EngineTest{T: t}
	engine.Default(false)

	engine.ContextF = snow.DefaultContextTest

	vdrs := validators.NewSet()
	err = vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	handler := &Handler{}
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)
	go handler.Dispatch()

	// Register requests for each request type
	msgs := []constants.MsgType{
		constants.GetMsg,
		constants.GetAncestorsMsg,
		constants.PullQueryMsg,
		constants.PushQueryMsg,
		constants.GetAcceptedMsg,
		constants.GetAcceptedFrontierMsg,
	}

	vID := ids.GenerateTestShortID()
	for i, msg := range msgs {
		chainRouter.RegisterRequest(vID, handler.ctx.ChainID, uint32(i), msg)
	}

	// Clear each timeout by simulating responses to the queries
	// Note: Depends on the ordering of [msgs]
	chainRouter.Put(vID, handler.ctx.ChainID, 0, ids.GenerateTestID(), nil, nil)
	chainRouter.MultiPut(vID, handler.ctx.ChainID, 1, nil, nil)
	chainRouter.Chits(vID, handler.ctx.ChainID, 2, nil, nil)
	chainRouter.Chits(vID, handler.ctx.ChainID, 3, nil, nil)
	chainRouter.Accepted(vID, handler.ctx.ChainID, 4, nil, nil)
	chainRouter.AcceptedFrontier(vID, handler.ctx.ChainID, 5, nil, nil)

	assert.Equal(t, chainRouter.timedRequests.Len(), 0)
}

func TestValidatorOnlyMessageDrops(t *testing.T) {
	// Create a timeout manager
	maxTimeout := 25 * time.Millisecond
	tm := timeout.Manager{}
	err := tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     10 * time.Millisecond,
			MinimumTimeout:     10 * time.Millisecond,
			MaximumTimeout:     maxTimeout,
			TimeoutCoefficient: 1,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist.NewNoBenchlist(),
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	// Create a router
	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	// Create an engine and handler
	engine := common.EngineTest{T: t}
	engine.Default(false)

	var (
		calledGetFailed, calledGetAncestorsFailed,
		calledQueryFailed, calledQueryFailed2,
		calledGetAcceptedFailed, calledGetAcceptedFrontierFailed bool

		wg = sync.WaitGroup{}
	)

	engine.GetFailedF = func(nodeID ids.ShortID, requestID uint32) error { wg.Done(); calledGetFailed = true; return nil }
	engine.GetAncestorsFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAncestorsFailed = true
		return nil
	}
	engine.QueryFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		if !calledQueryFailed {
			calledQueryFailed = true
			return nil
		}
		calledQueryFailed2 = true
		return nil
	}
	engine.GetAcceptedFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAcceptedFailed = true
		return nil
	}
	engine.GetAcceptedFrontierFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAcceptedFrontierFailed = true
		return nil
	}

	engine.ContextF = func() *snow.Context {
		ctx := snow.DefaultContextTest()
		ctx.SetValidatorOnly()
		return ctx
	}

	handler := &Handler{}
	vdrs := validators.NewSet()
	vID := ids.GenerateTestShortID()
	err = vdrs.AddWeight(vID, 1)
	assert.NoError(t, err)
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)
	go handler.Dispatch()

	// Register requests for each request type
	msgs := []constants.MsgType{
		constants.GetMsg,
		constants.GetAncestorsMsg,
		constants.PullQueryMsg,
		constants.PushQueryMsg,
		constants.GetAcceptedMsg,
		constants.GetAcceptedFrontierMsg,
	}

	// generate a non-validator ID
	nID := ids.GenerateTestShortID()
	for i, msg := range msgs {
		chainRouter.RegisterRequest(nID, handler.ctx.ChainID, uint32(i), msg)
	}
	// should not be called
	assert.False(t, calledGetFailed && calledGetAncestorsFailed && calledQueryFailed2 && calledGetAcceptedFailed && calledGetAcceptedFrontierFailed)

	// validator case
	wg.Add(len(msgs))

	for i, msg := range msgs {
		chainRouter.RegisterRequest(vID, handler.ctx.ChainID, uint32(i), msg)
	}

	wg.Wait()

	// should be called since this is a validator request
	assert.True(t, calledGetFailed && calledGetAncestorsFailed && calledQueryFailed2 && calledGetAcceptedFailed && calledGetAcceptedFrontierFailed)

	// register a validator request
	reqID := uint32(3333)
	chainRouter.RegisterRequest(vID, handler.ctx.ChainID, reqID, constants.GetMsg)
	assert.Equal(t, 1, chainRouter.timedRequests.Len())
	// remove it from validators
	err = handler.validators.Set(validators.NewSet().List())
	assert.NoError(t, err)
	calledFinished := false
	chainRouter.Put(vID, handler.ctx.ChainID, reqID, ids.GenerateTestID(), nil, func() { calledFinished = true })
	assert.True(t, calledFinished)
	// should clear out timed request
	assert.Equal(t, 0, chainRouter.timedRequests.Len())
}
<commit_msg>fix tests after removal of isValidator logic from failed messages<commit_after>// (c) 2019-2020, Ava Labs, Inc. All rights reserved.
// See the file LICENSE for licensing terms.

package router

import (
	"sync"
	"testing"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/stretchr/testify/assert"

	"github.com/ava-labs/avalanchego/ids"
	"github.com/ava-labs/avalanchego/snow"
	"github.com/ava-labs/avalanchego/snow/engine/common"
	"github.com/ava-labs/avalanchego/snow/networking/benchlist"
	"github.com/ava-labs/avalanchego/snow/networking/timeout"
	"github.com/ava-labs/avalanchego/snow/validators"
	"github.com/ava-labs/avalanchego/utils/constants"
	"github.com/ava-labs/avalanchego/utils/logging"
	"github.com/ava-labs/avalanchego/utils/timer"
)

func TestShutdown(t *testing.T) {
	vdrs := validators.NewSet()
	err := vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	benchlist := benchlist.NewNoBenchlist()
	tm := timeout.Manager{}
	err = tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     time.Millisecond,
			MinimumTimeout:     time.Millisecond,
			MaximumTimeout:     10 * time.Second,
			TimeoutCoefficient: 1.25,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist,
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Second, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	engine := common.EngineTest{T: t}
	engine.Default(false)

	shutdownCalled := make(chan struct{}, 1)

	engine.ContextF = snow.DefaultContextTest
	engine.ShutdownF = func() error { shutdownCalled <- struct{}{}; return nil }

	handler := &Handler{}
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	go handler.Dispatch()

	chainRouter.AddChain(handler)

	chainRouter.Shutdown()

	ticker := time.NewTicker(250 * time.Millisecond)
	select {
	case <-ticker.C:
		t.Fatalf("Handler shutdown was not called or timed out after 250ms during chainRouter shutdown")
	case <-shutdownCalled:
	}

	select {
	case <-handler.closed:
	default:
		t.Fatal("handler shutdown but never closed its closing channel")
	}
}

func TestShutdownTimesOut(t *testing.T) {
	vdrs := validators.NewSet()
	err := vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	benchlist := benchlist.NewNoBenchlist()
	tm := timeout.Manager{}
	// Ensure that the MultiPut request does not timeout
	err = tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     time.Second,
			MinimumTimeout:     500 * time.Millisecond,
			MaximumTimeout:     10 * time.Second,
			TimeoutCoefficient: 1.25,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist,
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	engine := common.EngineTest{T: t}
	engine.Default(false)

	engineFinished := make(chan struct{}, 1)

	// MultiPut blocks for two seconds
	engine.MultiPutF = func(nodeID ids.ShortID, requestID uint32, containers [][]byte) error {
		time.Sleep(2 * time.Second)
		engineFinished <- struct{}{}
		return nil
	}

	closed := new(int)

	engine.ContextF = snow.DefaultContextTest
	engine.ShutdownF = func() error { *closed++; return nil }

	handler := &Handler{}
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)

	go handler.Dispatch()

	shutdownFinished := make(chan struct{}, 1)

	go func() {
		handler.MultiPut(ids.ShortID{}, 1, nil, func() {})
		time.Sleep(50 * time.Millisecond) // Pause to ensure message gets processed

		chainRouter.Shutdown()
		shutdownFinished <- struct{}{}
	}()

	select {
	case <-engineFinished:
		t.Fatalf("Shutdown should have finished in one millisecond before timing out instead of waiting for engine to finish shutting down.")
	case <-shutdownFinished:
	}
}

// Ensure that a timeout fires if we don't get a response to a request
func TestRouterTimeout(t *testing.T) {
	// Create a timeout manager
	maxTimeout := 25 * time.Millisecond
	tm := timeout.Manager{}
	err := tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     10 * time.Millisecond,
			MinimumTimeout:     10 * time.Millisecond,
			MaximumTimeout:     maxTimeout,
			TimeoutCoefficient: 1,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist.NewNoBenchlist(),
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	// Create a router
	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	// Create an engine and handler
	engine := common.EngineTest{T: t}
	engine.Default(false)

	var (
		calledGetFailed, calledGetAncestorsFailed,
		calledQueryFailed, calledQueryFailed2,
		calledGetAcceptedFailed, calledGetAcceptedFrontierFailed bool

		wg = sync.WaitGroup{}
	)

	engine.GetFailedF = func(nodeID ids.ShortID, requestID uint32) error { wg.Done(); calledGetFailed = true; return nil }
	engine.GetAncestorsFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAncestorsFailed = true
		return nil
	}
	engine.QueryFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		if !calledQueryFailed {
			calledQueryFailed = true
			return nil
		}
		calledQueryFailed2 = true
		return nil
	}
	engine.GetAcceptedFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAcceptedFailed = true
		return nil
	}
	engine.GetAcceptedFrontierFailedF = func(nodeID ids.ShortID, requestID uint32) error {
		defer wg.Done()
		calledGetAcceptedFrontierFailed = true
		return nil
	}

	engine.ContextF = snow.DefaultContextTest

	handler := &Handler{}
	vdrs := validators.NewSet()
	err = vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)
	go handler.Dispatch()

	// Register requests for each request type
	msgs := []constants.MsgType{
		constants.GetMsg,
		constants.GetAncestorsMsg,
		constants.PullQueryMsg,
		constants.PushQueryMsg,
		constants.GetAcceptedMsg,
		constants.GetAcceptedFrontierMsg,
	}

	wg.Add(len(msgs))

	for i, msg := range msgs {
		chainRouter.RegisterRequest(ids.GenerateTestShortID(), handler.ctx.ChainID, uint32(i), msg)
	}

	wg.Wait()
	chainRouter.lock.Lock()
	defer chainRouter.lock.Unlock()
	assert.True(t, calledGetFailed && calledGetAncestorsFailed && calledQueryFailed2 && calledGetAcceptedFailed && calledGetAcceptedFrontierFailed)
}

func TestRouterClearTimeouts(t *testing.T) {
	// Create a timeout manager
	tm := timeout.Manager{}
	err := tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     3 * time.Second,
			MinimumTimeout:     3 * time.Second,
			MaximumTimeout:     5 * time.Minute,
			TimeoutCoefficient: 1,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist.NewNoBenchlist(),
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	// Create a router
	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	// Create an engine and handler
	engine := common.EngineTest{T: t}
	engine.Default(false)

	engine.ContextF = snow.DefaultContextTest

	vdrs := validators.NewSet()
	err = vdrs.AddWeight(ids.GenerateTestShortID(), 1)
	assert.NoError(t, err)
	handler := &Handler{}
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)
	go handler.Dispatch()

	// Register requests for each request type
	msgs := []constants.MsgType{
		constants.GetMsg,
		constants.GetAncestorsMsg,
		constants.PullQueryMsg,
		constants.PushQueryMsg,
		constants.GetAcceptedMsg,
		constants.GetAcceptedFrontierMsg,
	}

	vID := ids.GenerateTestShortID()
	for i, msg := range msgs {
		chainRouter.RegisterRequest(vID, handler.ctx.ChainID, uint32(i), msg)
	}

	// Clear each timeout by simulating responses to the queries
	// Note: Depends on the ordering of [msgs]
	chainRouter.Put(vID, handler.ctx.ChainID, 0, ids.GenerateTestID(), nil, nil)
	chainRouter.MultiPut(vID, handler.ctx.ChainID, 1, nil, nil)
	chainRouter.Chits(vID, handler.ctx.ChainID, 2, nil, nil)
	chainRouter.Chits(vID, handler.ctx.ChainID, 3, nil, nil)
	chainRouter.Accepted(vID, handler.ctx.ChainID, 4, nil, nil)
	chainRouter.AcceptedFrontier(vID, handler.ctx.ChainID, 5, nil, nil)

	assert.Equal(t, chainRouter.timedRequests.Len(), 0)
}

func TestValidatorOnlyMessageDrops(t *testing.T) {
	// Create a timeout manager
	maxTimeout := 25 * time.Millisecond
	tm := timeout.Manager{}
	err := tm.Initialize(
		&timer.AdaptiveTimeoutConfig{
			InitialTimeout:     10 * time.Millisecond,
			MinimumTimeout:     10 * time.Millisecond,
			MaximumTimeout:     maxTimeout,
			TimeoutCoefficient: 1,
			TimeoutHalflife:    5 * time.Minute,
		},
		benchlist.NewNoBenchlist(),
		"",
		prometheus.NewRegistry(),
	)
	if err != nil {
		t.Fatal(err)
	}
	go tm.Dispatch()

	// Create a router
	chainRouter := ChainRouter{}
	err = chainRouter.Initialize(ids.ShortEmpty, logging.NoLog{}, &tm, time.Hour, time.Millisecond, ids.Set{}, nil, HealthConfig{}, "", prometheus.NewRegistry())
	assert.NoError(t, err)

	// Create an engine and handler
	engine := common.EngineTest{T: t}
	engine.Default(false)

	calledF := new(bool)

	wg := sync.WaitGroup{}

	engine.PullQueryF = func(nodeID ids.ShortID, requestID uint32, containerID ids.ID) error {
		defer wg.Done()
		*calledF = true
		return nil
	}

	engine.ContextF = func() *snow.Context {
		ctx := snow.DefaultContextTest()
		ctx.SetValidatorOnly()
		return ctx
	}

	handler := &Handler{}
	vdrs := validators.NewSet()
	vID := ids.GenerateTestShortID()
	err = vdrs.AddWeight(vID, 1)
	assert.NoError(t, err)
	err = handler.Initialize(
		&engine,
		vdrs,
		nil,
		"",
		prometheus.NewRegistry(),
	)
	assert.NoError(t, err)

	chainRouter.AddChain(handler)
	go handler.Dispatch()

	// generate a non-validator ID
	nID := ids.GenerateTestShortID()

	*calledF = false
	timeout := chainRouter.clock.Time().Add(30 * time.Second)
	chainRouter.PullQuery(nID, handler.ctx.ChainID, uint32(1), timeout, ids.GenerateTestID(), func() {})
	// should not be called
	assert.False(t, *calledF)

	// validator case
	*calledF = false
	wg.Add(1)
	chainRouter.PullQuery(vID, handler.ctx.ChainID, uint32(2), timeout, ids.GenerateTestID(), func() {})
	wg.Wait()
	// should be called since this is a validator request
	assert.True(t, *calledF)

	// register a validator request
	reqID := uint32(3)
	chainRouter.RegisterRequest(vID, handler.ctx.ChainID, reqID, constants.GetMsg)
	assert.Equal(t, 1, chainRouter.timedRequests.Len())
	// remove it from validators
	err = handler.validators.Set(validators.NewSet().List())
	assert.NoError(t, err)
	calledFinished := false
	chainRouter.Put(vID, handler.ctx.ChainID, reqID, ids.GenerateTestID(), nil, func() { calledFinished = true })
	assert.True(t, calledFinished)
	// should clear out timed request
	assert.Equal(t, 0, chainRouter.timedRequests.Len())
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os/exec"
	"strings"

	"github.com/go-check/check"
)

// TestPullImageWithAliases pulls a specific image tag and verifies that any aliases (i.e., other
// tags for the same image) are not also pulled down.
//
// Ref: docker/docker#8141
func (s *DockerRegistrySuite) TestPullImageWithAliases(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	repos := []string{}
	for _, tag := range []string{"recent", "fresh"} {
		repos = append(repos, fmt.Sprintf("%v:%v", repoName, tag))
	}

	// Tag and push the same image multiple times.
	for _, repo := range repos {
		dockerCmd(c, "tag", "busybox", repo)
		dockerCmd(c, "push", repo)
	}

	// Clear local images store.
	args := append([]string{"rmi"}, repos...)
	dockerCmd(c, args...)

	// Pull a single tag and verify it doesn't bring down all aliases.
	dockerCmd(c, "pull", repos[0])
	dockerCmd(c, "inspect", repos[0])
	for _, repo := range repos[1:] {
		if _, _, err := dockerCmdWithError("inspect", repo); err == nil {
			c.Fatalf("Image %v shouldn't have been pulled down", repo)
		}
	}
}

// TestConcurrentPullWholeRepo pulls the same repo concurrently.
func (s *DockerRegistrySuite) TestConcurrentPullWholeRepo(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	repos := []string{}
	for _, tag := range []string{"recent", "fresh", "todays"} {
		repo := fmt.Sprintf("%v:%v", repoName, tag)
		_, err := buildImage(repo, fmt.Sprintf(`
		    FROM busybox
		    ENTRYPOINT ["/bin/echo"]
		    ENV FOO foo
		    ENV BAR bar
		    CMD echo %s
		`, repo), true)
		if err != nil {
			c.Fatal(err)
		}
		dockerCmd(c, "push", repo)
		repos = append(repos, repo)
	}

	// Clear local images store.
	args := append([]string{"rmi"}, repos...)
	dockerCmd(c, args...)

	// Run multiple re-pulls concurrently
	results := make(chan error)
	numPulls := 3

	for i := 0; i != numPulls; i++ {
		go func() {
			_, _, err := runCommandWithOutput(exec.Command(dockerBinary, "pull", "-a", repoName))
			results <- err
		}()
	}

	// These checks are separate from the loop above because the check
	// package is not goroutine-safe.
	for i := 0; i != numPulls; i++ {
		err := <-results
		c.Assert(err, check.IsNil, check.Commentf("concurrent pull failed with error: %v", err))
	}

	// Ensure all tags were pulled successfully
	for _, repo := range repos {
		dockerCmd(c, "inspect", repo)
		out, _ := dockerCmd(c, "run", "--rm", repo)
		if strings.TrimSpace(out) != "/bin/sh -c echo "+repo {
			c.Fatalf("CMD did not contain /bin/sh -c echo %s: %s", repo, out)
		}
	}
}

// TestConcurrentFailingPull tries a concurrent pull that doesn't succeed.
func (s *DockerRegistrySuite) TestConcurrentFailingPull(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	// Run multiple pulls concurrently
	results := make(chan error)
	numPulls := 3

	for i := 0; i != numPulls; i++ {
		go func() {
			_, _, err := runCommandWithOutput(exec.Command(dockerBinary, "pull", repoName+":asdfasdf"))
			results <- err
		}()
	}

	// These checks are separate from the loop above because the check
	// package is not goroutine-safe.
	for i := 0; i != numPulls; i++ {
		err := <-results
		if err == nil {
			c.Fatal("expected pull to fail")
		}
	}
}

// TestConcurrentPullMultipleTags pulls multiple tags from the same repo
// concurrently.
func (s *DockerRegistrySuite) TestConcurrentPullMultipleTags(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	repos := []string{}
	for _, tag := range []string{"recent", "fresh", "todays"} {
		repo := fmt.Sprintf("%v:%v", repoName, tag)
		_, err := buildImage(repo, fmt.Sprintf(`
		    FROM busybox
		    ENTRYPOINT ["/bin/echo"]
		    ENV FOO foo
		    ENV BAR bar
		    CMD echo %s
		`, repo), true)
		if err != nil {
			c.Fatal(err)
		}
		dockerCmd(c, "push", repo)
		repos = append(repos, repo)
	}

	// Clear local images store.
	args := append([]string{"rmi"}, repos...)
	dockerCmd(c, args...)

	// Re-pull individual tags, in parallel
	results := make(chan error)

	for _, repo := range repos {
		go func(repo string) {
			_, _, err := runCommandWithOutput(exec.Command(dockerBinary, "pull", repo))
			results <- err
		}(repo)
	}

	// These checks are separate from the loop above because the check
	// package is not goroutine-safe.
	for range repos {
		err := <-results
		c.Assert(err, check.IsNil, check.Commentf("concurrent pull failed with error: %v", err))
	}

	// Ensure all tags were pulled successfully
	for _, repo := range repos {
		dockerCmd(c, "inspect", repo)
		out, _ := dockerCmd(c, "run", "--rm", repo)
		if strings.TrimSpace(out) != "/bin/sh -c echo "+repo {
			c.Fatalf("CMD did not contain /bin/sh -c echo %s: %s", repo, out)
		}
	}
}
<commit_msg>update pull-local#16756<commit_after>package main

import (
	"fmt"
	"os/exec"
	"strings"

	"github.com/go-check/check"
)

// TestPullImageWithAliases pulls a specific image tag and verifies that any aliases (i.e., other
// tags for the same image) are not also pulled down.
//
// Ref: docker/docker#8141
func (s *DockerRegistrySuite) TestPullImageWithAliases(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	repos := []string{}
	for _, tag := range []string{"recent", "fresh"} {
		repos = append(repos, fmt.Sprintf("%v:%v", repoName, tag))
	}

	// Tag and push the same image multiple times.
	for _, repo := range repos {
		dockerCmd(c, "tag", "busybox", repo)
		dockerCmd(c, "push", repo)
	}

	// Clear local images store.
	args := append([]string{"rmi"}, repos...)
	dockerCmd(c, args...)

	// Pull a single tag and verify it doesn't bring down all aliases.
	dockerCmd(c, "pull", repos[0])
	dockerCmd(c, "inspect", repos[0])
	for _, repo := range repos[1:] {
		_, _, err := dockerCmdWithError("inspect", repo)
		c.Assert(err, check.NotNil, check.Commentf("Image %v shouldn't have been pulled down", repo))
	}
}

// TestConcurrentPullWholeRepo pulls the same repo concurrently.
func (s *DockerRegistrySuite) TestConcurrentPullWholeRepo(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	repos := []string{}
	for _, tag := range []string{"recent", "fresh", "todays"} {
		repo := fmt.Sprintf("%v:%v", repoName, tag)
		_, err := buildImage(repo, fmt.Sprintf(`
		    FROM busybox
		    ENTRYPOINT ["/bin/echo"]
		    ENV FOO foo
		    ENV BAR bar
		    CMD echo %s
		`, repo), true)
		c.Assert(err, check.IsNil)
		dockerCmd(c, "push", repo)
		repos = append(repos, repo)
	}

	// Clear local images store.
	args := append([]string{"rmi"}, repos...)
	dockerCmd(c, args...)

	// Run multiple re-pulls concurrently
	results := make(chan error)
	numPulls := 3

	for i := 0; i != numPulls; i++ {
		go func() {
			_, _, err := runCommandWithOutput(exec.Command(dockerBinary, "pull", "-a", repoName))
			results <- err
		}()
	}

	// These checks are separate from the loop above because the check
	// package is not goroutine-safe.
	for i := 0; i != numPulls; i++ {
		err := <-results
		c.Assert(err, check.IsNil, check.Commentf("concurrent pull failed with error: %v", err))
	}

	// Ensure all tags were pulled successfully
	for _, repo := range repos {
		dockerCmd(c, "inspect", repo)
		out, _ := dockerCmd(c, "run", "--rm", repo)
		c.Assert(strings.TrimSpace(out), check.Equals, "/bin/sh -c echo "+repo, check.Commentf("CMD did not contain /bin/sh -c echo %s: %s", repo, out))
	}
}

// TestConcurrentFailingPull tries a concurrent pull that doesn't succeed.
func (s *DockerRegistrySuite) TestConcurrentFailingPull(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	// Run multiple pulls concurrently
	results := make(chan error)
	numPulls := 3

	for i := 0; i != numPulls; i++ {
		go func() {
			_, _, err := runCommandWithOutput(exec.Command(dockerBinary, "pull", repoName+":asdfasdf"))
			results <- err
		}()
	}

	// These checks are separate from the loop above because the check
	// package is not goroutine-safe.
	for i := 0; i != numPulls; i++ {
		err := <-results
		c.Assert(err, check.NotNil, check.Commentf("expected pull to fail"))
	}
}

// TestConcurrentPullMultipleTags pulls multiple tags from the same repo
// concurrently.
func (s *DockerRegistrySuite) TestConcurrentPullMultipleTags(c *check.C) {
	repoName := fmt.Sprintf("%v/dockercli/busybox", privateRegistryURL)

	repos := []string{}
	for _, tag := range []string{"recent", "fresh", "todays"} {
		repo := fmt.Sprintf("%v:%v", repoName, tag)
		_, err := buildImage(repo, fmt.Sprintf(`
		    FROM busybox
		    ENTRYPOINT ["/bin/echo"]
		    ENV FOO foo
		    ENV BAR bar
		    CMD echo %s
		`, repo), true)
		c.Assert(err, check.IsNil)
		dockerCmd(c, "push", repo)
		repos = append(repos, repo)
	}

	// Clear local images store.
	args := append([]string{"rmi"}, repos...)
	dockerCmd(c, args...)

	// Re-pull individual tags, in parallel
	results := make(chan error)

	for _, repo := range repos {
		go func(repo string) {
			_, _, err := runCommandWithOutput(exec.Command(dockerBinary, "pull", repo))
			results <- err
		}(repo)
	}

	// These checks are separate from the loop above because the check
	// package is not goroutine-safe.
	for range repos {
		err := <-results
		c.Assert(err, check.IsNil, check.Commentf("concurrent pull failed with error: %v", err))
	}

	// Ensure all tags were pulled successfully
	for _, repo := range repos {
		dockerCmd(c, "inspect", repo)
		out, _ := dockerCmd(c, "run", "--rm", repo)
		c.Assert(strings.TrimSpace(out), check.Equals, "/bin/sh -c echo "+repo, check.Commentf("CMD did not contain /bin/sh -c echo %s; %s", repo, out))
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"gonzbee/nntp"
	"gonzbee/nzb"
	"gonzbee/yenc"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type job struct {
	dir string
	n   *nzb.Nzb
}

type messagejob struct {
	group string
	msgId string
	ch    chan []byte
}

func init() {
	go poolHandler()
}

var download = make(chan *messagejob)
var downloadMux = make(chan *messagejob)
var reaper = make(chan int)

func newConnection() error {
	s := config.Server.GetAddressStr()
	var err error
	var n *nntp.Conn
	if config.Server.TLS {
		n, err = nntp.DialTLS(s)
	} else {
		n, err = nntp.Dial(s)
	}
	if err != nil {
		return err
	}

	err = n.Authenticate(config.Server.Username, config.Server.Password)
	if err != nil {
		n.Close()
		return err
	}
	log.Println("spun up nntp connection")
	go func() {
		defer n.Close()
		for {
			select {
			case m := <-downloadMux:
				err = n.SwitchGroup(m.group)
				if err != nil {
					panic(err)
				}
				b, err := n.GetMessage(m.msgId)
				if err != nil {
					log.Print("Error getting Message ", m.msgId, ": ", err.Error())
				}
				m.ch <- b
			case <-(time.After(10 * time.Second)):
				reaper <- 1
				return
			}
		}
	}()
	return nil
}

func poolHandler() {
	var number int
	for {
		select {
		case msg := <-download:
			if number < 10 {
				err := newConnection()
				if err == nil {
					number++
				}
			}
			downloadMux <- msg
		case <-reaper:
			number--
		}
	}
}

func (j *job) handle() {
	wg := new(sync.WaitGroup)
	for _, f := range j.n.File {
		ch := make(chan []byte)
		wg.Add(1)
		partsLeft := len(f.Segments)
		go func() {
			var file *os.File
			var part *yenc.Part
			var err error
			defer wg.Done()
			for ; partsLeft > 0; partsLeft-- {
				s := <-ch
				if s == nil {
					continue
				}
				m := bytes.NewReader(s)
				part, err = yenc.NewPart(m)
				if err != nil {
					log.Print(err.Error())
					continue
				}
				if file == nil {
					file, err = os.Create(filepath.Join(j.dir, part.Filename))
					if err != nil {
						panic("could not create file: " + err.Error())
					}
					defer file.Close()
				}
				file.Seek(part.Begin, os.SEEK_SET)
				part.Decode(file)
			}
			if part != nil {
				log.Print("Done Decoding file " + part.Filename)
			} else {
				log.Print("Could not decode entire file")
			}
		}()
		for _, seg := range f.Segments {
			msg := &messagejob{
				msgId: seg.MsgId,
				group: f.Groups[0],
				ch:    ch,
			}
			download <- msg
		}
	}
	wg.Wait()
}

func jobStart(n *nzb.Nzb, name string) {
	incDir := config.GetIncompleteDir()
	workDir := filepath.Join(incDir, name)
	os.Mkdir(workDir, 0777)
	j := &job{
		dir: workDir,
		n:   n,
	}
	j.handle()
}
<commit_msg>Reduce potential contention<commit_after>package main

import (
	"bytes"
	"gonzbee/nntp"
	"gonzbee/nzb"
	"gonzbee/yenc"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type job struct {
	dir string
	n   *nzb.Nzb
}

type messagejob struct {
	group string
	msgId string
	ch    chan []byte
}

func init() {
	go poolHandler()
}

var download = make(chan *messagejob)
var downloadMux = make(chan *messagejob)
var reaper = make(chan int)

func newConnection() error {
	s := config.Server.GetAddressStr()
	var err error
	var n *nntp.Conn
	if config.Server.TLS {
		n, err = nntp.DialTLS(s)
	} else {
		n, err = nntp.Dial(s)
	}
	if err != nil {
		return err
	}

	err = n.Authenticate(config.Server.Username, config.Server.Password)
	if err != nil {
		n.Close()
		return err
	}
	log.Println("spun up nntp connection")
	go func() {
		defer n.Close()
		for {
			select {
			case m := <-downloadMux:
				err = n.SwitchGroup(m.group)
				if err != nil {
					panic(err)
				}
				b, err := n.GetMessage(m.msgId)
				if err != nil {
					log.Print("Error getting Message ", m.msgId, ": ", err.Error())
				}
				m.ch <- b
			case <-(time.After(10 * time.Second)):
				reaper <- 1
				return
			}
		}
	}()
	return nil
}

func poolHandler() {
	var number int
	for {
		select {
		case msg := <-download:
			if number < 10 {
				err := newConnection()
				if err == nil {
					number++
				}
			}
			downloadMux <- msg
		case <-reaper:
			number--
		}
	}
}

func (j *job) handle() {
	wg := new(sync.WaitGroup)
	for _, f := range j.n.File {
		ch := make(chan []byte, 1024)
		wg.Add(1)
		partsLeft := len(f.Segments)
		go func() {
			var file *os.File
			var part *yenc.Part
			var err error
			defer wg.Done()
			for ; partsLeft > 0; partsLeft-- {
				s := <-ch
				if s == nil {
					continue
				}
				m := bytes.NewReader(s)
				part, err = yenc.NewPart(m)
				if err != nil {
					log.Print(err.Error())
					continue
				}
				if file == nil {
					file, err = os.Create(filepath.Join(j.dir, part.Filename))
					if err != nil {
						panic("could not create file: " + err.Error())
					}
					defer file.Close()
				}
				file.Seek(part.Begin, os.SEEK_SET)
				part.Decode(file)
			}
			if part != nil {
				log.Print("Done Decoding file " + part.Filename)
			} else {
				log.Print("Could not decode entire file")
			}
		}()
		for _, seg := range f.Segments {
			msg := &messagejob{
				msgId: seg.MsgId,
				group: f.Groups[0],
				ch:    ch,
			}
			download <- msg
		}
	}
	wg.Wait()
}

func jobStart(n *nzb.Nzb, name string) {
	incDir := config.GetIncompleteDir()
	workDir := filepath.Join(incDir, name)
	os.Mkdir(workDir, 0777)
	j := &job{
		dir: workDir,
		n:   n,
	}
	j.handle()
}
<|endoftext|>


<commit_before>package mithril

import "fmt"

var (
	// SelfClosingTagToken is a list of self closing elements tag
	SelfClosingTagToken = []string{"area", "base", "br", "col", "command", "embed", "hr", "img", "input", "keygen", "link", "meta", "param", "source", "track", "wbr", "!doctype"}
)

// VirtualElement is an object which can convert to real DOM element
type VirtualElement struct {
	Tag      string
	Attrs    []Attribute
	Children interface{}
}

// Attr returns or sets the attribute
func (el *VirtualElement) Attr(key string, val ...string) string {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			if len(val) > 0 {
				attr.Val(val[0])
			}
			return attr.Val()
		}
	}
	if len(val) > 0 {
		el.Attrs = append(el.Attrs, NewStringAttr(key, val[0]))
		return val[0]
	}
	return ""
}

// HasAttr returns true if VirtualElement has this attribute in attributes array
func (el *VirtualElement) HasAttr(key string) bool {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			return true
		}
	}
	return false
}

// IsSelfClosingTag returns true if Tag is a self closing tag element
func (el *VirtualElement) IsSelfClosingTag() bool {
	for _, val := range SelfClosingTagToken {
		if val == el.Tag {
			return true
		}
	}
	return false
}

// Render returns html string of anything
func (el *VirtualElement) Render(item interface{}) string {
	str := ""
	switch obj := item.(type) {
	case []interface{}:
		for _, item := range obj {
			str += el.Render(item)
		}
	case *VirtualElement:
		str += obj.String()
	case string:
		str += obj
	case int, int32, int64:
		str += fmt.Sprintf("%d", obj)
	case float32, float64:
		str += fmt.Sprintf("%f", obj)
	}
	return str
}

// String returns the html string of the element
func (el *VirtualElement) String() string {
	res := make(chan string)
	go func() {
		html := "<"
		if el.Tag != "" {
			html += el.Tag
		} else {
			html += "div"
		}
		for _, attr := range el.Attrs {
			html += (" " + attr.Key())
			if len(attr.Val()) > 0 {
				html += ("=\"" + attr.Val() + "\"")
			}
		}
		if el.IsSelfClosingTag() {
			html += " />"
		} else {
			html += ">"
			html += el.Render(el.Children)
			html += "</"
			if el.Tag != "" {
				html += el.Tag
			} else {
				html += "div"
			}
			html += ">"
		}
		res <- html
	}()
	return <-res
}
<commit_msg>remove unused code<commit_after>package mithril

var (
	// SelfClosingTagToken is a list of self closing elements tag
	SelfClosingTagToken = []string{"area", "base", "br", "col", "command", "embed", "hr", "img", "input", "keygen", "link", "meta", "param", "source", "track", "wbr", "!doctype"}
)

// VirtualElement is an object which can convert to real DOM element
type VirtualElement struct {
	Tag      string
	Attrs    []Attribute
	Children interface{}
}

// Attr returns or sets the attribute
func (el *VirtualElement) Attr(key string, val ...string) string {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			if len(val) > 0 {
				attr.Val(val[0])
			}
			return attr.Val()
		}
	}
	if len(val) > 0 {
		el.Attrs = append(el.Attrs, NewStringAttr(key, val[0]))
		return val[0]
	}
	return ""
}

// HasAttr returns true if VirtualElement has this attribute in attributes array
func (el *VirtualElement) HasAttr(key string) bool {
	for _, attr := range el.Attrs {
		if attr.Key() == key {
			return true
		}
	}
	return false
}

// IsSelfClosingTag returns true if Tag is a self closing tag element
func (el *VirtualElement) IsSelfClosingTag() bool {
	for _, val := range SelfClosingTagToken {
		if val == el.Tag {
			return true
		}
	}
	return false
}

// String returns the html string of the element
func (el *VirtualElement) String() string {
	res := make(chan string)
	go func() {
		html := "<"
		if el.Tag != "" {
			html += el.Tag
		} else {
			html += "div"
		}
		for _, attr := range el.Attrs {
			html += (" " + attr.Key())
			if len(attr.Val()) > 0 {
				html += ("=\"" + attr.Val() + "\"")
			}
		}
		if el.IsSelfClosingTag() {
			html += " />"
		} else {
			html += ">"
			html += Render(el.Children)
			html += "</"
			if el.Tag != "" {
				html += el.Tag
			} else {
				html += "div"
			}
			html += ">"
		}
		res <- html
	}()
	return <-res
}
<|endoftext|>


<commit_before>package jwp

func New(addr string) (j *Session, err error) {
	j = &Session{
		Server: addr,
	}
	err = j.Session(&ApiDesiredCapabilities{
		DesiredCapabilities: DesiredCapabilities{
			Proxy: Proxy{
				ProxyType: "dirct",
			},
		},
	})
	return
}
<commit_msg>Rename New method to Dial and refactoring.<commit_after>package jwp

func Dial(addr string) *Client {
	return &Client{
		Server: addr,
	}
}
<|endoftext|>