<commit_before>package support

import (
	"context"
	"fmt"
	"io"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
	watchtools "k8s.io/client-go/tools/watch"
)

// NewAcceptAvailablePods makes a new acceptAvailablePods from a real client.
func NewAcceptAvailablePods(
	out io.Writer,
	kclient corev1client.ReplicationControllersGetter,
	timeout time.Duration,
) *acceptAvailablePods {
	return &acceptAvailablePods{
		out:     out,
		kclient: kclient,
		timeout: timeout,
	}
}

// acceptAvailablePods will accept a replication controller if all the pods
// for the replication controller become available.
type acceptAvailablePods struct {
	out     io.Writer
	kclient corev1client.ReplicationControllersGetter
	// timeout is how long to wait for pods to become available from ready state.
	timeout time.Duration
}

// Accept all pods for a replication controller once they are available.
func (c *acceptAvailablePods) Accept(rc *corev1.ReplicationController) error {
	allReplicasAvailable := func(r *corev1.ReplicationController) bool {
		return r.Status.AvailableReplicas == *r.Spec.Replicas
	}

	if allReplicasAvailable(rc) {
		return nil
	}

	watcher, err := c.kclient.ReplicationControllers(rc.Namespace).Watch(metav1.SingleObject(metav1.ObjectMeta{Name: rc.Name, ResourceVersion: rc.ResourceVersion}))
	if err != nil {
		return fmt.Errorf("acceptAvailablePods failed to watch ReplicationController %s/%s: %v", rc.Namespace, rc.Name, err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	_, err = watchtools.UntilWithoutRetry(ctx, watcher, func(event watch.Event) (bool, error) {
		if t := event.Type; t != watch.Modified {
			return false, fmt.Errorf("acceptAvailablePods failed watching for ReplicationController %s/%s: received event %v", rc.Namespace, rc.Name, t)
		}
		newRc, ok := event.Object.(*corev1.ReplicationController)
		if !ok {
			return false, fmt.Errorf("unknown event object %#v", event.Object)
		}
		return allReplicasAvailable(newRc), nil
	})
	// Handle acceptance failure.
	if err != nil {
		if err == wait.ErrWaitTimeout {
			return fmt.Errorf("pods for rc '%s/%s' took longer than %.f seconds to become available", rc.Namespace, rc.Name, c.timeout.Seconds())
		}
		return err
	}
	return nil
}
<commit_msg>Fix acceptor WATCH<commit_after>package support

import (
	"context"
	"fmt"
	"io"
	"time"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/fields"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/apimachinery/pkg/watch"
	corev1client "k8s.io/client-go/kubernetes/typed/core/v1"
	"k8s.io/client-go/tools/cache"
	watchtools "k8s.io/client-go/tools/watch"
)

// NewAcceptAvailablePods makes a new acceptAvailablePods from a real client.
func NewAcceptAvailablePods(
	out io.Writer,
	kclient corev1client.ReplicationControllersGetter,
	timeout time.Duration,
) *acceptAvailablePods {
	return &acceptAvailablePods{
		out:     out,
		kclient: kclient,
		timeout: timeout,
	}
}

// acceptAvailablePods will accept a replication controller if all the pods
// for the replication controller become available.
type acceptAvailablePods struct {
	out     io.Writer
	kclient corev1client.ReplicationControllersGetter
	// timeout is how long to wait for pods to become available from ready state.
	timeout time.Duration
}

// Accept all pods for a replication controller once they are available.
func (c *acceptAvailablePods) Accept(rc *corev1.ReplicationController) error {
	allReplicasAvailable := func(r *corev1.ReplicationController) bool {
		return r.Status.AvailableReplicas == *r.Spec.Replicas
	}

	if allReplicasAvailable(rc) {
		return nil
	}

	fieldSelector := fields.OneTermEqualSelector("metadata.name", rc.Name).String()
	lw := &cache.ListWatch{
		ListFunc: func(options metav1.ListOptions) (runtime.Object, error) {
			options.FieldSelector = fieldSelector
			return c.kclient.ReplicationControllers(rc.Namespace).List(options)
		},
		WatchFunc: func(options metav1.ListOptions) (watch.Interface, error) {
			options.FieldSelector = fieldSelector
			return c.kclient.ReplicationControllers(rc.Namespace).Watch(options)
		},
	}

	preconditionFunc := func(store cache.Store) (bool, error) {
		item, exists, err := store.Get(&metav1.ObjectMeta{Namespace: rc.Namespace, Name: rc.Name})
		if err != nil {
			return true, err
		}
		if !exists {
			// We need to make sure we see the object in the cache before we start waiting for events
			// or we would be waiting for the timeout if such object didn't exist.
			return true, fmt.Errorf("%s '%s/%s' not found", corev1.Resource("replicationcontrollers"), rc.Namespace, rc.Name)
		}

		// Check that the objects UID match for cases of recreation
		storeRc, ok := item.(*corev1.ReplicationController)
		if !ok {
			return true, fmt.Errorf("unexpected store item type: %#v", item)
		}
		if rc.UID != storeRc.UID {
			return true, fmt.Errorf("%s '%s/%s' no longer exists, expected UID %q, got UID %q", corev1.Resource("replicationcontrollers"), rc.Namespace, rc.Name, rc.UID, storeRc.UID)
		}

		return false, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
	defer cancel()
	_, err := watchtools.UntilWithSync(ctx, lw, &corev1.ReplicationController{}, preconditionFunc, func(event watch.Event) (bool, error) {
		switch event.Type {
		case watch.Added, watch.Modified:
			newRc, ok := event.Object.(*corev1.ReplicationController)
			if !ok {
				return true, fmt.Errorf("unknown event object %#v", event.Object)
			}
			return allReplicasAvailable(newRc), nil

		case watch.Deleted:
			return true, fmt.Errorf("replicationController got deleted %#v", event.Object)

		case watch.Error:
			return true, fmt.Errorf("unexpected error %#v", event.Object)

		default:
			return true, fmt.Errorf("unexpected event type: %T", event.Type)
		}
	})
	// Handle acceptance failure.
	if err == wait.ErrWaitTimeout {
		return fmt.Errorf("pods for rc '%s/%s' took longer than %.f seconds to become available", rc.Namespace, rc.Name, c.timeout.Seconds())
	}

	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package sync

import (
	"context"
	"fmt"
	"net"
	"time"

	"github.com/golang/glog"

	"k8s.io/api/core/v1"
	"k8s.io/kubernetes/pkg/controller/node/ipam/cidrset"
)

const (
	// InvalidPodCIDR is the event recorded when a node is found with an
	// invalid PodCIDR.
	InvalidPodCIDR = "CloudCIDRAllocatorInvalidPodCIDR"
	// InvalidModeEvent is the event recorded when the CIDR range cannot be
	// sync'd due to the cluster running in the wrong mode.
	InvalidModeEvent = "CloudCIDRAllocatorInvalidMode"
	// MismatchEvent is the event recorded when the CIDR range allocated in the
	// node spec does not match what has been allocated in the cloud.
	MismatchEvent = "CloudCIDRAllocatorMismatch"
)

// cloudAlias is the interface to the cloud platform APIs.
type cloudAlias interface {
	// Alias returns the IP alias for the node.
	Alias(ctx context.Context, nodeName string) (*net.IPNet, error)
	// AddAlias adds an alias to the node.
	AddAlias(ctx context.Context, nodeName string, cidrRange *net.IPNet) error
}

// kubeAPI is the interface to the Kubernetes APIs.
type kubeAPI interface {
	// Node returns the spec for the Node object.
	Node(ctx context.Context, name string) (*v1.Node, error)
	// UpdateNodePodCIDR updates the PodCIDR in the Node spec.
	UpdateNodePodCIDR(ctx context.Context, node *v1.Node, cidrRange *net.IPNet) error
	// UpdateNodeNetworkUnavailable updates the network unavailable status for the node.
	UpdateNodeNetworkUnavailable(nodeName string, unavailable bool) error
	// EmitNodeWarningEvent emits an event for the given node.
	EmitNodeWarningEvent(nodeName, reason, fmt string, args ...interface{})
}

// controller is the interface to the controller.
type controller interface {
	// ReportResult updates the controller with the result of the latest
	// sync operation.
	ReportResult(err error)
	// ResyncTimeout returns the amount of time to wait before retrying
	// a sync with a node.
	ResyncTimeout() time.Duration
}

// NodeSyncMode is the mode the cloud CIDR allocator runs in.
type NodeSyncMode string

var (
	// SyncFromCloud is the mode that synchronizes the IP allocation from the cloud
	// platform to the node.
	SyncFromCloud NodeSyncMode = "SyncFromCloud"
	// SyncFromCluster is the mode that synchronizes the IP allocation determined
	// by the k8s controller to the cloud provider.
	SyncFromCluster NodeSyncMode = "SyncFromCluster"
)

// IsValidMode returns true if the given mode is valid.
func IsValidMode(m NodeSyncMode) bool {
	switch m {
	case SyncFromCloud:
	case SyncFromCluster:
	default:
		return false
	}
	return true
}

// NodeSync synchronizes the state for a single node in the cluster.
type NodeSync struct {
	c          controller
	cloudAlias cloudAlias
	kubeAPI    kubeAPI
	mode       NodeSyncMode
	nodeName   string
	opChan     chan syncOp
	set        *cidrset.CidrSet
}

// New returns a new syncer for a given node.
func New(c controller, cloudAlias cloudAlias, kubeAPI kubeAPI, mode NodeSyncMode, nodeName string, set *cidrset.CidrSet) *NodeSync {
	return &NodeSync{
		c:          c,
		cloudAlias: cloudAlias,
		kubeAPI:    kubeAPI,
		mode:       mode,
		nodeName:   nodeName,
		opChan:     make(chan syncOp, 1),
		set:        set,
	}
}

// Loop runs the sync loop for a given node. done is an optional channel that
// is closed when the Loop() returns.
func (sync *NodeSync) Loop(done chan struct{}) {
	glog.V(2).Infof("Starting sync loop for node %q", sync.nodeName)

	defer func() {
		if done != nil {
			close(done)
		}
	}()

	timeout := sync.c.ResyncTimeout()
	delayTimer := time.NewTimer(timeout)
	glog.V(4).Infof("Resync node %q in %v", sync.nodeName, timeout)

	for {
		select {
		case op, more := <-sync.opChan:
			if !more {
				glog.V(2).Infof("Stopping sync loop")
				return
			}
			sync.c.ReportResult(op.run(sync))
			if !delayTimer.Stop() {
				<-delayTimer.C
			}
		case <-delayTimer.C:
			glog.V(4).Infof("Running resync for node %q", sync.nodeName)
			sync.c.ReportResult((&updateOp{}).run(sync))
		}

		timeout := sync.c.ResyncTimeout()
		delayTimer.Reset(timeout)
		glog.V(4).Infof("Resync node %q in %v", sync.nodeName, timeout)
	}
}

// Update causes an update operation on the given node. If node is nil, then
// the syncer will fetch the node spec from the API server before syncing.
//
// This method is safe to call from multiple goroutines.
func (sync *NodeSync) Update(node *v1.Node) {
	sync.opChan <- &updateOp{node}
}

// Delete performs the sync operations necessary to remove the node from the
// IPAM state.
//
// This method is safe to call from multiple goroutines.
func (sync *NodeSync) Delete(node *v1.Node) {
	sync.opChan <- &deleteOp{node}
	close(sync.opChan)
}

// syncOp is the interface for generic sync operation.
type syncOp interface {
	// Run the requested sync operation.
	run(sync *NodeSync) error
}

// updateOp handles creation and updates of a node.
type updateOp struct {
	node *v1.Node
}

func (op *updateOp) String() string {
	if op.node == nil {
		return fmt.Sprintf("updateOp(nil)")
	}
	return fmt.Sprintf("updateOp(%q,%v)", op.node.Name, op.node.Spec.PodCIDR)
}

func (op *updateOp) run(sync *NodeSync) error {
	glog.V(3).Infof("Running updateOp %+v", op)

	ctx := context.Background()

	if op.node == nil {
		glog.V(3).Infof("Getting node spec for %q", sync.nodeName)
		node, err := sync.kubeAPI.Node(ctx, sync.nodeName)
		if err != nil {
			glog.Errorf("Error getting node %q spec: %v", sync.nodeName, err)
			return err
		}
		op.node = node
	}

	aliasRange, err := sync.cloudAlias.Alias(ctx, sync.nodeName)
	if err != nil {
		glog.Errorf("Error getting cloud alias for node %q: %v", sync.nodeName, err)
		return err
	}

	switch {
	case op.node.Spec.PodCIDR == "" && aliasRange == nil:
		err = op.allocateRange(ctx, sync, op.node)
	case op.node.Spec.PodCIDR == "" && aliasRange != nil:
		err = op.updateNodeFromAlias(ctx, sync, op.node, aliasRange)
	case op.node.Spec.PodCIDR != "" && aliasRange == nil:
		err = op.updateAliasFromNode(ctx, sync, op.node)
	case op.node.Spec.PodCIDR != "" && aliasRange != nil:
		err = op.validateRange(ctx, sync, op.node, aliasRange)
	}

	return err
}

// validateRange checks that the allocated range and the alias range
// match.
func (op *updateOp) validateRange(ctx context.Context, sync *NodeSync, node *v1.Node, aliasRange *net.IPNet) error {
	if node.Spec.PodCIDR != aliasRange.String() {
		glog.Errorf("Inconsistency detected between node PodCIDR and node alias (%v != %v)",
			node.Spec.PodCIDR, aliasRange)
		sync.kubeAPI.EmitNodeWarningEvent(node.Name, MismatchEvent,
			"Node.Spec.PodCIDR != cloud alias (%v != %v)", node.Spec.PodCIDR, aliasRange)
		// User intervention is required in this case, as this is most likely due
		// to the user mucking around with their VM aliases on the side.
	} else {
		glog.V(4).Infof("Node %q CIDR range %v is matches cloud assignment", node.Name, node.Spec.PodCIDR)
	}
	return nil
}

// updateNodeFromAlias updates the the node from the cloud allocated
// alias.
func (op *updateOp) updateNodeFromAlias(ctx context.Context, sync *NodeSync, node *v1.Node, aliasRange *net.IPNet) error {
	if sync.mode != SyncFromCloud {
		sync.kubeAPI.EmitNodeWarningEvent(node.Name, InvalidModeEvent,
			"Cannot sync from cloud in mode %q", sync.mode)
		return fmt.Errorf("cannot sync from cloud in mode %q", sync.mode)
	}

	glog.V(2).Infof("Updating node spec with alias range, node.PodCIDR = %v", aliasRange)

	if err := sync.set.Occupy(aliasRange); err != nil {
		glog.Errorf("Error occupying range %v for node %v", aliasRange, sync.nodeName)
		return err
	}

	if err := sync.kubeAPI.UpdateNodePodCIDR(ctx, node, aliasRange); err != nil {
		glog.Errorf("Could not update node %q PodCIDR to %v: %v", node.Name, aliasRange, err)
		return err
	}

	glog.V(2).Infof("Node %q PodCIDR set to %v", node.Name, aliasRange)

	if err := sync.kubeAPI.UpdateNodeNetworkUnavailable(node.Name, false); err != nil {
		glog.Errorf("Could not update node NetworkUnavailable status to false: %v", err)
		return err
	}

	glog.V(2).Infof("Updated node %q PodCIDR from cloud alias %v", node.Name, aliasRange)

	return nil
}

// updateAliasFromNode updates the cloud alias given the node allocation.
func (op *updateOp) updateAliasFromNode(ctx context.Context, sync *NodeSync, node *v1.Node) error {
	if sync.mode != SyncFromCluster {
		sync.kubeAPI.EmitNodeWarningEvent(
			node.Name, InvalidModeEvent, "Cannot sync to cloud in mode %q", sync.mode)
		return fmt.Errorf("cannot sync to cloud in mode %q", sync.mode)
	}

	_, aliasRange, err := net.ParseCIDR(node.Spec.PodCIDR)
	if err != nil {
		glog.Errorf("Could not parse PodCIDR (%q) for node %q: %v",
			node.Spec.PodCIDR, node.Name, err)
		return err
	}

	if err := sync.set.Occupy(aliasRange); err != nil {
		glog.Errorf("Error occupying range %v for node %v", aliasRange, sync.nodeName)
		return err
	}

	if err := sync.cloudAlias.AddAlias(ctx, node.Name, aliasRange); err != nil {
		glog.Errorf("Could not add alias %v for node %q: %v", aliasRange, node.Name, err)
		return err
	}

	if err := sync.kubeAPI.UpdateNodeNetworkUnavailable(node.Name, false); err != nil {
		glog.Errorf("Could not update node NetworkUnavailable status to false: %v", err)
		return err
	}

	glog.V(2).Infof("Updated node %q cloud alias with node spec, node.PodCIDR = %v",
		node.Name, node.Spec.PodCIDR)

	return nil
}

// allocateRange allocates a new range and updates both the cloud
// platform and the node allocation.
func (op *updateOp) allocateRange(ctx context.Context, sync *NodeSync, node *v1.Node) error {
	if sync.mode != SyncFromCluster {
		sync.kubeAPI.EmitNodeWarningEvent(node.Name, InvalidModeEvent,
			"Cannot allocate CIDRs in mode %q", sync.mode)
		return fmt.Errorf("controller cannot allocate CIDRS in mode %q", sync.mode)
	}

	cidrRange, err := sync.set.AllocateNext()
	if err != nil {
		return err
	}
	// If addAlias returns a hard error, cidrRange will be leaked as there
	// is no durable record of the range. The missing space will be
	// recovered on the next restart of the controller.
	if err := sync.cloudAlias.AddAlias(ctx, node.Name, cidrRange); err != nil {
		glog.Errorf("Could not add alias %v for node %q: %v", cidrRange, node.Name, err)
		return err
	}

	if err := sync.kubeAPI.UpdateNodePodCIDR(ctx, node, cidrRange); err != nil {
		glog.Errorf("Could not update node %q PodCIDR to %v: %v", node.Name, cidrRange, err)
		return err
	}

	if err := sync.kubeAPI.UpdateNodeNetworkUnavailable(node.Name, false); err != nil {
		glog.Errorf("Could not update node NetworkUnavailable status to false: %v", err)
		return err
	}

	glog.V(2).Infof("Allocated PodCIDR %v for node %q", cidrRange, node.Name)

	return nil
}

// deleteOp handles deletion of a node.
type deleteOp struct {
	node *v1.Node
}

func (op *deleteOp) String() string {
	if op.node == nil {
		return fmt.Sprintf("deleteOp(nil)")
	}
	return fmt.Sprintf("deleteOp(%q,%v)", op.node.Name, op.node.Spec.PodCIDR)
}

func (op *deleteOp) run(sync *NodeSync) error {
	glog.V(3).Infof("Running deleteOp %+v", op)
	if op.node.Spec.PodCIDR == "" {
		glog.V(2).Infof("Node %q was deleted, node had no PodCIDR range assigned", op.node.Name)
		return nil
	}

	_, cidrRange, err := net.ParseCIDR(op.node.Spec.PodCIDR)
	if err != nil {
		glog.Errorf("Deleted node %q has an invalid podCIDR %q: %v",
			op.node.Name, op.node.Spec.PodCIDR, err)
		sync.kubeAPI.EmitNodeWarningEvent(op.node.Name, InvalidPodCIDR,
			"Node %q has an invalid PodCIDR: %q", op.node.Name, op.node.Spec.PodCIDR)
		return nil
	}

	sync.set.Release(cidrRange)
	glog.V(2).Infof("Node %q was deleted, releasing CIDR range %v",
		op.node.Name, op.node.Spec.PodCIDR)

	return nil
}
<commit_msg>revert the comment change from Run to run<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package sync

import (
	"context"
	"fmt"
	"net"
	"time"

	"github.com/golang/glog"

	"k8s.io/api/core/v1"
	"k8s.io/kubernetes/pkg/controller/node/ipam/cidrset"
)

const (
	// InvalidPodCIDR is the event recorded when a node is found with an
	// invalid PodCIDR.
	InvalidPodCIDR = "CloudCIDRAllocatorInvalidPodCIDR"
	// InvalidModeEvent is the event recorded when the CIDR range cannot be
	// sync'd due to the cluster running in the wrong mode.
	InvalidModeEvent = "CloudCIDRAllocatorInvalidMode"
	// MismatchEvent is the event recorded when the CIDR range allocated in the
	// node spec does not match what has been allocated in the cloud.
	MismatchEvent = "CloudCIDRAllocatorMismatch"
)

// cloudAlias is the interface to the cloud platform APIs.
type cloudAlias interface {
	// Alias returns the IP alias for the node.
	Alias(ctx context.Context, nodeName string) (*net.IPNet, error)
	// AddAlias adds an alias to the node.
	AddAlias(ctx context.Context, nodeName string, cidrRange *net.IPNet) error
}

// kubeAPI is the interface to the Kubernetes APIs.
type kubeAPI interface {
	// Node returns the spec for the Node object.
	Node(ctx context.Context, name string) (*v1.Node, error)
	// UpdateNodePodCIDR updates the PodCIDR in the Node spec.
	UpdateNodePodCIDR(ctx context.Context, node *v1.Node, cidrRange *net.IPNet) error
	// UpdateNodeNetworkUnavailable updates the network unavailable status for the node.
	UpdateNodeNetworkUnavailable(nodeName string, unavailable bool) error
	// EmitNodeWarningEvent emits an event for the given node.
	EmitNodeWarningEvent(nodeName, reason, fmt string, args ...interface{})
}

// controller is the interface to the controller.
type controller interface {
	// ReportResult updates the controller with the result of the latest
	// sync operation.
	ReportResult(err error)
	// ResyncTimeout returns the amount of time to wait before retrying
	// a sync with a node.
	ResyncTimeout() time.Duration
}

// NodeSyncMode is the mode the cloud CIDR allocator runs in.
type NodeSyncMode string

var (
	// SyncFromCloud is the mode that synchronizes the IP allocation from the cloud
	// platform to the node.
	SyncFromCloud NodeSyncMode = "SyncFromCloud"
	// SyncFromCluster is the mode that synchronizes the IP allocation determined
	// by the k8s controller to the cloud provider.
	SyncFromCluster NodeSyncMode = "SyncFromCluster"
)

// IsValidMode returns true if the given mode is valid.
func IsValidMode(m NodeSyncMode) bool {
	switch m {
	case SyncFromCloud:
	case SyncFromCluster:
	default:
		return false
	}
	return true
}

// NodeSync synchronizes the state for a single node in the cluster.
type NodeSync struct {
	c          controller
	cloudAlias cloudAlias
	kubeAPI    kubeAPI
	mode       NodeSyncMode
	nodeName   string
	opChan     chan syncOp
	set        *cidrset.CidrSet
}

// New returns a new syncer for a given node.
func New(c controller, cloudAlias cloudAlias, kubeAPI kubeAPI, mode NodeSyncMode, nodeName string, set *cidrset.CidrSet) *NodeSync {
	return &NodeSync{
		c:          c,
		cloudAlias: cloudAlias,
		kubeAPI:    kubeAPI,
		mode:       mode,
		nodeName:   nodeName,
		opChan:     make(chan syncOp, 1),
		set:        set,
	}
}

// Loop runs the sync loop for a given node. done is an optional channel that
// is closed when the Loop() returns.
func (sync *NodeSync) Loop(done chan struct{}) {
	glog.V(2).Infof("Starting sync loop for node %q", sync.nodeName)

	defer func() {
		if done != nil {
			close(done)
		}
	}()

	timeout := sync.c.ResyncTimeout()
	delayTimer := time.NewTimer(timeout)
	glog.V(4).Infof("Resync node %q in %v", sync.nodeName, timeout)

	for {
		select {
		case op, more := <-sync.opChan:
			if !more {
				glog.V(2).Infof("Stopping sync loop")
				return
			}
			sync.c.ReportResult(op.run(sync))
			if !delayTimer.Stop() {
				<-delayTimer.C
			}
		case <-delayTimer.C:
			glog.V(4).Infof("Running resync for node %q", sync.nodeName)
			sync.c.ReportResult((&updateOp{}).run(sync))
		}

		timeout := sync.c.ResyncTimeout()
		delayTimer.Reset(timeout)
		glog.V(4).Infof("Resync node %q in %v", sync.nodeName, timeout)
	}
}

// Update causes an update operation on the given node. If node is nil, then
// the syncer will fetch the node spec from the API server before syncing.
//
// This method is safe to call from multiple goroutines.
func (sync *NodeSync) Update(node *v1.Node) {
	sync.opChan <- &updateOp{node}
}

// Delete performs the sync operations necessary to remove the node from the
// IPAM state.
//
// This method is safe to call from multiple goroutines.
func (sync *NodeSync) Delete(node *v1.Node) {
	sync.opChan <- &deleteOp{node}
	close(sync.opChan)
}

// syncOp is the interface for generic sync operation.
type syncOp interface {
	// run the requested sync operation.
	run(sync *NodeSync) error
}

// updateOp handles creation and updates of a node.
type updateOp struct {
	node *v1.Node
}

func (op *updateOp) String() string {
	if op.node == nil {
		return fmt.Sprintf("updateOp(nil)")
	}
	return fmt.Sprintf("updateOp(%q,%v)", op.node.Name, op.node.Spec.PodCIDR)
}

func (op *updateOp) run(sync *NodeSync) error {
	glog.V(3).Infof("Running updateOp %+v", op)

	ctx := context.Background()

	if op.node == nil {
		glog.V(3).Infof("Getting node spec for %q", sync.nodeName)
		node, err := sync.kubeAPI.Node(ctx, sync.nodeName)
		if err != nil {
			glog.Errorf("Error getting node %q spec: %v", sync.nodeName, err)
			return err
		}
		op.node = node
	}

	aliasRange, err := sync.cloudAlias.Alias(ctx, sync.nodeName)
	if err != nil {
		glog.Errorf("Error getting cloud alias for node %q: %v", sync.nodeName, err)
		return err
	}

	switch {
	case op.node.Spec.PodCIDR == "" && aliasRange == nil:
		err = op.allocateRange(ctx, sync, op.node)
	case op.node.Spec.PodCIDR == "" && aliasRange != nil:
		err = op.updateNodeFromAlias(ctx, sync, op.node, aliasRange)
	case op.node.Spec.PodCIDR != "" && aliasRange == nil:
		err = op.updateAliasFromNode(ctx, sync, op.node)
	case op.node.Spec.PodCIDR != "" && aliasRange != nil:
		err = op.validateRange(ctx, sync, op.node, aliasRange)
	}

	return err
}

// validateRange checks that the allocated range and the alias range
// match.
func (op *updateOp) validateRange(ctx context.Context, sync *NodeSync, node *v1.Node, aliasRange *net.IPNet) error {
	if node.Spec.PodCIDR != aliasRange.String() {
		glog.Errorf("Inconsistency detected between node PodCIDR and node alias (%v != %v)",
			node.Spec.PodCIDR, aliasRange)
		sync.kubeAPI.EmitNodeWarningEvent(node.Name, MismatchEvent,
			"Node.Spec.PodCIDR != cloud alias (%v != %v)", node.Spec.PodCIDR, aliasRange)
		// User intervention is required in this case, as this is most likely due
		// to the user mucking around with their VM aliases on the side.
	} else {
		glog.V(4).Infof("Node %q CIDR range %v is matches cloud assignment", node.Name, node.Spec.PodCIDR)
	}
	return nil
}

// updateNodeFromAlias updates the the node from the cloud allocated
// alias.
func (op *updateOp) updateNodeFromAlias(ctx context.Context, sync *NodeSync, node *v1.Node, aliasRange *net.IPNet) error {
	if sync.mode != SyncFromCloud {
		sync.kubeAPI.EmitNodeWarningEvent(node.Name, InvalidModeEvent,
			"Cannot sync from cloud in mode %q", sync.mode)
		return fmt.Errorf("cannot sync from cloud in mode %q", sync.mode)
	}

	glog.V(2).Infof("Updating node spec with alias range, node.PodCIDR = %v", aliasRange)

	if err := sync.set.Occupy(aliasRange); err != nil {
		glog.Errorf("Error occupying range %v for node %v", aliasRange, sync.nodeName)
		return err
	}

	if err := sync.kubeAPI.UpdateNodePodCIDR(ctx, node, aliasRange); err != nil {
		glog.Errorf("Could not update node %q PodCIDR to %v: %v", node.Name, aliasRange, err)
		return err
	}

	glog.V(2).Infof("Node %q PodCIDR set to %v", node.Name, aliasRange)

	if err := sync.kubeAPI.UpdateNodeNetworkUnavailable(node.Name, false); err != nil {
		glog.Errorf("Could not update node NetworkUnavailable status to false: %v", err)
		return err
	}

	glog.V(2).Infof("Updated node %q PodCIDR from cloud alias %v", node.Name, aliasRange)

	return nil
}

// updateAliasFromNode updates the cloud alias given the node allocation.
func (op *updateOp) updateAliasFromNode(ctx context.Context, sync *NodeSync, node *v1.Node) error {
	if sync.mode != SyncFromCluster {
		sync.kubeAPI.EmitNodeWarningEvent(
			node.Name, InvalidModeEvent, "Cannot sync to cloud in mode %q", sync.mode)
		return fmt.Errorf("cannot sync to cloud in mode %q", sync.mode)
	}

	_, aliasRange, err := net.ParseCIDR(node.Spec.PodCIDR)
	if err != nil {
		glog.Errorf("Could not parse PodCIDR (%q) for node %q: %v",
			node.Spec.PodCIDR, node.Name, err)
		return err
	}

	if err := sync.set.Occupy(aliasRange); err != nil {
		glog.Errorf("Error occupying range %v for node %v", aliasRange, sync.nodeName)
		return err
	}

	if err := sync.cloudAlias.AddAlias(ctx, node.Name, aliasRange); err != nil {
		glog.Errorf("Could not add alias %v for node %q: %v", aliasRange, node.Name, err)
		return err
	}

	if err := sync.kubeAPI.UpdateNodeNetworkUnavailable(node.Name, false); err != nil {
		glog.Errorf("Could not update node NetworkUnavailable status to false: %v", err)
		return err
	}

	glog.V(2).Infof("Updated node %q cloud alias with node spec, node.PodCIDR = %v",
		node.Name, node.Spec.PodCIDR)

	return nil
}

// allocateRange allocates a new range and updates both the cloud
// platform and the node allocation.
func (op *updateOp) allocateRange(ctx context.Context, sync *NodeSync, node *v1.Node) error {
	if sync.mode != SyncFromCluster {
		sync.kubeAPI.EmitNodeWarningEvent(node.Name, InvalidModeEvent,
			"Cannot allocate CIDRs in mode %q", sync.mode)
		return fmt.Errorf("controller cannot allocate CIDRS in mode %q", sync.mode)
	}

	cidrRange, err := sync.set.AllocateNext()
	if err != nil {
		return err
	}
	// If addAlias returns a hard error, cidrRange will be leaked as there
	// is no durable record of the range. The missing space will be
	// recovered on the next restart of the controller.
	if err := sync.cloudAlias.AddAlias(ctx, node.Name, cidrRange); err != nil {
		glog.Errorf("Could not add alias %v for node %q: %v", cidrRange, node.Name, err)
		return err
	}

	if err := sync.kubeAPI.UpdateNodePodCIDR(ctx, node, cidrRange); err != nil {
		glog.Errorf("Could not update node %q PodCIDR to %v: %v", node.Name, cidrRange, err)
		return err
	}

	if err := sync.kubeAPI.UpdateNodeNetworkUnavailable(node.Name, false); err != nil {
		glog.Errorf("Could not update node NetworkUnavailable status to false: %v", err)
		return err
	}

	glog.V(2).Infof("Allocated PodCIDR %v for node %q", cidrRange, node.Name)

	return nil
}

// deleteOp handles deletion of a node.
type deleteOp struct {
	node *v1.Node
}

func (op *deleteOp) String() string {
	if op.node == nil {
		return fmt.Sprintf("deleteOp(nil)")
	}
	return fmt.Sprintf("deleteOp(%q,%v)", op.node.Name, op.node.Spec.PodCIDR)
}

func (op *deleteOp) run(sync *NodeSync) error {
	glog.V(3).Infof("Running deleteOp %+v", op)
	if op.node.Spec.PodCIDR == "" {
		glog.V(2).Infof("Node %q was deleted, node had no PodCIDR range assigned", op.node.Name)
		return nil
	}

	_, cidrRange, err := net.ParseCIDR(op.node.Spec.PodCIDR)
	if err != nil {
		glog.Errorf("Deleted node %q has an invalid podCIDR %q: %v",
			op.node.Name, op.node.Spec.PodCIDR, err)
		sync.kubeAPI.EmitNodeWarningEvent(op.node.Name, InvalidPodCIDR,
			"Node %q has an invalid PodCIDR: %q", op.node.Name, op.node.Spec.PodCIDR)
		return nil
	}

	sync.set.Release(cidrRange)
	glog.V(2).Infof("Node %q was deleted, releasing CIDR range %v",
		op.node.Name, op.node.Spec.PodCIDR)

	return nil
}
<|endoftext|>


<commit_before>package image

import (
	"context"
	"errors"
	"fmt"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/sync/singleflight"

	"github.com/grafana/grafana/pkg/components/imguploader"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/dashboards"
	"github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/services/ngalert/store"
	"github.com/grafana/grafana/pkg/services/rendering"
	"github.com/grafana/grafana/pkg/services/screenshot"
	"github.com/grafana/grafana/pkg/setting"
)

const (
	screenshotTimeout  = 10 * time.Second
	screenshotCacheTTL = 60 * time.Second
)

var (
	// ErrNoDashboard is returned when the alert rule does not have a Dashboard UID
	// in its annotations or the dashboard does not exist.
	ErrNoDashboard = errors.New("no dashboard")

	// ErrNoPanel is returned when the alert rule does not have a PanelID in its
	// annotations.
	ErrNoPanel = errors.New("no panel")
)

// DeleteExpiredService is a service to delete expired images.
type DeleteExpiredService struct {
	store store.ImageAdminStore
}

func (s *DeleteExpiredService) DeleteExpired(ctx context.Context) (int64, error) {
	return s.store.DeleteExpiredImages(ctx)
}

func ProvideDeleteExpiredService(store *store.DBstore) *DeleteExpiredService {
	return &DeleteExpiredService{store: store}
}

//go:generate mockgen -destination=mock.go -package=image github.com/grafana/grafana/pkg/services/ngalert/image ImageService
type ImageService interface {
	// NewImage returns a new image for the alert instance.
	NewImage(ctx context.Context, r *models.AlertRule) (*models.Image, error)
}

// ScreenshotImageService takes screenshots of the alert rule and saves the
// image in the store. The image contains a unique token that can be passed
// as an annotation or label to the Alertmanager. This service cannot take
// screenshots of alert rules that are not associated with a dashboard panel.
type ScreenshotImageService struct {
	limiter      screenshot.RateLimiter
	logger       log.Logger
	screenshots  screenshot.ScreenshotService
	singleflight singleflight.Group
	store        store.ImageStore
	uploads      *UploadingService
}

// NewScreenshotImageService returns a new ScreenshotImageService.
func NewScreenshotImageService(
	limiter screenshot.RateLimiter,
	logger log.Logger,
	screenshots screenshot.ScreenshotService,
	store store.ImageStore,
	uploads *UploadingService) ImageService {
	return &ScreenshotImageService{
		limiter:     limiter,
		logger:      logger,
		screenshots: screenshots,
		store:       store,
		uploads:     uploads,
	}
}

// NewScreenshotImageServiceFromCfg returns a new ScreenshotImageService
// from the configuration.
func NewScreenshotImageServiceFromCfg(cfg *setting.Cfg, db *store.DBstore, ds dashboards.DashboardService,
	rs rendering.Service, r prometheus.Registerer) (ImageService, error) {
	var (
		limiter     screenshot.RateLimiter       = &screenshot.NoOpRateLimiter{}
		screenshots screenshot.ScreenshotService = &screenshot.ScreenshotUnavailableService{}
		uploads     *UploadingService            = nil
	)

	// If screenshots are enabled
	if cfg.UnifiedAlerting.Screenshots.Capture {
		limiter = screenshot.NewTokenRateLimiter(cfg.UnifiedAlerting.Screenshots.MaxConcurrentScreenshots)
		screenshots = screenshot.NewHeadlessScreenshotService(ds, rs, r)

		// Image uploading is an optional feature
		if cfg.UnifiedAlerting.Screenshots.UploadExternalImageStorage {
			m, err := imguploader.NewImageUploader()
			if err != nil {
				return nil, fmt.Errorf("failed to initialize uploading screenshot service: %w", err)
			}
			uploads = NewUploadingService(m, r)
		}
	}

	return NewScreenshotImageService(limiter, cfg.Logger, screenshots, db, uploads), nil
}

// NewImage returns a screenshot of the alert rule or an error.
//
// The alert rule must be associated with a dashboard panel for a screenshot to be
// taken. If the alert rule does not have a Dashboard UID in its annotations,
// or the dashboard does not exist, an ErrNoDashboard error is returned. If the
// alert rule has a Dashboard UID and the dashboard exists, but does not have a
// Panel ID in its annotations then an ErrNoPanel error is returned.
func (s *ScreenshotImageService) NewImage(ctx context.Context, r *models.AlertRule) (*models.Image, error) {
	if r.DashboardUID == nil {
		return nil, ErrNoDashboard
	}

	if r.PanelID == nil || *r.PanelID == 0 {
		return nil, ErrNoPanel
	}

	ctx, cancelFunc := context.WithTimeout(ctx, screenshotTimeout)
	defer cancelFunc()

	opts := screenshot.ScreenshotOptions{
		DashboardUID: *r.DashboardUID,
		PanelID:      *r.PanelID,
		Timeout:      screenshotTimeout,
	}

	k := fmt.Sprintf("%s-%d-%s", opts.DashboardUID, opts.PanelID, opts.Theme)
	result, err, _ := s.singleflight.Do(k, func() (interface{}, error) {
		screenshot, err := s.limiter.Do(ctx, opts, s.screenshots.Take)
		if err != nil {
			if errors.Is(err, dashboards.ErrDashboardNotFound) {
				return nil, ErrNoDashboard
			}
			return nil, err
		}
		image := models.Image{Path: screenshot.Path}
		if s.uploads != nil {
			if image, err = s.uploads.Upload(ctx, image); err != nil {
				s.logger.Warn("failed to upload image", "path", image.Path, "error", err)
			}
		}
		if err := s.store.SaveImage(ctx, &image); err != nil {
			return nil, fmt.Errorf("failed to save image: %w", err)
		}
		return image, nil
	})
	if err != nil {
		return nil, err
	}

	image := result.(models.Image)
	return &image, nil
}

// NotAvailableImageService is a service that returns ErrScreenshotsUnavailable.
type NotAvailableImageService struct{}

func (s *NotAvailableImageService) NewImage(_ context.Context, _ *models.AlertRule) (*models.Image, error) {
	return nil, screenshot.ErrScreenshotsUnavailable
}

// NoopImageService is a no-op image service.
type NoopImageService struct{}

func (s *NoopImageService) NewImage(_ context.Context, _ *models.AlertRule) (*models.Image, error) {
	return &models.Image{}, nil
}
<commit_msg>Alerting: Use hash of opts in singleflight (#58474)<commit_after>package image

import (
	"context"
	"encoding/base64"
	"errors"
	"fmt"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/sync/singleflight"

	"github.com/grafana/grafana/pkg/components/imguploader"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/services/dashboards"
	"github.com/grafana/grafana/pkg/services/ngalert/models"
	"github.com/grafana/grafana/pkg/services/ngalert/store"
	"github.com/grafana/grafana/pkg/services/rendering"
	"github.com/grafana/grafana/pkg/services/screenshot"
	"github.com/grafana/grafana/pkg/setting"
)

const (
	screenshotTimeout  = 10 * time.Second
	screenshotCacheTTL = 60 * time.Second
)

var (
	// ErrNoDashboard is returned when the alert rule does not have a Dashboard UID
	// in its annotations or the dashboard does not exist.
	ErrNoDashboard = errors.New("no dashboard")

	// ErrNoPanel is returned when the alert rule does not have a PanelID in its
	// annotations.
	ErrNoPanel = errors.New("no panel")
)

// DeleteExpiredService is a service to delete expired images.
type DeleteExpiredService struct {
	store store.ImageAdminStore
}

func (s *DeleteExpiredService) DeleteExpired(ctx context.Context) (int64, error) {
	return s.store.DeleteExpiredImages(ctx)
}

func ProvideDeleteExpiredService(store *store.DBstore) *DeleteExpiredService {
	return &DeleteExpiredService{store: store}
}

//go:generate mockgen -destination=mock.go -package=image github.com/grafana/grafana/pkg/services/ngalert/image ImageService
type ImageService interface {
	// NewImage returns a new image for the alert instance.
	NewImage(ctx context.Context, r *models.AlertRule) (*models.Image, error)
}

// ScreenshotImageService takes screenshots of the alert rule and saves the
// image in the store. The image contains a unique token that can be passed
// as an annotation or label to the Alertmanager. This service cannot take
// screenshots of alert rules that are not associated with a dashboard panel.
type ScreenshotImageService struct {
	limiter      screenshot.RateLimiter
	logger       log.Logger
	screenshots  screenshot.ScreenshotService
	singleflight singleflight.Group
	store        store.ImageStore
	uploads      *UploadingService
}

// NewScreenshotImageService returns a new ScreenshotImageService.
func NewScreenshotImageService(
	limiter screenshot.RateLimiter,
	logger log.Logger,
	screenshots screenshot.ScreenshotService,
	store store.ImageStore,
	uploads *UploadingService) ImageService {
	return &ScreenshotImageService{
		limiter:     limiter,
		logger:      logger,
		screenshots: screenshots,
		store:       store,
		uploads:     uploads,
	}
}

// NewScreenshotImageServiceFromCfg returns a new ScreenshotImageService
// from the configuration.
func NewScreenshotImageServiceFromCfg(cfg *setting.Cfg, db *store.DBstore, ds dashboards.DashboardService,
	rs rendering.Service, r prometheus.Registerer) (ImageService, error) {
	var (
		limiter     screenshot.RateLimiter       = &screenshot.NoOpRateLimiter{}
		screenshots screenshot.ScreenshotService = &screenshot.ScreenshotUnavailableService{}
		uploads     *UploadingService            = nil
	)

	// If screenshots are enabled
	if cfg.UnifiedAlerting.Screenshots.Capture {
		limiter = screenshot.NewTokenRateLimiter(cfg.UnifiedAlerting.Screenshots.MaxConcurrentScreenshots)
		screenshots = screenshot.NewHeadlessScreenshotService(ds, rs, r)

		// Image uploading is an optional feature
		if cfg.UnifiedAlerting.Screenshots.UploadExternalImageStorage {
			m, err := imguploader.NewImageUploader()
			if err != nil {
				return nil, fmt.Errorf("failed to initialize uploading screenshot service: %w", err)
			}
			uploads = NewUploadingService(m, r)
		}
	}

	return NewScreenshotImageService(limiter, cfg.Logger, screenshots, db, uploads), nil
}

// NewImage returns a screenshot of the alert rule or an error.
//
// The alert rule must be associated with a dashboard panel for a screenshot to be
// taken. If the alert rule does not have a Dashboard UID in its annotations,
// or the dashboard does not exist, an ErrNoDashboard error is returned. If the
// alert rule has a Dashboard UID and the dashboard exists, but does not have a
// Panel ID in its annotations then an ErrNoPanel error is returned.
func (s *ScreenshotImageService) NewImage(ctx context.Context, r *models.AlertRule) (*models.Image, error) {
	if r.DashboardUID == nil {
		return nil, ErrNoDashboard
	}

	if r.PanelID == nil || *r.PanelID == 0 {
		return nil, ErrNoPanel
	}

	ctx, cancelFunc := context.WithTimeout(ctx, screenshotTimeout)
	defer cancelFunc()

	opts := screenshot.ScreenshotOptions{
		DashboardUID: *r.DashboardUID,
		PanelID:      *r.PanelID,
		Timeout:      screenshotTimeout,
	}

	optsHash := base64.StdEncoding.EncodeToString(opts.Hash())
	result, err, _ := s.singleflight.Do(optsHash, func() (interface{}, error) {
		screenshot, err := s.limiter.Do(ctx, opts, s.screenshots.Take)
		if err != nil {
			if errors.Is(err, dashboards.ErrDashboardNotFound) {
				return nil, ErrNoDashboard
			}
			return nil, err
		}
		image := models.Image{Path: screenshot.Path}
		if s.uploads != nil {
			if image, err = s.uploads.Upload(ctx, image); err != nil {
				s.logger.Warn("failed to upload image", "path", image.Path, "error", err)
			}
		}
		if err := s.store.SaveImage(ctx, &image); err != nil {
			return nil, fmt.Errorf("failed to save image: %w", err)
		}
		return image, nil
	})
	if err != nil {
		return nil, err
	}

	image := result.(models.Image)
	return &image, nil
}

// NotAvailableImageService is a service that returns ErrScreenshotsUnavailable.
type NotAvailableImageService struct{}

func (s *NotAvailableImageService) NewImage(_ context.Context, _ *models.AlertRule) (*models.Image, error) {
	return nil, screenshot.ErrScreenshotsUnavailable
}

// NoopImageService is a no-op image service.
type NoopImageService struct{}

func (s *NoopImageService) NewImage(_ context.Context, _ *models.AlertRule) (*models.Image, error) {
	return &models.Image{}, nil
}
<|endoftext|>


<commit_before>// Copyright Istio 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 endpoint

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/google/uuid"
	"google.golang.org/grpc"
	"google.golang.org/grpc/admin"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	xdscreds "google.golang.org/grpc/credentials/xds"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/reflection"
	"google.golang.org/grpc/xds"
	"k8s.io/utils/env"

	"istio.io/istio/pkg/istio-agent/grpcxds"
	"istio.io/istio/pkg/test/echo/common"
	"istio.io/istio/pkg/test/echo/common/response"
	"istio.io/istio/pkg/test/echo/proto"
	"istio.io/istio/pkg/test/echo/server/forwarder"
	"istio.io/istio/pkg/test/util/retry"
)

var _ Instance = &grpcInstance{}

// grpcServer is the intersection of used methods for grpc.Server and xds.GRPCServer
type grpcServer interface {
	reflection.GRPCServer
	Serve(listener net.Listener) error
	RegisterService(sd *grpc.ServiceDesc, ss interface{})
	Stop()
}

type grpcInstance struct {
	Config
	server   grpcServer
	cleanups []func()
}

func newGRPC(config Config) Instance {
	return &grpcInstance{
		Config: config,
	}
}

func (s *grpcInstance) GetConfig() Config {
	return s.Config
}

func (s *grpcInstance) newServer(opts ...grpc.ServerOption) grpcServer {
	if s.Port.XDSServer {
		if len(s.Port.XDSTestBootstrap) > 0 {
			opts = append(opts, xds.BootstrapContentsForTesting(s.Port.XDSTestBootstrap))
		}
		epLog.Infof("Using xDS for serverside gRPC on %d", s.Port.Port)
		return xds.NewGRPCServer(opts...)
	}
	return grpc.NewServer(opts...)
}

func (s *grpcInstance) Start(onReady OnReadyFunc) error {
	// Listen on the given port and update the port if it changed from what was passed in.
	listener, p, err := listenOnAddress(s.ListenerIP, s.Port.Port)
	if err != nil {
		return err
	}
	// Store the actual listening port back to the argument.
	s.Port.Port = p

	var opts []grpc.ServerOption
	if s.Port.TLS {
		epLog.Infof("Listening GRPC (over TLS) on %v", p)
		// Create the TLS credentials
		creds, errCreds := credentials.NewServerTLSFromFile(s.TLSCert, s.TLSKey)
		if errCreds != nil {
			epLog.Errorf("could not load TLS keys: %s", errCreds)
		}
		opts = append(opts, grpc.Creds(creds))
	} else if s.Port.XDSServer {
		epLog.Infof("Listening GRPC (over xDS-configured mTLS) on %v", p)
		creds, err := xdscreds.NewServerCredentials(xdscreds.ServerOptions{
			FallbackCreds: insecure.NewCredentials(),
		})
		if err != nil {
			return err
		}
		opts = append(opts, grpc.Creds(creds))
	} else {
		epLog.Infof("Listening GRPC on %v", p)
	}
	s.server = s.newServer(opts...)

	proto.RegisterEchoTestServiceServer(s.server, &grpcHandler{
		Config: s.Config,
	})
	reflection.Register(s.server)
	if val, _ := env.GetBool("EXPOSE_GRPC_ADMIN", false); val {
		cleanup, err := admin.Register(s.server)
		if err != nil {
			return err
		}
		s.cleanups = append(s.cleanups, cleanup)
	}
	// Start serving GRPC traffic.
	go func() {
		err := s.server.Serve(listener)
		epLog.Warnf("Port %d listener terminated with error: %v", p, err)
	}()

	// Notify the WaitGroup once the port has transitioned to ready.
	go s.awaitReady(onReady, listener)

	return nil
}

func (s *grpcInstance) awaitReady(onReady OnReadyFunc, listener net.Listener) {
	defer onReady()

	err := retry.UntilSuccess(func() error {
		cert, key, ca, err := s.certsFromBootstrapForReady()
		if err != nil {
			return err
		}
		f, err := forwarder.New(forwarder.Config{
			XDSTestBootstrap: s.Port.XDSTestBootstrap,
			Request: &proto.ForwardEchoRequest{
				Url:           "grpc://" + listener.Addr().String(),
				Message:       "hello",
				TimeoutMicros: common.DurationToMicros(readyInterval),
				CertFile:      cert,
				KeyFile:       key,
				CaCertFile:    ca,
			},
		})
		defer func() {
			_ = f.Close()
		}()

		if err != nil {
			return err
		}
		_, err = f.Run(context.Background())
		return err
	}, retry.Timeout(readyTimeout), retry.Delay(readyInterval))
	if err != nil {
		epLog.Errorf("readiness failed for GRPC endpoint %s: %v", listener.Addr().String(), err)
	} else {
		epLog.Infof("ready for GRPC endpoint %s", listener.Addr().String())
	}
}

// TODO (hack) we have to send certs OR use xds:///fqdn. We don't know our own fqdn, and even if we did
// we could send traffic to another instance. Instead we look into gRPC internals to authenticate with ourself.
func (s *grpcInstance) certsFromBootstrapForReady() (cert string, key string, ca string, err error) {
	if !s.Port.XDSServer {
		return
	}

	var bootstrapData []byte
	if data := s.Port.XDSTestBootstrap; len(data) > 0 {
		bootstrapData = data
	} else if path := os.Getenv("GRPC_XDS_BOOTSTRAP"); len(path) > 0 {
		bootstrapData, err = os.ReadFile(path)
	} else if data := os.Getenv("GRPC_XDS_BOOTSTRAP_CONFIG"); len(data) > 0 {
		bootstrapData = []byte(data)
	}
	var bootstrap grpcxds.Bootstrap
	if uerr := json.Unmarshal(bootstrapData, &bootstrap); uerr != nil {
		err = uerr
		return
	}
	certs := bootstrap.FileWatcherProvider()
	if certs == nil {
		err = fmt.Errorf("no certs found in bootstrap")
		return
	}
	cert = certs.CertificateFile
	key = certs.PrivateKeyFile
	ca = certs.CACertificateFile
	return
}

func (s *grpcInstance) Close() error {
	if s.server != nil {
		s.server.Stop()
	}
	for _, cleanup := range s.cleanups {
		cleanup()
	}
	return nil
}

type grpcHandler struct {
	proto.UnimplementedEchoTestServiceServer
	Config
}

func (h *grpcHandler) Echo(ctx context.Context, req *proto.EchoRequest) (*proto.EchoResponse, error) {
	defer common.Metrics.GrpcRequests.With(common.PortLabel.Value(strconv.Itoa(h.Port.Port))).Increment()
	body := bytes.Buffer{}
	md, ok := metadata.FromIncomingContext(ctx)
	if ok {
		for key, values := range md {
			if strings.HasSuffix(key, "-bin") {
				continue
			}
			field := response.Field(key)
			if key == ":authority" {
				field = response.HostField
			}
			for _, value := range values {
				writeField(&body, field, value)
			}
		}
	}

	id := uuid.New()
	epLog.WithLabels("message", req.GetMessage(), "headers", md, "id", id).Infof("GRPC Request")

	portNumber := 0
	if h.Port != nil {
		portNumber = h.Port.Port
	}

	ip := "0.0.0.0"
	if peerInfo, ok := peer.FromContext(ctx); ok {
		ip, _, _ = net.SplitHostPort(peerInfo.Addr.String())
	}

	writeField(&body, response.StatusCodeField, response.StatusCodeOK)
	writeField(&body, response.ServiceVersionField, h.Version)
	writeField(&body, response.ServicePortField, strconv.Itoa(portNumber))
	writeField(&body, response.ClusterField, h.Cluster)
	writeField(&body, response.IPField, ip)
	writeField(&body, response.IstioVersionField, h.IstioVersion)
	writeField(&body, "Echo", req.GetMessage())

	if hostname, err := os.Hostname(); err == nil {
		writeField(&body, response.HostnameField, hostname)
	}

	epLog.WithLabels("id", id).Infof("GRPC Response")
	return &proto.EchoResponse{Message: body.String()}, nil
}

func (h *grpcHandler) ForwardEcho(ctx context.Context, req *proto.ForwardEchoRequest) (*proto.ForwardEchoResponse, error) {
	id := uuid.New()
	l := epLog.WithLabels("url", req.Url, "id", id)
	l.Infof("ForwardEcho request")
	t0 := time.Now()
	instance, err := forwarder.New(forwarder.Config{
		Request: req,
		Dialer:  h.Dialer,
	})
	if err != nil {
		return nil, err
	}
	defer instance.Close()

	ret, err := instance.Run(ctx)
	if err == nil {
		l.WithLabels("latency", time.Since(t0)).Infof("ForwardEcho response complete: %v", ret.GetOutput())
	} else {
		l.WithLabels("latency", time.Since(t0)).Infof("ForwardEcho response failed: %v", err)
	}
	return ret, err
}
<commit_msg>Remove duplicate method RegisterService (#34958)<commit_after>// Copyright Istio 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 endpoint

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/google/uuid"
	"google.golang.org/grpc"
	"google.golang.org/grpc/admin"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	xdscreds "google.golang.org/grpc/credentials/xds"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
	"google.golang.org/grpc/reflection"
	"google.golang.org/grpc/xds"
	"k8s.io/utils/env"

	"istio.io/istio/pkg/istio-agent/grpcxds"
	"istio.io/istio/pkg/test/echo/common"
	"istio.io/istio/pkg/test/echo/common/response"
	"istio.io/istio/pkg/test/echo/proto"
	"istio.io/istio/pkg/test/echo/server/forwarder"
	"istio.io/istio/pkg/test/util/retry"
)

var _ Instance = &grpcInstance{}

// grpcServer is the intersection of used methods for grpc.Server and xds.GRPCServer
type grpcServer interface {
	reflection.GRPCServer
	Serve(listener net.Listener) error
	Stop()
}

type grpcInstance struct {
	Config
	server   grpcServer
	cleanups []func()
}

func newGRPC(config Config) Instance {
	return &grpcInstance{
		Config: config,
	}
}

func (s *grpcInstance) GetConfig() Config {
	return s.Config
}

func (s *grpcInstance) newServer(opts ...grpc.ServerOption) grpcServer {
	if s.Port.XDSServer {
		if len(s.Port.XDSTestBootstrap) > 0 {
			opts = append(opts, xds.BootstrapContentsForTesting(s.Port.XDSTestBootstrap))
		}
		epLog.Infof("Using xDS for serverside gRPC on %d", s.Port.Port)
		return xds.NewGRPCServer(opts...)
	}
	return grpc.NewServer(opts...)
}

func (s *grpcInstance) Start(onReady OnReadyFunc) error {
	// Listen on the given port and update the port if it changed from what was passed in.
	listener, p, err := listenOnAddress(s.ListenerIP, s.Port.Port)
	if err != nil {
		return err
	}
	// Store the actual listening port back to the argument.
	s.Port.Port = p

	var opts []grpc.ServerOption
	if s.Port.TLS {
		epLog.Infof("Listening GRPC (over TLS) on %v", p)
		// Create the TLS credentials
		creds, errCreds := credentials.NewServerTLSFromFile(s.TLSCert, s.TLSKey)
		if errCreds != nil {
			epLog.Errorf("could not load TLS keys: %s", errCreds)
		}
		opts = append(opts, grpc.Creds(creds))
	} else if s.Port.XDSServer {
		epLog.Infof("Listening GRPC (over xDS-configured mTLS) on %v", p)
		creds, err := xdscreds.NewServerCredentials(xdscreds.ServerOptions{
			FallbackCreds: insecure.NewCredentials(),
		})
		if err != nil {
			return err
		}
		opts = append(opts, grpc.Creds(creds))
	} else {
		epLog.Infof("Listening GRPC on %v", p)
	}
	s.server = s.newServer(opts...)

	proto.RegisterEchoTestServiceServer(s.server, &grpcHandler{
		Config: s.Config,
	})
	reflection.Register(s.server)
	if val, _ := env.GetBool("EXPOSE_GRPC_ADMIN", false); val {
		cleanup, err := admin.Register(s.server)
		if err != nil {
			return err
		}
		s.cleanups = append(s.cleanups, cleanup)
	}
	// Start serving GRPC traffic.
	go func() {
		err := s.server.Serve(listener)
		epLog.Warnf("Port %d listener terminated with error: %v", p, err)
	}()

	// Notify the WaitGroup once the port has transitioned to ready.
	go s.awaitReady(onReady, listener)

	return nil
}

func (s *grpcInstance) awaitReady(onReady OnReadyFunc, listener net.Listener) {
	defer onReady()

	err := retry.UntilSuccess(func() error {
		cert, key, ca, err := s.certsFromBootstrapForReady()
		if err != nil {
			return err
		}
		f, err := forwarder.New(forwarder.Config{
			XDSTestBootstrap: s.Port.XDSTestBootstrap,
			Request: &proto.ForwardEchoRequest{
				Url:           "grpc://" + listener.Addr().String(),
				Message:       "hello",
				TimeoutMicros: common.DurationToMicros(readyInterval),
				CertFile:      cert,
				KeyFile:       key,
				CaCertFile:    ca,
			},
		})
		defer func() {
			_ = f.Close()
		}()

		if err != nil {
			return err
		}
		_, err = f.Run(context.Background())
		return err
	}, retry.Timeout(readyTimeout), retry.Delay(readyInterval))
	if err != nil {
		epLog.Errorf("readiness failed for GRPC endpoint %s: %v", listener.Addr().String(), err)
	} else {
		epLog.Infof("ready for GRPC endpoint %s", listener.Addr().String())
	}
}

// TODO (hack) we have to send certs OR use xds:///fqdn. We don't know our own fqdn, and even if we did
// we could send traffic to another instance. Instead we look into gRPC internals to authenticate with ourself.
func (s *grpcInstance) certsFromBootstrapForReady() (cert string, key string, ca string, err error) {
	if !s.Port.XDSServer {
		return
	}

	var bootstrapData []byte
	if data := s.Port.XDSTestBootstrap; len(data) > 0 {
		bootstrapData = data
	} else if path := os.Getenv("GRPC_XDS_BOOTSTRAP"); len(path) > 0 {
		bootstrapData, err = os.ReadFile(path)
	} else if data := os.Getenv("GRPC_XDS_BOOTSTRAP_CONFIG"); len(data) > 0 {
		bootstrapData = []byte(data)
	}
	var bootstrap grpcxds.Bootstrap
	if uerr := json.Unmarshal(bootstrapData, &bootstrap); uerr != nil {
		err = uerr
		return
	}
	certs := bootstrap.FileWatcherProvider()
	if certs == nil {
		err = fmt.Errorf("no certs found in bootstrap")
		return
	}
	cert = certs.CertificateFile
	key = certs.PrivateKeyFile
	ca = certs.CACertificateFile
	return
}

func (s *grpcInstance) Close() error {
	if s.server != nil {
		s.server.Stop()
	}
	for _, cleanup := range s.cleanups {
		cleanup()
	}
	return nil
}

type grpcHandler struct {
	proto.UnimplementedEchoTestServiceServer
	Config
}

func (h *grpcHandler) Echo(ctx context.Context, req *proto.EchoRequest) (*proto.EchoResponse, error) {
	defer common.Metrics.GrpcRequests.With(common.PortLabel.Value(strconv.Itoa(h.Port.Port))).Increment()
	body := bytes.Buffer{}
	md, ok := metadata.FromIncomingContext(ctx)
	if ok {
		for key, values := range md {
			if strings.HasSuffix(key, "-bin") {
				continue
			}
			field := response.Field(key)
			if key == ":authority" {
				field = response.HostField
			}
			for _, value := range values {
				writeField(&body, field, value)
			}
		}
	}

	id := uuid.New()
	epLog.WithLabels("message", req.GetMessage(), "headers", md, "id", id).Infof("GRPC Request")

	portNumber := 0
	if h.Port != nil {
		portNumber = h.Port.Port
	}

	ip := "0.0.0.0"
	if peerInfo, ok := peer.FromContext(ctx); ok {
		ip, _, _ = net.SplitHostPort(peerInfo.Addr.String())
	}

	writeField(&body, response.StatusCodeField, response.StatusCodeOK)
	writeField(&body, response.ServiceVersionField, h.Version)
	writeField(&body, response.ServicePortField, strconv.Itoa(portNumber))
	writeField(&body, response.ClusterField, h.Cluster)
	writeField(&body, response.IPField, ip)
	writeField(&body, response.IstioVersionField, h.IstioVersion)
	writeField(&body, "Echo", req.GetMessage())

	if hostname, err := os.Hostname(); err == nil {
		writeField(&body, response.HostnameField, hostname)
	}

	epLog.WithLabels("id", id).Infof("GRPC Response")
	return &proto.EchoResponse{Message: body.String()}, nil
}

func (h *grpcHandler) ForwardEcho(ctx context.Context, req *proto.ForwardEchoRequest) (*proto.ForwardEchoResponse, error) {
	id := uuid.New()
	l := epLog.WithLabels("url", req.Url, "id", id)
	l.Infof("ForwardEcho request")
	t0 := time.Now()
	instance, err := forwarder.New(forwarder.Config{
		Request: req,
		Dialer:  h.Dialer,
	})
	if err != nil {
		return nil, err
	}
	defer instance.Close()

	ret, err := instance.Run(ctx)
	if err == nil {
		l.WithLabels("latency", time.Since(t0)).Infof("ForwardEcho response complete: %v", ret.GetOutput())
	} else {
		l.WithLabels("latency", time.Since(t0)).Infof("ForwardEcho response failed: %v", err)
	}
	return ret, err
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package feature

import (
	"fmt"
	"strings"
	"testing"

	"github.com/spf13/pflag"
)

func TestFeatureGateFlag(t *testing.T) {
	// gates for testing
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	tests := []struct {
		arg        string
		expect     map[Feature]bool
		parseError string
	}{
		{
			arg: "",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "fooBarBaz=maybeidk",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
			parseError: "unrecognized key: fooBarBaz",
		},
		{
			arg: "AllAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "AllAlpha=true",
			expect: map[Feature]bool{
				allAlphaGate:  true,
				testAlphaGate: true,
				testBetaGate:  false,
			},
		},
		{
			arg: "AllAlpha=banana",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
			parseError: "invalid value of AllAlpha",
		},
		{
			arg: "AllAlpha=false,TestAlpha=true",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: true,
				testBetaGate:  false,
			},
		},
		{
			arg: "TestAlpha=true,AllAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: true,
				testBetaGate:  false,
			},
		},
		{
			arg: "AllAlpha=true,TestAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  true,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "TestAlpha=false,AllAlpha=true",
			expect: map[Feature]bool{
				allAlphaGate:  true,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "TestBeta=true,AllAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  true,
			},
		},
	}
	for i, test := range tests {
		fs := pflag.NewFlagSet("testfeaturegateflag", pflag.ContinueOnError)
		f := NewFeatureGate()
		f.Add(map[Feature]FeatureSpec{
			testAlphaGate: {Default: false, PreRelease: Alpha},
			testBetaGate:  {Default: false, PreRelease: Beta},
		})
		f.AddFlag(fs)

		err := fs.Parse([]string{fmt.Sprintf("--%s=%s", flagName, test.arg)})
		if test.parseError != "" {
			if !strings.Contains(err.Error(), test.parseError) {
				t.Errorf("%d: Parse() Expected %v, Got %v", i, test.parseError, err)
			}
		} else if err != nil {
			t.Errorf("%d: Parse() Expected nil, Got %v", i, err)
		}
		for k, v := range test.expect {
			if actual := f.enabled.Load().(map[Feature]bool)[k]; actual != v {
				t.Errorf("%d: expected %s=%v, Got %v", i, k, v, actual)
			}
		}
	}
}

func TestFeatureGateOverride(t *testing.T) {
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	// Don't parse the flag, assert defaults are used.
	var f FeatureGate = NewFeatureGate()
	f.Add(map[Feature]FeatureSpec{
		testAlphaGate: {Default: false, PreRelease: Alpha},
		testBetaGate:  {Default: false, PreRelease: Beta},
	})

	f.Set("TestAlpha=true,TestBeta=true")
	if f.Enabled(testAlphaGate) != true {
		t.Errorf("Expected true")
	}
	if f.Enabled(testBetaGate) != true {
		t.Errorf("Expected true")
	}

	f.Set("TestAlpha=false")
	if f.Enabled(testAlphaGate) != false {
		t.Errorf("Expected false")
	}
	if f.Enabled(testBetaGate) != true {
		t.Errorf("Expected true")
	}
}

func TestFeatureGateFlagDefaults(t *testing.T) {
	// gates for testing
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	// Don't parse the flag, assert defaults are used.
	var f FeatureGate = NewFeatureGate()
	f.Add(map[Feature]FeatureSpec{
		testAlphaGate: {Default: false, PreRelease: Alpha},
		testBetaGate:  {Default: true, PreRelease: Beta},
	})

	if f.Enabled(testAlphaGate) != false {
		t.Errorf("Expected false")
	}
	if f.Enabled(testBetaGate) != true {
		t.Errorf("Expected true")
	}
}
<commit_msg>add unit test for function FeatureGateSetFromMap and FeatureGateString<commit_after>/*
Copyright 2016 The Kubernetes Authors.

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

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

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

package feature

import (
	"fmt"
	"strings"
	"testing"

	"github.com/spf13/pflag"
)

func TestFeatureGateFlag(t *testing.T) {
	// gates for testing
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	tests := []struct {
		arg        string
		expect     map[Feature]bool
		parseError string
	}{
		{
			arg: "",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "fooBarBaz=maybeidk",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
			parseError: "unrecognized key: fooBarBaz",
		},
		{
			arg: "AllAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "AllAlpha=true",
			expect: map[Feature]bool{
				allAlphaGate:  true,
				testAlphaGate: true,
				testBetaGate:  false,
			},
		},
		{
			arg: "AllAlpha=banana",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  false,
			},
			parseError: "invalid value of AllAlpha",
		},
		{
			arg: "AllAlpha=false,TestAlpha=true",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: true,
				testBetaGate:  false,
			},
		},
		{
			arg: "TestAlpha=true,AllAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: true,
				testBetaGate:  false,
			},
		},
		{
			arg: "AllAlpha=true,TestAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  true,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "TestAlpha=false,AllAlpha=true",
			expect: map[Feature]bool{
				allAlphaGate:  true,
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			arg: "TestBeta=true,AllAlpha=false",
			expect: map[Feature]bool{
				allAlphaGate:  false,
				testAlphaGate: false,
				testBetaGate:  true,
			},
		},
	}
	for i, test := range tests {
		fs := pflag.NewFlagSet("testfeaturegateflag", pflag.ContinueOnError)
		f := NewFeatureGate()
		f.Add(map[Feature]FeatureSpec{
			testAlphaGate: {Default: false, PreRelease: Alpha},
			testBetaGate:  {Default: false, PreRelease: Beta},
		})
		f.AddFlag(fs)

		err := fs.Parse([]string{fmt.Sprintf("--%s=%s", flagName, test.arg)})
		if test.parseError != "" {
			if !strings.Contains(err.Error(), test.parseError) {
				t.Errorf("%d: Parse() Expected %v, Got %v", i, test.parseError, err)
			}
		} else if err != nil {
			t.Errorf("%d: Parse() Expected nil, Got %v", i, err)
		}
		for k, v := range test.expect {
			if actual := f.enabled.Load().(map[Feature]bool)[k]; actual != v {
				t.Errorf("%d: expected %s=%v, Got %v", i, k, v, actual)
			}
		}
	}
}

func TestFeatureGateOverride(t *testing.T) {
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	// Don't parse the flag, assert defaults are used.
	var f FeatureGate = NewFeatureGate()
	f.Add(map[Feature]FeatureSpec{
		testAlphaGate: {Default: false, PreRelease: Alpha},
		testBetaGate:  {Default: false, PreRelease: Beta},
	})

	f.Set("TestAlpha=true,TestBeta=true")
	if f.Enabled(testAlphaGate) != true {
		t.Errorf("Expected true")
	}
	if f.Enabled(testBetaGate) != true {
		t.Errorf("Expected true")
	}

	f.Set("TestAlpha=false")
	if f.Enabled(testAlphaGate) != false {
		t.Errorf("Expected false")
	}
	if f.Enabled(testBetaGate) != true {
		t.Errorf("Expected true")
	}
}

func TestFeatureGateFlagDefaults(t *testing.T) {
	// gates for testing
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	// Don't parse the flag, assert defaults are used.
	var f FeatureGate = NewFeatureGate()
	f.Add(map[Feature]FeatureSpec{
		testAlphaGate: {Default: false, PreRelease: Alpha},
		testBetaGate:  {Default: true, PreRelease: Beta},
	})

	if f.Enabled(testAlphaGate) != false {
		t.Errorf("Expected false")
	}
	if f.Enabled(testBetaGate) != true {
		t.Errorf("Expected true")
	}
}

func TestFeatureGateSetFromMap(t *testing.T) {
	// gates for testing
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"

	tests := []struct {
		name        string
		setmap      map[string]bool
		expect      map[Feature]bool
		setmapError string
	}{
		{
			name: "set TestAlpha and TestBeta true",
			setmap: map[string]bool{
				"TestAlpha": true,
				"TestBeta":  true,
			},
			expect: map[Feature]bool{
				testAlphaGate: true,
				testBetaGate:  true,
			},
		},
		{
			name: "set TestBeta true",
			setmap: map[string]bool{
				"TestBeta": true,
			},
			expect: map[Feature]bool{
				testAlphaGate: false,
				testBetaGate:  true,
			},
		},
		{
			name: "set TestAlpha false",
			setmap: map[string]bool{
				"TestAlpha": false,
			},
			expect: map[Feature]bool{
				testAlphaGate: false,
				testBetaGate:  false,
			},
		},
		{
			name: "set TestInvaild true",
			setmap: map[string]bool{
				"TestInvaild": true,
			},
			expect: map[Feature]bool{
				testAlphaGate: false,
				testBetaGate:  false,
			},
			setmapError: "unrecognized key:",
		},
	}
	for i, test := range tests {
		t.Run(fmt.Sprintf("SetFromMap %s", test.name), func(t *testing.T) {
			f := NewFeatureGate()
			f.Add(map[Feature]FeatureSpec{
				testAlphaGate: {Default: false, PreRelease: Alpha},
				testBetaGate:  {Default: false, PreRelease: Beta},
			})
			err := f.SetFromMap(test.setmap)
			if test.setmapError != "" {
				if !strings.Contains(err.Error(), test.setmapError) {
					t.Errorf("%d: SetFromMap(%#v) Expected err:%v, Got err:%v", i, test.setmap, test.setmapError, err)
				}
			} else if err != nil {
				t.Errorf("%d: SetFromMap(%#v) Expected success, Got err:%v", i, test.setmap, err)
			}
			for k, v := range test.expect {
				if actual := f.Enabled(k); actual != v {
					t.Errorf("%d: SetFromMap(%#v) Expected %s=%v, Got %s=%v", i, test.setmap, k, v, k, actual)
				}
			}
		})
	}
}

func TestFeatureGateString(t *testing.T) {
	// gates for testing
	const testAlphaGate Feature = "TestAlpha"
	const testBetaGate Feature = "TestBeta"
	const testGAGate Feature = "TestGA"

	featuremap := map[Feature]FeatureSpec{
		testGAGate:    {Default: true, PreRelease: GA},
		testAlphaGate: {Default: false, PreRelease: Alpha},
		testBetaGate:  {Default: true, PreRelease: Beta},
	}

	tests := []struct {
		setmap map[string]bool
		expect string
	}{
		{
			setmap: map[string]bool{
				"TestAlpha": false,
			},
			expect: "TestAlpha=false",
		},
		{
			setmap: map[string]bool{
				"TestAlpha": false,
				"TestBeta":  true,
			},
			expect: "TestAlpha=false,TestBeta=true",
		},
		{
			setmap: map[string]bool{
				"TestGA":    true,
				"TestAlpha": false,
				"TestBeta":  true,
			},
			expect: "TestAlpha=false,TestBeta=true,TestGA=true",
		},
	}
	for i, test := range tests {
		t.Run(fmt.Sprintf("SetFromMap %s", test.expect), func(t *testing.T) {
			f := NewFeatureGate()
			f.Add(featuremap)
			f.SetFromMap(test.setmap)
			result := f.String()
			if result != test.expect {
				t.Errorf("%d: SetFromMap(%#v) Expected %s, Got %s", i, test.setmap, test.expect, result)
			}
		})
	}
}
<|endoftext|>


<commit_before>package parser

import (
	"testing"

	"github.com/Tamrin007/monkey/ast"
	"github.com/Tamrin007/monkey/lexer"
)

func TestLetStatements(t *testing.T) {
	input := `
let x = 5;
let y = 10;
let foobar = 838383;
	`

	l := lexer.New(input)
	p := New(l)

	program := p.ParseProgram()
	checkParserErrors(t, p)
	if program == nil {
		t.Fatalf("ParseProgram() returned nil")
	}
	if len(program.Statements) != 3 {
		t.Fatalf("program.Statements does not contain 3 statements. got=%d", len(program.Statements))
	}

	tests := []struct {
		expectedIdentifier string
	}{
		{"x"},
		{"y"},
		{"foobar"},
	}

	for i, tt := range tests {
		stmt := program.Statements[i]
		if !testLetStatement(t, stmt, tt.expectedIdentifier) {
			return
		}
	}
}

func checkParserErrors(t *testing.T, p *Parser) {
	errors := p.Errors()
	if len(errors) == 0 {
		return
	}

	t.Errorf("Parser has %d errors", len(errors))
	for _, msg := range errors {
		t.Errorf("parser error: %q", msg)
	}
	t.FailNow()
}

func testLetStatement(t *testing.T, s ast.Statement, name string) bool {
	if s.TokenLiteral() != "let" {
		t.Errorf("s.TokenLiteral not 'let'. got=%q", s.TokenLiteral())
		return false
	}

	letStmt, ok := s.(*ast.LetStatement)
	if !ok {
		t.Errorf("s not *ast.LetStatement. got=%T", s)
		return false
	}

	if letStmt.Name.TokenLiteral() != name {
		t.Errorf("s.Name not '%s'. got=%s", name, letStmt.Name)
		return false
	}

	return true
}

func TestReturnStatement(t *testing.T) {
	input := `
return 5;
return 10;
return 993322;
`

	l := lexer.New(input)
	p := New(l)

	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 3 {
		t.Fatalf("program.Statements does not contain 3 statements. got=%d", len(program.Statements))
	}

	for _, stmt := range program.Statements {
		returnStmt, ok := stmt.(*ast.ReturnStatement)
		if !ok {
			t.Errorf("stmt not *ast.returnStatement. got=%T", stmt)
			continue
		}
		if returnStmt.TokenLiteral() != "return" {
			t.Errorf("returnStmt.TokenLiteral not 'return', got %q", returnStmt.TokenLiteral())
		}
	}
}

func TestIdentifierExpression(t *testing.T) {
	input := "foobar;"

	l := lexer.New(input)
	p := New(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("program has not enough statements. got=%d", len(program.Statements))
	}
	stmt, ok := program.Statements[0].(*ast.ExpressionStatement)
	if !ok {
		t.Fatalf("program.Statements[0] is not ast.ExpressionStatement. got=%T", program.Statements[0])
	}

	ident, ok := stmt.Expression.(*ast.Identifier)
	if !ok {
		t.Fatalf("exp not *ast.Identifier. got=%T", stmt.Expression)
	}
	if ident.Value != "foobar" {
		t.Errorf("ident.Value not %s. got=%s", "foobar", ident.Value)
	}
	if ident.TokenLiteral() != "foobar" {
		t.Errorf("ident.TokenLiteral not %s. got=%s", "foobar", ident.TokenLiteral())
	}
}
<commit_msg>add a parser test for integer literal<commit_after>package parser

import (
	"testing"

	"github.com/Tamrin007/monkey/ast"
	"github.com/Tamrin007/monkey/lexer"
)

func TestLetStatements(t *testing.T) {
	input := `
let x = 5;
let y = 10;
let foobar = 838383;
	`

	l := lexer.New(input)
	p := New(l)

	program := p.ParseProgram()
	checkParserErrors(t, p)
	if program == nil {
		t.Fatalf("ParseProgram() returned nil")
	}
	if len(program.Statements) != 3 {
		t.Fatalf("program.Statements does not contain 3 statements. got=%d", len(program.Statements))
	}

	tests := []struct {
		expectedIdentifier string
	}{
		{"x"},
		{"y"},
		{"foobar"},
	}

	for i, tt := range tests {
		stmt := program.Statements[i]
		if !testLetStatement(t, stmt, tt.expectedIdentifier) {
			return
		}
	}
}

func checkParserErrors(t *testing.T, p *Parser) {
	errors := p.Errors()
	if len(errors) == 0 {
		return
	}

	t.Errorf("Parser has %d errors", len(errors))
	for _, msg := range errors {
		t.Errorf("parser error: %q", msg)
	}
	t.FailNow()
}

func testLetStatement(t *testing.T, s ast.Statement, name string) bool {
	if s.TokenLiteral() != "let" {
		t.Errorf("s.TokenLiteral not 'let'. got=%q", s.TokenLiteral())
		return false
	}

	letStmt, ok := s.(*ast.LetStatement)
	if !ok {
		t.Errorf("s not *ast.LetStatement. got=%T", s)
		return false
	}

	if letStmt.Name.TokenLiteral() != name {
		t.Errorf("s.Name not '%s'. got=%s", name, letStmt.Name)
		return false
	}

	return true
}

func TestReturnStatement(t *testing.T) {
	input := `
return 5;
return 10;
return 993322;
`

	l := lexer.New(input)
	p := New(l)

	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 3 {
		t.Fatalf("program.Statements does not contain 3 statements. got=%d", len(program.Statements))
	}

	for _, stmt := range program.Statements {
		returnStmt, ok := stmt.(*ast.ReturnStatement)
		if !ok {
			t.Errorf("stmt not *ast.returnStatement. got=%T", stmt)
			continue
		}
		if returnStmt.TokenLiteral() != "return" {
			t.Errorf("returnStmt.TokenLiteral not 'return', got %q", returnStmt.TokenLiteral())
		}
	}
}

func TestIdentifierExpression(t *testing.T) {
	input := "foobar;"

	l := lexer.New(input)
	p := New(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("program has not enough statements. got=%d", len(program.Statements))
	}
	stmt, ok := program.Statements[0].(*ast.ExpressionStatement)
	if !ok {
		t.Fatalf("program.Statements[0] is not ast.ExpressionStatement. got=%T", program.Statements[0])
	}

	ident, ok := stmt.Expression.(*ast.Identifier)
	if !ok {
		t.Fatalf("exp not *ast.Identifier. got=%T", stmt.Expression)
	}
	if ident.Value != "foobar" {
		t.Errorf("ident.Value not %s. got=%s", "foobar", ident.Value)
	}
	if ident.TokenLiteral() != "foobar" {
		t.Errorf("ident.TokenLiteral not %s. got=%s", "foobar", ident.TokenLiteral())
	}
}

func TestIntegerLiteralExpression(t *testing.T) {
	input := "5;"

	l := lexer.New(input)
	p := New(l)
	program := p.ParseProgram()
	checkParserErrors(t, p)

	if len(program.Statements) != 1 {
		t.Fatalf("program has not enough statements. got=%d", len(program.Statements))
	}
	stmt, ok := program.Statements[0].(*ast.ExpressionStatement)
	if !ok {
		t.Fatalf("program.Statements[0] is not ast.ExpressionStatement. got=%T", program.Statements[0])
	}

	literal, ok := stmt.Expression.(*ast.IntegerLiteral)
	if !ok {
		t.Fatalf("exp not *ast.IntegerLiteral. got=%T", stmt.Expression)
	}
	if literal.Value != 5 {
		t.Errorf("literal.Value not %d. got=%d", 5, literal.Value)
	}
	if literal.TokenLiteral() != "5" {
		t.Errorf("literal.TokenLiteral not %s. got=%s", "5", literal.TokenLiteral())
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// 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 spanstore

import (
	"context"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/uber/jaeger-lib/metrics"
	"go.uber.org/zap"
	"gopkg.in/olivere/elastic.v5"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/model/converter/json"
	jModel "github.com/jaegertracing/jaeger/model/json"
	"github.com/jaegertracing/jaeger/pkg/cache"
	"github.com/jaegertracing/jaeger/pkg/es"
	storageMetrics "github.com/jaegertracing/jaeger/storage/spanstore/metrics"
)

const (
	spanType    = "span"
	serviceType = "service"

	defaultNumShards   = 5
	defaultNumReplicas = 1
)

type spanWriterMetrics struct {
	indexCreate *storageMetrics.WriteMetrics
	spans       *storageMetrics.WriteMetrics
}

type serviceWriter func(string, *jModel.Span) error

// SpanWriter is a wrapper around elastic.Client
type SpanWriter struct {
	ctx           context.Context
	client        es.Client
	logger        *zap.Logger
	writerMetrics spanWriterMetrics // TODO: build functions to wrap around each Do fn
	indexCache    cache.Cache
	serviceWriter serviceWriter
	numShards     int64
	numReplicas   int64
}

// Service is the JSON struct for service:operation documents in ElasticSearch
type Service struct {
	ServiceName   string `json:"serviceName"`
	OperationName string `json:"operationName"`
}

// Span adds a StartTimeMillis field to the standard JSON span.
// ElasticSearch does not support a UNIX Epoch timestamp in microseconds,
// so Jaeger maps StartTime to a 'long' type. This extra StartTimeMillis field
// works around this issue, enabling timerange queries.
type Span struct {
	*jModel.Span
	StartTimeMillis uint64 `json:"startTimeMillis"`
}

// NewSpanWriter creates a new SpanWriter for use
func NewSpanWriter(
	client es.Client,
	logger *zap.Logger,
	metricsFactory metrics.Factory,
	numShards int64,
	numReplicas int64,
) *SpanWriter {
	ctx := context.Background()
	if numShards == 0 {
		numShards = defaultNumShards
	}
	if numReplicas == 0 {
		numReplicas = defaultNumReplicas
	}
	// TODO: Configurable TTL
	serviceOperationStorage := NewServiceOperationStorage(ctx, client, metricsFactory, logger, time.Hour*12)
	return &SpanWriter{
		ctx:    ctx,
		client: client,
		logger: logger,
		writerMetrics: spanWriterMetrics{
			indexCreate: storageMetrics.NewWriteMetrics(metricsFactory, "IndexCreate"),
			spans:       storageMetrics.NewWriteMetrics(metricsFactory, "Spans"),
		},
		serviceWriter: serviceOperationStorage.Write,
		indexCache: cache.NewLRUWithOptions(
			5,
			&cache.Options{
				TTL: 48 * time.Hour,
			},
		),
		numShards:   numShards,
		numReplicas: numReplicas,
	}
}

// WriteSpan writes a span and its corresponding service:operation in ElasticSearch
func (s *SpanWriter) WriteSpan(span *model.Span) error {
	spanIndexName, serviceIndexName := indexNames(span)
	// Convert model.Span into json.Span
	jsonSpan := json.FromDomainEmbedProcess(span)

	if err := s.createIndex(serviceIndexName, serviceMapping, jsonSpan); err != nil {
		return err
	}
	if err := s.writeService(serviceIndexName, jsonSpan); err != nil {
		return err
	}
	if err := s.createIndex(spanIndexName, spanMapping, jsonSpan); err != nil {
		return err
	}
	return s.writeSpan(spanIndexName, jsonSpan)
}

func indexNames(span *model.Span) (string, string) {
	spanDate := span.StartTime.Format("2006-01-02")
	return spanIndexPrefix + spanDate, serviceIndexPrefix + spanDate
}

func (s *SpanWriter) createIndex(indexName string, mapping string, jsonSpan *jModel.Span) error {
	if !keyInCache(indexName, s.indexCache) {
		start := time.Now()
		exists, _ := s.client.IndexExists(indexName).Do(s.ctx) // don't need to check the error because the exists variable will be false anyway if there is an error
		if !exists {
			// if there are multiple collectors writing to the same elasticsearch host a race condition can occur - create the index multiple times
			// we check for the error type to minimize errors
			_, err := s.client.CreateIndex(indexName).Body(s.fixMapping(mapping)).Do(s.ctx)
			s.writerMetrics.indexCreate.Emit(err, time.Since(start))
			if err != nil {
				eErr, ok := err.(*elastic.Error)
				if !ok || eErr.Details != nil && eErr.Details.Type != "index_already_exists_exception" {
					return s.logError(jsonSpan, err, "Failed to create index", s.logger)
				}
			}
		}
		writeCache(indexName, s.indexCache)
	}
	return nil
}

func keyInCache(key string, c cache.Cache) bool {
	return c.Get(key) != nil
}

func writeCache(key string, c cache.Cache) {
	c.Put(key, key)
}

func (s *SpanWriter) fixMapping(mapping string) string {
	mapping = strings.Replace(mapping, "${__NUMBER_OF_SHARDS__}", strconv.FormatInt(s.numShards, 10), 1)
	mapping = strings.Replace(mapping, "${__NUMBER_OF_REPLICAS__}", strconv.FormatInt(s.numReplicas, 10), 1)
	return mapping
}

func (s *SpanWriter) writeService(indexName string, jsonSpan *jModel.Span) error {
	return s.serviceWriter(indexName, jsonSpan)
}

func (s *SpanWriter) writeSpan(indexName string, jsonSpan *jModel.Span) error {
	start := time.Now()
	elasticSpan := Span{Span: jsonSpan, StartTimeMillis: jsonSpan.StartTime / 1000} // Microseconds to milliseconds
	_, err := s.client.Index().Index(indexName).Type(spanType).BodyJson(&elasticSpan).Do(s.ctx)
	s.writerMetrics.spans.Emit(err, time.Since(start))
	if err != nil {
		return s.logError(jsonSpan, err, "Failed to insert span", s.logger)
	}
	return nil
}

func (s *SpanWriter) logError(span *jModel.Span, err error, msg string, logger *zap.Logger) error {
	logger.
		With(zap.String("trace_id", string(span.TraceID))).
		With(zap.String("span_id", string(span.SpanID))).
		With(zap.Error(err)).
		Error(msg)
	return errors.Wrap(err, msg)
}
<commit_msg>Use UTC timezone for ES indices (#646)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// 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 spanstore

import (
	"context"
	"strconv"
	"strings"
	"time"

	"github.com/pkg/errors"
	"github.com/uber/jaeger-lib/metrics"
	"go.uber.org/zap"
	"gopkg.in/olivere/elastic.v5"

	"github.com/jaegertracing/jaeger/model"
	"github.com/jaegertracing/jaeger/model/converter/json"
	jModel "github.com/jaegertracing/jaeger/model/json"
	"github.com/jaegertracing/jaeger/pkg/cache"
	"github.com/jaegertracing/jaeger/pkg/es"
	storageMetrics "github.com/jaegertracing/jaeger/storage/spanstore/metrics"
)

const (
	spanType    = "span"
	serviceType = "service"

	defaultNumShards   = 5
	defaultNumReplicas = 1
)

type spanWriterMetrics struct {
	indexCreate *storageMetrics.WriteMetrics
	spans       *storageMetrics.WriteMetrics
}

type serviceWriter func(string, *jModel.Span) error

// SpanWriter is a wrapper around elastic.Client
type SpanWriter struct {
	ctx           context.Context
	client        es.Client
	logger        *zap.Logger
	writerMetrics spanWriterMetrics // TODO: build functions to wrap around each Do fn
	indexCache    cache.Cache
	serviceWriter serviceWriter
	numShards     int64
	numReplicas   int64
}

// Service is the JSON struct for service:operation documents in ElasticSearch
type Service struct {
	ServiceName   string `json:"serviceName"`
	OperationName string `json:"operationName"`
}

// Span adds a StartTimeMillis field to the standard JSON span.
// ElasticSearch does not support a UNIX Epoch timestamp in microseconds,
// so Jaeger maps StartTime to a 'long' type. This extra StartTimeMillis field
// works around this issue, enabling timerange queries.
type Span struct {
	*jModel.Span
	StartTimeMillis uint64 `json:"startTimeMillis"`
}

// NewSpanWriter creates a new SpanWriter for use
func NewSpanWriter(
	client es.Client,
	logger *zap.Logger,
	metricsFactory metrics.Factory,
	numShards int64,
	numReplicas int64,
) *SpanWriter {
	ctx := context.Background()
	if numShards == 0 {
		numShards = defaultNumShards
	}
	if numReplicas == 0 {
		numReplicas = defaultNumReplicas
	}
	// TODO: Configurable TTL
	serviceOperationStorage := NewServiceOperationStorage(ctx, client, metricsFactory, logger, time.Hour*12)
	return &SpanWriter{
		ctx:    ctx,
		client: client,
		logger: logger,
		writerMetrics: spanWriterMetrics{
			indexCreate: storageMetrics.NewWriteMetrics(metricsFactory, "IndexCreate"),
			spans:       storageMetrics.NewWriteMetrics(metricsFactory, "Spans"),
		},
		serviceWriter: serviceOperationStorage.Write,
		indexCache: cache.NewLRUWithOptions(
			5,
			&cache.Options{
				TTL: 48 * time.Hour,
			},
		),
		numShards:   numShards,
		numReplicas: numReplicas,
	}
}

// WriteSpan writes a span and its corresponding service:operation in ElasticSearch
func (s *SpanWriter) WriteSpan(span *model.Span) error {
	spanIndexName, serviceIndexName := indexNames(span)
	// Convert model.Span into json.Span
	jsonSpan := json.FromDomainEmbedProcess(span)

	if err := s.createIndex(serviceIndexName, serviceMapping, jsonSpan); err != nil {
		return err
	}
	if err := s.writeService(serviceIndexName, jsonSpan); err != nil {
		return err
	}
	if err := s.createIndex(spanIndexName, spanMapping, jsonSpan); err != nil {
		return err
	}
	return s.writeSpan(spanIndexName, jsonSpan)
}

func indexNames(span *model.Span) (string, string) {
	spanDate := span.StartTime.UTC().Format("2006-01-02")
	return spanIndexPrefix + spanDate, serviceIndexPrefix + spanDate
}

func (s *SpanWriter) createIndex(indexName string, mapping string, jsonSpan *jModel.Span) error {
	if !keyInCache(indexName, s.indexCache) {
		start := time.Now()
		exists, _ := s.client.IndexExists(indexName).Do(s.ctx) // don't need to check the error because the exists variable will be false anyway if there is an error
		if !exists {
			// if there are multiple collectors writing to the same elasticsearch host a race condition can occur - create the index multiple times
			// we check for the error type to minimize errors
			_, err := s.client.CreateIndex(indexName).Body(s.fixMapping(mapping)).Do(s.ctx)
			s.writerMetrics.indexCreate.Emit(err, time.Since(start))
			if err != nil {
				eErr, ok := err.(*elastic.Error)
				if !ok || eErr.Details != nil && eErr.Details.Type != "index_already_exists_exception" {
					return s.logError(jsonSpan, err, "Failed to create index", s.logger)
				}
			}
		}
		writeCache(indexName, s.indexCache)
	}
	return nil
}

func keyInCache(key string, c cache.Cache) bool {
	return c.Get(key) != nil
}

func writeCache(key string, c cache.Cache) {
	c.Put(key, key)
}

func (s *SpanWriter) fixMapping(mapping string) string {
	mapping = strings.Replace(mapping, "${__NUMBER_OF_SHARDS__}", strconv.FormatInt(s.numShards, 10), 1)
	mapping = strings.Replace(mapping, "${__NUMBER_OF_REPLICAS__}", strconv.FormatInt(s.numReplicas, 10), 1)
	return mapping
}

func (s *SpanWriter) writeService(indexName string, jsonSpan *jModel.Span) error {
	return s.serviceWriter(indexName, jsonSpan)
}

func (s *SpanWriter) writeSpan(indexName string, jsonSpan *jModel.Span) error {
	start := time.Now()
	elasticSpan := Span{Span: jsonSpan, StartTimeMillis: jsonSpan.StartTime / 1000} // Microseconds to milliseconds
	_, err := s.client.Index().Index(indexName).Type(spanType).BodyJson(&elasticSpan).Do(s.ctx)
	s.writerMetrics.spans.Emit(err, time.Since(start))
	if err != nil {
		return s.logError(jsonSpan, err, "Failed to insert span", s.logger)
	}
	return nil
}

func (s *SpanWriter) logError(span *jModel.Span, err error, msg string, logger *zap.Logger) error {
	logger.
		With(zap.String("trace_id", string(span.TraceID))).
		With(zap.String("span_id", string(span.SpanID))).
		With(zap.Error(err)).
		Error(msg)
	return errors.Wrap(err, msg)
}
<|endoftext|>


<commit_before>package parser

import (
  "fmt"
  "github.com/kedebug/LispEx/scope"
  "testing"
)

func test(exprs string) string {
  block := ParseFromString("Parser", exprs)
  return block.Eval(scope.NewRootScope()).String()
}

func testIf() bool {
  var exprs string = `
    (if 1 2 invalid)
    (if #f invalid 'ok)
    (if #t 1)
    (if #f 1)
  `
  expected := "2\nok\n1"
  return expected == test(exprs)
}

func testDefine() bool {
  var exprs string = `
    (define x 3) x (+ x x)
    (define x 1) x (define x (+ x 1)) x
    (define y 2) ((lambda (x) (define y 1) (+ x y)) 3) y
    (define f (lambda () (+ 1 2))) (f)
    (define add3 (lambda (x) (+ x 3))) (add3 3)
    (define first car) (first '(1 2))
    (define (x y . z) (cons y z)) (x 1 2 3)
    (define (f x) (+ x y)) (define y 1) (f 1)
    (define plus (lambda (x) (+ x y))) (define y 1) (plus 3)
    (define x 0) (define z 1) (define (f x y) (set! z 2) (+ x y)) (f 1 2) x z
    (define x -2) x (set! x (* x x)) x
  `
  expected := "3\n6\n1\n2\n4\n2\n3\n6\n1\n(1 2 3)\n2\n4\n3\n0\n2\n-2\n4"
  return expected == test(exprs)
}

func testLambda() bool {
  var exprs string = `
    (lambda x 1 2 3)
    (lambda (x) 1 2 3)
    (lambda (x y) 1 2 3)
    (lambda (x . y) 1 2 3)
    ((lambda (x) x) 'a)
    ((lambda x x) 'a)
    ((lambda x x) 'a 'b)
    ((lambda (x y) (+ x y)) 3 5)
    ((lambda (x . y) (+ x (car y))) 1 2 5)
    ((lambda (x y . z) (+ x y (car z))) 1 2 5 11)
    (define x 10) ((lambda (x) x) 5) x
  `
  expected := `(lambda x 1 2 3)
(lambda (x) 1 2 3)
(lambda (x y) 1 2 3)
(lambda (x . y) 1 2 3)
a
(a)
(a b)
8
3
8
5
10`

  return expected == test(exprs)
}

func testQuasiquote() bool {
  var exprs string
  exprs += "`()"
  exprs += "`(())"
  exprs += "`(+ 2 3)"
  exprs += "`(+ 2 ,(+ 3 4))"
  exprs += "`(a b (,(+ 2 3) c) d)"
  exprs += "'`,(cons 'a 'b)"
  exprs += "`',(cons 'a 'b)"
  exprs += "`(+ ,@(cdr '(* 2 3)))"
  exprs += "`(1 2 `(3 4 ,@(5 6 8 9 10) 11 12) 13 14)"
  exprs += "``(+ ,,(+ 1 2) 2 3)"
  exprs += "`(1 2 `(10 ,',(+ 2 3)))"
  exprs += "`(+ 2 `(10 ,(+ 2 3)))"
  exprs += "`(1 2 `(10 ,,(+ 2 3)))"
  exprs += "`(1 `,(+ 1 ,(+ 2 3)) 4)"

  var expected string
  expected += "()"
  expected += "\n(())"
  expected += "\n(+ 2 3)"
  expected += "\n(+ 2 7)"
  expected += "\n(a b (5 c) d)"
  expected += "\n`,(cons 'a 'b)"
  expected += "\n'(a . b)"
  expected += "\n(+ 2 3)"
  expected += "\n(1 2 `(3 4 ,@(5 6 8 9 10) 11 12) 13 14)"
  expected += "\n`(+ ,3 2 3)"
  expected += "\n(1 2 `(10 ,'5))"
  expected += "\n(+ 2 `(10 ,(+ 2 3)))"
  expected += "\n(1 2 `(10 ,5))"
  expected += "\n(1 `,(+ 1 5) 4)"

  return expected == test(exprs)
}

func TestParser(t *testing.T) {
  if testDefine() {
    fmt.Println("TEST define:       PASS")
  } else {
    fmt.Println("TEST define:       FAILED")
  }
  if testLambda() {
    fmt.Println("TEST lambda:       PASS")
  } else {
    fmt.Println("TEST lambda:       FAILED")
  }
  if testQuasiquote() {
    fmt.Println("TEST quasiquote:   PASS")
  } else {
    fmt.Println("TEST quasiquote:   FAILED")
  }
  if testIf() {
    fmt.Println("TEST if:           PASS")
  } else {
    fmt.Println("TEST if:           FAILED")
  }
}
<commit_msg>add more nested quasiquote test cases<commit_after>package parser

import (
  "fmt"
  "github.com/kedebug/LispEx/scope"
  "testing"
)

func test(exprs string) string {
  block := ParseFromString("Parser", exprs)
  return block.Eval(scope.NewRootScope()).String()
}

func testIf() bool {
  var exprs string = `
    (if 1 2 invalid)
    (if #f invalid 'ok)
    (if #t 1)
    (if #f 1)
  `
  expected := "2\nok\n1"
  return expected == test(exprs)
}

func testDefine() bool {
  var exprs string = `
    (define x 3) x (+ x x)
    (define x 1) x (define x (+ x 1)) x
    (define y 2) ((lambda (x) (define y 1) (+ x y)) 3) y
    (define f (lambda () (+ 1 2))) (f)
    (define add3 (lambda (x) (+ x 3))) (add3 3)
    (define first car) (first '(1 2))
    (define (x y . z) (cons y z)) (x 1 2 3)
    (define (f x) (+ x y)) (define y 1) (f 1)
    (define plus (lambda (x) (+ x y))) (define y 1) (plus 3)
    (define x 0) (define z 1) (define (f x y) (set! z 2) (+ x y)) (f 1 2) x z
    (define x -2) x (set! x (* x x)) x
  `
  expected := "3\n6\n1\n2\n4\n2\n3\n6\n1\n(1 2 3)\n2\n4\n3\n0\n2\n-2\n4"
  return expected == test(exprs)
}

func testLambda() bool {
  var exprs string = `
    (lambda x 1 2 3)
    (lambda (x) 1 2 3)
    (lambda (x y) 1 2 3)
    (lambda (x . y) 1 2 3)
    ((lambda (x) x) 'a)
    ((lambda x x) 'a)
    ((lambda x x) 'a 'b)
    ((lambda (x y) (+ x y)) 3 5)
    ((lambda (x . y) (+ x (car y))) 1 2 5)
    ((lambda (x y . z) (+ x y (car z))) 1 2 5 11)
    (define x 10) ((lambda (x) x) 5) x
  `
  expected := `(lambda x 1 2 3)
(lambda (x) 1 2 3)
(lambda (x y) 1 2 3)
(lambda (x . y) 1 2 3)
a
(a)
(a b)
8
3
8
5
10`

  return expected == test(exprs)
}

func testQuasiquote() bool {
  var exprs string
  exprs += "`()"
  exprs += "`(())"
  exprs += "`(+ 2 3)"
  exprs += "`(+ 2 ,(+ 3 4))"
  exprs += "`(a b (,(+ 2 3) c) d)"
  exprs += "'`,(cons 'a 'b)"
  exprs += "`',(cons 'a 'b)"
  exprs += "`(+ ,@(cdr '(* 2 3)))"
  exprs += "`(1 2 `(3 4 ,@(5 6 8 9 10) 11 12) 13 14)"
  exprs += "`(1 2 `(3 4 ,@(5 6 ,@(cdr '(6 7 8)) 9 10) 11 12) 13 14)"
  exprs += "``(+ ,,(+ 1 2) 2 3)"
  exprs += "`(1 2 `(10 ,',(+ 2 3)))"
  exprs += "`(+ 2 `(10 ,(+ 2 3)))"
  exprs += "`(1 2 `(10 ,,(+ 2 3)))"
  exprs += "`(1 `,(+ 1 ,(+ 2 3)) 4)"

  var expected string
  expected += "()"
  expected += "\n(())"
  expected += "\n(+ 2 3)"
  expected += "\n(+ 2 7)"
  expected += "\n(a b (5 c) d)"
  expected += "\n`,(cons 'a 'b)"
  expected += "\n'(a . b)"
  expected += "\n(+ 2 3)"
  expected += "\n(1 2 `(3 4 ,@(5 6 8 9 10) 11 12) 13 14)"
  expected += "\n(1 2 `(3 4 ,@(5 6 7 8 9 10) 11 12) 13 14)"
  expected += "\n`(+ ,3 2 3)"
  expected += "\n(1 2 `(10 ,'5))"
  expected += "\n(+ 2 `(10 ,(+ 2 3)))"
  expected += "\n(1 2 `(10 ,5))"
  expected += "\n(1 `,(+ 1 5) 4)"

  return expected == test(exprs)
}

func TestParser(t *testing.T) {
  if testDefine() {
    fmt.Println("TEST define:       PASS")
  } else {
    fmt.Println("TEST define:       FAILED")
  }
  if testLambda() {
    fmt.Println("TEST lambda:       PASS")
  } else {
    fmt.Println("TEST lambda:       FAILED")
  }
  if testQuasiquote() {
    fmt.Println("TEST quasiquote:   PASS")
  } else {
    fmt.Println("TEST quasiquote:   FAILED")
  }
  if testIf() {
    fmt.Println("TEST if:           PASS")
  } else {
    fmt.Println("TEST if:           FAILED")
  }
}
<|endoftext|>