<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 aws

import (
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/autoscaler/cluster-autoscaler/cloudprovider"
	"k8s.io/autoscaler/cluster-autoscaler/config"
	"k8s.io/autoscaler/cluster-autoscaler/utils/errors"
	klog "k8s.io/klog/v2"
	schedulerframework "k8s.io/kubernetes/pkg/scheduler/framework"
)

const (
	// GPULabel is the label added to nodes with GPU resource.
	GPULabel = "k8s.amazonaws.com/accelerator"
)

var (
	availableGPUTypes = map[string]struct{}{
		"nvidia-tesla-k80":  {},
		"nvidia-tesla-p100": {},
		"nvidia-tesla-v100": {},
		"nvidia-tesla-t4":   {},
	}
)

// awsCloudProvider implements CloudProvider interface.
type awsCloudProvider struct {
	awsManager      *AwsManager
	resourceLimiter *cloudprovider.ResourceLimiter
}

// BuildAwsCloudProvider builds CloudProvider implementation for AWS.
func BuildAwsCloudProvider(awsManager *AwsManager, resourceLimiter *cloudprovider.ResourceLimiter) (cloudprovider.CloudProvider, error) {
	aws := &awsCloudProvider{
		awsManager:      awsManager,
		resourceLimiter: resourceLimiter,
	}
	return aws, nil
}

// Cleanup stops the go routine that is handling the current view of the ASGs in the form of a cache
func (aws *awsCloudProvider) Cleanup() error {
	aws.awsManager.Cleanup()
	return nil
}

// Name returns name of the cloud provider.
func (aws *awsCloudProvider) Name() string {
	return cloudprovider.AwsProviderName
}

// GPULabel returns the label added to nodes with GPU resource.
func (aws *awsCloudProvider) GPULabel() string {
	return GPULabel
}

// GetAvailableGPUTypes return all available GPU types cloud provider supports
func (aws *awsCloudProvider) GetAvailableGPUTypes() map[string]struct{} {
	return availableGPUTypes
}

// NodeGroups returns all node groups configured for this cloud provider.
func (aws *awsCloudProvider) NodeGroups() []cloudprovider.NodeGroup {
	asgs := aws.awsManager.getAsgs()
	ngs := make([]cloudprovider.NodeGroup, len(asgs))
	for i, asg := range asgs {
		ngs[i] = &AwsNodeGroup{
			asg:        asg,
			awsManager: aws.awsManager,
		}
	}

	return ngs
}

// NodeGroupForNode returns the node group for the given node.
func (aws *awsCloudProvider) NodeGroupForNode(node *apiv1.Node) (cloudprovider.NodeGroup, error) {
	if len(node.Spec.ProviderID) == 0 {
		klog.Warningf("Node %v has no providerId", node.Name)
		return nil, nil
	}
	ref, err := AwsRefFromProviderId(node.Spec.ProviderID)
	if err != nil {
		return nil, err
	}
	asg := aws.awsManager.GetAsgForInstance(*ref)

	if asg == nil {
		return nil, nil
	}

	return &AwsNodeGroup{
		asg:        asg,
		awsManager: aws.awsManager,
	}, nil
}

// Pricing returns pricing model for this cloud provider or error if not available.
func (aws *awsCloudProvider) Pricing() (cloudprovider.PricingModel, errors.AutoscalerError) {
	return nil, cloudprovider.ErrNotImplemented
}

// GetAvailableMachineTypes get all machine types that can be requested from the cloud provider.
func (aws *awsCloudProvider) GetAvailableMachineTypes() ([]string, error) {
	return []string{}, nil
}

// NewNodeGroup builds a theoretical node group based on the node definition provided. The node group is not automatically
// created on the cloud provider side. The node group is not returned by NodeGroups() until it is created.
func (aws *awsCloudProvider) NewNodeGroup(machineType string, labels map[string]string, systemLabels map[string]string,
	taints []apiv1.Taint, extraResources map[string]resource.Quantity) (cloudprovider.NodeGroup, error) {
	return nil, cloudprovider.ErrNotImplemented
}

// GetResourceLimiter returns struct containing limits (max, min) for resources (cores, memory etc.).
func (aws *awsCloudProvider) GetResourceLimiter() (*cloudprovider.ResourceLimiter, error) {
	return aws.resourceLimiter, nil
}

// Refresh is called before every main loop and can be used to dynamically update cloud provider state.
// In particular the list of node groups returned by NodeGroups can change as a result of CloudProvider.Refresh().
func (aws *awsCloudProvider) Refresh() error {
	return aws.awsManager.Refresh()
}

// AwsRef contains a reference to some entity in AWS world.
type AwsRef struct {
	Name string
}

// AwsInstanceRef contains a reference to an instance in the AWS world.
type AwsInstanceRef struct {
	ProviderID string
	Name       string
}

var validAwsRefIdRegex = regexp.MustCompile(fmt.Sprintf(`^aws\:\/\/\/[-0-9a-z]*\/[-0-9a-z]*(\/[-0-9a-z\.]*)?$|aws\:\/\/\/[-0-9a-z]*\/%s.*$`, placeholderInstanceNamePrefix))

// AwsRefFromProviderId creates AwsInstanceRef object from provider id which
// must be in format: aws:///zone/name
func AwsRefFromProviderId(id string) (*AwsInstanceRef, error) {
	if validAwsRefIdRegex.FindStringSubmatch(id) == nil {
		return nil, fmt.Errorf("wrong id: expected format aws:///<zone>/<name>, got %v", id)
	}
	splitted := strings.Split(id[7:], "/")
	return &AwsInstanceRef{
		ProviderID: id,
		Name:       splitted[1],
	}, nil
}

// AwsNodeGroup implements NodeGroup interface.
type AwsNodeGroup struct {
	awsManager *AwsManager
	asg        *asg
}

// MaxSize returns maximum size of the node group.
func (ng *AwsNodeGroup) MaxSize() int {
	return ng.asg.maxSize
}

// MinSize returns minimum size of the node group.
func (ng *AwsNodeGroup) MinSize() int {
	return ng.asg.minSize
}

// TargetSize returns the current TARGET size of the node group. It is possible that the
// number is different from the number of nodes registered in Kubernetes.
func (ng *AwsNodeGroup) TargetSize() (int, error) {
	return ng.asg.curSize, nil
}

// Exist checks if the node group really exists on the cloud provider side. Allows to tell the
// theoretical node group from the real one.
func (ng *AwsNodeGroup) Exist() bool {
	return true
}

// Create creates the node group on the cloud provider side.
func (ng *AwsNodeGroup) Create() (cloudprovider.NodeGroup, error) {
	return nil, cloudprovider.ErrAlreadyExist
}

// Autoprovisioned returns true if the node group is autoprovisioned.
func (ng *AwsNodeGroup) Autoprovisioned() bool {
	return false
}

// Delete deletes the node group on the cloud provider side.
// This will be executed only for autoprovisioned node groups, once their size drops to 0.
func (ng *AwsNodeGroup) Delete() error {
	return cloudprovider.ErrNotImplemented
}

// GetOptions returns NodeGroupAutoscalingOptions that should be used for this particular
// NodeGroup. Returning a nil will result in using default options.
func (ng *AwsNodeGroup) GetOptions(defaults config.NodeGroupAutoscalingOptions) (*config.NodeGroupAutoscalingOptions, error) {
	return nil, cloudprovider.ErrNotImplemented
}

// IncreaseSize increases Asg size
func (ng *AwsNodeGroup) IncreaseSize(delta int) error {
	if delta <= 0 {
		return fmt.Errorf("size increase must be positive")
	}
	size := ng.asg.curSize
	if size+delta > ng.asg.maxSize {
		return fmt.Errorf("size increase too large - desired:%d max:%d", size+delta, ng.asg.maxSize)
	}
	return ng.awsManager.SetAsgSize(ng.asg, size+delta)
}

// DecreaseTargetSize decreases the target size of the node group. This function
// doesn't permit to delete any existing node and can be used only to reduce the
// request for new nodes that have not been yet fulfilled. Delta should be negative.
// It is assumed that cloud provider will not delete the existing nodes if the size
// when there is an option to just decrease the target.
func (ng *AwsNodeGroup) DecreaseTargetSize(delta int) error {
	if delta >= 0 {
		return fmt.Errorf("size decrease size must be negative")
	}

	size := ng.asg.curSize
	nodes, err := ng.awsManager.GetAsgNodes(ng.asg.AwsRef)
	if err != nil {
		return err
	}
	if int(size)+delta < len(nodes) {
		return fmt.Errorf("attempt to delete existing nodes targetSize:%d delta:%d existingNodes: %d",
			size, delta, len(nodes))
	}
	return ng.awsManager.SetAsgSize(ng.asg, size+delta)
}

// Belongs returns true if the given node belongs to the NodeGroup.
func (ng *AwsNodeGroup) Belongs(node *apiv1.Node) (bool, error) {
	ref, err := AwsRefFromProviderId(node.Spec.ProviderID)
	if err != nil {
		return false, err
	}
	targetAsg := ng.awsManager.GetAsgForInstance(*ref)
	if targetAsg == nil {
		return false, fmt.Errorf("%s doesn't belong to a known asg", node.Name)
	}
	if targetAsg.AwsRef != ng.asg.AwsRef {
		return false, nil
	}
	return true, nil
}

// DeleteNodes deletes the nodes from the group.
func (ng *AwsNodeGroup) DeleteNodes(nodes []*apiv1.Node) error {
	size := ng.asg.curSize
	if int(size) <= ng.MinSize() {
		return fmt.Errorf("min size reached, nodes will not be deleted")
	}
	refs := make([]*AwsInstanceRef, 0, len(nodes))
	for _, node := range nodes {
		belongs, err := ng.Belongs(node)
		if err != nil {
			return err
		}
		if belongs != true {
			return fmt.Errorf("%s belongs to a different asg than %s", node.Name, ng.Id())
		}
		awsref, err := AwsRefFromProviderId(node.Spec.ProviderID)
		if err != nil {
			return err
		}
		refs = append(refs, awsref)
	}
	return ng.awsManager.DeleteInstances(refs)
}

// Id returns asg id.
func (ng *AwsNodeGroup) Id() string {
	return ng.asg.Name
}

// Debug returns a debug string for the Asg.
func (ng *AwsNodeGroup) Debug() string {
	return fmt.Sprintf("%s (%d:%d)", ng.Id(), ng.MinSize(), ng.MaxSize())
}

// Nodes returns a list of all nodes that belong to this node group.
func (ng *AwsNodeGroup) Nodes() ([]cloudprovider.Instance, error) {
	asgNodes, err := ng.awsManager.GetAsgNodes(ng.asg.AwsRef)
	if err != nil {
		return nil, err
	}

	instances := make([]cloudprovider.Instance, len(asgNodes))

	for i, asgNode := range asgNodes {
		instances[i] = cloudprovider.Instance{Id: asgNode.ProviderID}
	}
	return instances, nil
}

// TemplateNodeInfo returns a node template for this node group.
func (ng *AwsNodeGroup) TemplateNodeInfo() (*schedulerframework.NodeInfo, error) {
	template, err := ng.awsManager.getAsgTemplate(ng.asg)
	if err != nil {
		return nil, err
	}

	node, err := ng.awsManager.buildNodeFromTemplate(ng.asg, template)
	if err != nil {
		return nil, err
	}

	nodeInfo := schedulerframework.NewNodeInfo(cloudprovider.BuildKubeProxy(ng.asg.Name))
	nodeInfo.SetNode(node)
	return nodeInfo, nil
}

// BuildAWS builds AWS cloud provider, manager etc.
func BuildAWS(opts config.AutoscalingOptions, do cloudprovider.NodeGroupDiscoveryOptions, rl *cloudprovider.ResourceLimiter) cloudprovider.CloudProvider {
	var config io.ReadCloser
	if opts.CloudConfig != "" {
		var err error
		config, err = os.Open(opts.CloudConfig)
		if err != nil {
			klog.Fatalf("Couldn't open cloud provider configuration %s: %#v", opts.CloudConfig, err)
		}
		defer config.Close()
	}

	// Generate EC2 list
	instanceTypes, lastUpdateTime := GetStaticEC2InstanceTypes()
	if opts.AWSUseStaticInstanceList {
		klog.Warningf("Use static EC2 Instance Types and list could be outdated. Last update time: %s", lastUpdateTime)
	} else {
		region, err := GetCurrentAwsRegion()
		if err != nil {
			klog.Fatalf("Failed to get AWS Region: %v", err)
		}

		generatedInstanceTypes, err := GenerateEC2InstanceTypes(region)
		if err != nil {
			klog.Fatalf("Failed to generate AWS EC2 Instance Types: %v", err)
		}
		// fallback on the static list if we miss any instance types in the generated output
		// credits to: https://github.com/lyft/cni-ipvlan-vpc-k8s/pull/80
		for k, v := range instanceTypes {
			_, ok := generatedInstanceTypes[k]
			if ok {
				continue
			}
			klog.Infof("Using static instance type %s", k)
			generatedInstanceTypes[k] = v
		}
		instanceTypes = generatedInstanceTypes

		keys := make([]string, 0, len(instanceTypes))
		for key := range instanceTypes {
			keys = append(keys, key)
		}

		klog.Infof("Successfully load %d EC2 Instance Types %s", len(keys), keys)
	}

	manager, err := CreateAwsManager(config, do, instanceTypes)
	if err != nil {
		klog.Fatalf("Failed to create AWS Manager: %v", err)
	}

	provider, err := BuildAwsCloudProvider(manager, rl)
	if err != nil {
		klog.Fatalf("Failed to create AWS cloud provider: %v", err)
	}
	RegisterMetrics()
	return provider
}
<commit_msg>Add support for nvidia-tesla-a100 gpu for aws cloud provider<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 aws

import (
	"fmt"
	"io"
	"os"
	"regexp"
	"strings"

	apiv1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	"k8s.io/autoscaler/cluster-autoscaler/cloudprovider"
	"k8s.io/autoscaler/cluster-autoscaler/config"
	"k8s.io/autoscaler/cluster-autoscaler/utils/errors"
	klog "k8s.io/klog/v2"
	schedulerframework "k8s.io/kubernetes/pkg/scheduler/framework"
)

const (
	// GPULabel is the label added to nodes with GPU resource.
	GPULabel = "k8s.amazonaws.com/accelerator"
)

var (
	availableGPUTypes = map[string]struct{}{
		"nvidia-tesla-k80":  {},
		"nvidia-tesla-p100": {},
		"nvidia-tesla-v100": {},
		"nvidia-tesla-t4":   {},
		"nvidia-tesla-a100": {},
	}
)

// awsCloudProvider implements CloudProvider interface.
type awsCloudProvider struct {
	awsManager      *AwsManager
	resourceLimiter *cloudprovider.ResourceLimiter
}

// BuildAwsCloudProvider builds CloudProvider implementation for AWS.
func BuildAwsCloudProvider(awsManager *AwsManager, resourceLimiter *cloudprovider.ResourceLimiter) (cloudprovider.CloudProvider, error) {
	aws := &awsCloudProvider{
		awsManager:      awsManager,
		resourceLimiter: resourceLimiter,
	}
	return aws, nil
}

// Cleanup stops the go routine that is handling the current view of the ASGs in the form of a cache
func (aws *awsCloudProvider) Cleanup() error {
	aws.awsManager.Cleanup()
	return nil
}

// Name returns name of the cloud provider.
func (aws *awsCloudProvider) Name() string {
	return cloudprovider.AwsProviderName
}

// GPULabel returns the label added to nodes with GPU resource.
func (aws *awsCloudProvider) GPULabel() string {
	return GPULabel
}

// GetAvailableGPUTypes return all available GPU types cloud provider supports
func (aws *awsCloudProvider) GetAvailableGPUTypes() map[string]struct{} {
	return availableGPUTypes
}

// NodeGroups returns all node groups configured for this cloud provider.
func (aws *awsCloudProvider) NodeGroups() []cloudprovider.NodeGroup {
	asgs := aws.awsManager.getAsgs()
	ngs := make([]cloudprovider.NodeGroup, len(asgs))
	for i, asg := range asgs {
		ngs[i] = &AwsNodeGroup{
			asg:        asg,
			awsManager: aws.awsManager,
		}
	}

	return ngs
}

// NodeGroupForNode returns the node group for the given node.
func (aws *awsCloudProvider) NodeGroupForNode(node *apiv1.Node) (cloudprovider.NodeGroup, error) {
	if len(node.Spec.ProviderID) == 0 {
		klog.Warningf("Node %v has no providerId", node.Name)
		return nil, nil
	}
	ref, err := AwsRefFromProviderId(node.Spec.ProviderID)
	if err != nil {
		return nil, err
	}
	asg := aws.awsManager.GetAsgForInstance(*ref)

	if asg == nil {
		return nil, nil
	}

	return &AwsNodeGroup{
		asg:        asg,
		awsManager: aws.awsManager,
	}, nil
}

// Pricing returns pricing model for this cloud provider or error if not available.
func (aws *awsCloudProvider) Pricing() (cloudprovider.PricingModel, errors.AutoscalerError) {
	return nil, cloudprovider.ErrNotImplemented
}

// GetAvailableMachineTypes get all machine types that can be requested from the cloud provider.
func (aws *awsCloudProvider) GetAvailableMachineTypes() ([]string, error) {
	return []string{}, nil
}

// NewNodeGroup builds a theoretical node group based on the node definition provided. The node group is not automatically
// created on the cloud provider side. The node group is not returned by NodeGroups() until it is created.
func (aws *awsCloudProvider) NewNodeGroup(machineType string, labels map[string]string, systemLabels map[string]string,
	taints []apiv1.Taint, extraResources map[string]resource.Quantity) (cloudprovider.NodeGroup, error) {
	return nil, cloudprovider.ErrNotImplemented
}

// GetResourceLimiter returns struct containing limits (max, min) for resources (cores, memory etc.).
func (aws *awsCloudProvider) GetResourceLimiter() (*cloudprovider.ResourceLimiter, error) {
	return aws.resourceLimiter, nil
}

// Refresh is called before every main loop and can be used to dynamically update cloud provider state.
// In particular the list of node groups returned by NodeGroups can change as a result of CloudProvider.Refresh().
func (aws *awsCloudProvider) Refresh() error {
	return aws.awsManager.Refresh()
}

// AwsRef contains a reference to some entity in AWS world.
type AwsRef struct {
	Name string
}

// AwsInstanceRef contains a reference to an instance in the AWS world.
type AwsInstanceRef struct {
	ProviderID string
	Name       string
}

var validAwsRefIdRegex = regexp.MustCompile(fmt.Sprintf(`^aws\:\/\/\/[-0-9a-z]*\/[-0-9a-z]*(\/[-0-9a-z\.]*)?$|aws\:\/\/\/[-0-9a-z]*\/%s.*$`, placeholderInstanceNamePrefix))

// AwsRefFromProviderId creates AwsInstanceRef object from provider id which
// must be in format: aws:///zone/name
func AwsRefFromProviderId(id string) (*AwsInstanceRef, error) {
	if validAwsRefIdRegex.FindStringSubmatch(id) == nil {
		return nil, fmt.Errorf("wrong id: expected format aws:///<zone>/<name>, got %v", id)
	}
	splitted := strings.Split(id[7:], "/")
	return &AwsInstanceRef{
		ProviderID: id,
		Name:       splitted[1],
	}, nil
}

// AwsNodeGroup implements NodeGroup interface.
type AwsNodeGroup struct {
	awsManager *AwsManager
	asg        *asg
}

// MaxSize returns maximum size of the node group.
func (ng *AwsNodeGroup) MaxSize() int {
	return ng.asg.maxSize
}

// MinSize returns minimum size of the node group.
func (ng *AwsNodeGroup) MinSize() int {
	return ng.asg.minSize
}

// TargetSize returns the current TARGET size of the node group. It is possible that the
// number is different from the number of nodes registered in Kubernetes.
func (ng *AwsNodeGroup) TargetSize() (int, error) {
	return ng.asg.curSize, nil
}

// Exist checks if the node group really exists on the cloud provider side. Allows to tell the
// theoretical node group from the real one.
func (ng *AwsNodeGroup) Exist() bool {
	return true
}

// Create creates the node group on the cloud provider side.
func (ng *AwsNodeGroup) Create() (cloudprovider.NodeGroup, error) {
	return nil, cloudprovider.ErrAlreadyExist
}

// Autoprovisioned returns true if the node group is autoprovisioned.
func (ng *AwsNodeGroup) Autoprovisioned() bool {
	return false
}

// Delete deletes the node group on the cloud provider side.
// This will be executed only for autoprovisioned node groups, once their size drops to 0.
func (ng *AwsNodeGroup) Delete() error {
	return cloudprovider.ErrNotImplemented
}

// GetOptions returns NodeGroupAutoscalingOptions that should be used for this particular
// NodeGroup. Returning a nil will result in using default options.
func (ng *AwsNodeGroup) GetOptions(defaults config.NodeGroupAutoscalingOptions) (*config.NodeGroupAutoscalingOptions, error) {
	return nil, cloudprovider.ErrNotImplemented
}

// IncreaseSize increases Asg size
func (ng *AwsNodeGroup) IncreaseSize(delta int) error {
	if delta <= 0 {
		return fmt.Errorf("size increase must be positive")
	}
	size := ng.asg.curSize
	if size+delta > ng.asg.maxSize {
		return fmt.Errorf("size increase too large - desired:%d max:%d", size+delta, ng.asg.maxSize)
	}
	return ng.awsManager.SetAsgSize(ng.asg, size+delta)
}

// DecreaseTargetSize decreases the target size of the node group. This function
// doesn't permit to delete any existing node and can be used only to reduce the
// request for new nodes that have not been yet fulfilled. Delta should be negative.
// It is assumed that cloud provider will not delete the existing nodes if the size
// when there is an option to just decrease the target.
func (ng *AwsNodeGroup) DecreaseTargetSize(delta int) error {
	if delta >= 0 {
		return fmt.Errorf("size decrease size must be negative")
	}

	size := ng.asg.curSize
	nodes, err := ng.awsManager.GetAsgNodes(ng.asg.AwsRef)
	if err != nil {
		return err
	}
	if int(size)+delta < len(nodes) {
		return fmt.Errorf("attempt to delete existing nodes targetSize:%d delta:%d existingNodes: %d",
			size, delta, len(nodes))
	}
	return ng.awsManager.SetAsgSize(ng.asg, size+delta)
}

// Belongs returns true if the given node belongs to the NodeGroup.
func (ng *AwsNodeGroup) Belongs(node *apiv1.Node) (bool, error) {
	ref, err := AwsRefFromProviderId(node.Spec.ProviderID)
	if err != nil {
		return false, err
	}
	targetAsg := ng.awsManager.GetAsgForInstance(*ref)
	if targetAsg == nil {
		return false, fmt.Errorf("%s doesn't belong to a known asg", node.Name)
	}
	if targetAsg.AwsRef != ng.asg.AwsRef {
		return false, nil
	}
	return true, nil
}

// DeleteNodes deletes the nodes from the group.
func (ng *AwsNodeGroup) DeleteNodes(nodes []*apiv1.Node) error {
	size := ng.asg.curSize
	if int(size) <= ng.MinSize() {
		return fmt.Errorf("min size reached, nodes will not be deleted")
	}
	refs := make([]*AwsInstanceRef, 0, len(nodes))
	for _, node := range nodes {
		belongs, err := ng.Belongs(node)
		if err != nil {
			return err
		}
		if belongs != true {
			return fmt.Errorf("%s belongs to a different asg than %s", node.Name, ng.Id())
		}
		awsref, err := AwsRefFromProviderId(node.Spec.ProviderID)
		if err != nil {
			return err
		}
		refs = append(refs, awsref)
	}
	return ng.awsManager.DeleteInstances(refs)
}

// Id returns asg id.
func (ng *AwsNodeGroup) Id() string {
	return ng.asg.Name
}

// Debug returns a debug string for the Asg.
func (ng *AwsNodeGroup) Debug() string {
	return fmt.Sprintf("%s (%d:%d)", ng.Id(), ng.MinSize(), ng.MaxSize())
}

// Nodes returns a list of all nodes that belong to this node group.
func (ng *AwsNodeGroup) Nodes() ([]cloudprovider.Instance, error) {
	asgNodes, err := ng.awsManager.GetAsgNodes(ng.asg.AwsRef)
	if err != nil {
		return nil, err
	}

	instances := make([]cloudprovider.Instance, len(asgNodes))

	for i, asgNode := range asgNodes {
		instances[i] = cloudprovider.Instance{Id: asgNode.ProviderID}
	}
	return instances, nil
}

// TemplateNodeInfo returns a node template for this node group.
func (ng *AwsNodeGroup) TemplateNodeInfo() (*schedulerframework.NodeInfo, error) {
	template, err := ng.awsManager.getAsgTemplate(ng.asg)
	if err != nil {
		return nil, err
	}

	node, err := ng.awsManager.buildNodeFromTemplate(ng.asg, template)
	if err != nil {
		return nil, err
	}

	nodeInfo := schedulerframework.NewNodeInfo(cloudprovider.BuildKubeProxy(ng.asg.Name))
	nodeInfo.SetNode(node)
	return nodeInfo, nil
}

// BuildAWS builds AWS cloud provider, manager etc.
func BuildAWS(opts config.AutoscalingOptions, do cloudprovider.NodeGroupDiscoveryOptions, rl *cloudprovider.ResourceLimiter) cloudprovider.CloudProvider {
	var config io.ReadCloser
	if opts.CloudConfig != "" {
		var err error
		config, err = os.Open(opts.CloudConfig)
		if err != nil {
			klog.Fatalf("Couldn't open cloud provider configuration %s: %#v", opts.CloudConfig, err)
		}
		defer config.Close()
	}

	// Generate EC2 list
	instanceTypes, lastUpdateTime := GetStaticEC2InstanceTypes()
	if opts.AWSUseStaticInstanceList {
		klog.Warningf("Use static EC2 Instance Types and list could be outdated. Last update time: %s", lastUpdateTime)
	} else {
		region, err := GetCurrentAwsRegion()
		if err != nil {
			klog.Fatalf("Failed to get AWS Region: %v", err)
		}

		generatedInstanceTypes, err := GenerateEC2InstanceTypes(region)
		if err != nil {
			klog.Fatalf("Failed to generate AWS EC2 Instance Types: %v", err)
		}
		// fallback on the static list if we miss any instance types in the generated output
		// credits to: https://github.com/lyft/cni-ipvlan-vpc-k8s/pull/80
		for k, v := range instanceTypes {
			_, ok := generatedInstanceTypes[k]
			if ok {
				continue
			}
			klog.Infof("Using static instance type %s", k)
			generatedInstanceTypes[k] = v
		}
		instanceTypes = generatedInstanceTypes

		keys := make([]string, 0, len(instanceTypes))
		for key := range instanceTypes {
			keys = append(keys, key)
		}

		klog.Infof("Successfully load %d EC2 Instance Types %s", len(keys), keys)
	}

	manager, err := CreateAwsManager(config, do, instanceTypes)
	if err != nil {
		klog.Fatalf("Failed to create AWS Manager: %v", err)
	}

	provider, err := BuildAwsCloudProvider(manager, rl)
	if err != nil {
		klog.Fatalf("Failed to create AWS cloud provider: %v", err)
	}
	RegisterMetrics()
	return provider
}
<|endoftext|>


<commit_before>package ultron

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

type (
	// HandleResult type of result handle function
	HandleResult func(*AttackResult)
	// HandleReport type of report handle function
	HandleReport func(map[string]*StatsReport)

	resultHandleChain struct {
		handles []HandleResult
		ch      chan *AttackResult
		wg      sync.WaitGroup
	}

	reportHandleChain struct {
		handles []HandleReport
		ch      chan map[string]*StatsReport
		wg      sync.WaitGroup
	}
)

const (
	// StatsReportInterval time interval
	StatsReportInterval = time.Second * 5
)

var (
	// ResultHandleChain handlers for each request
	ResultHandleChain *resultHandleChain
	// ReportHandleChain handlers for each timing stats report
	ReportHandleChain *reportHandleChain
)

func (rc *resultHandleChain) AddHandle(fn HandleResult) {
	rc.handles = append(rc.handles, fn)
}

func (rc *resultHandleChain) channel() chan *AttackResult {
	return rc.ch
}

func (rc *resultHandleChain) listening() {
	for msg := range rc.ch {
		rc.wg.Add(1)
		go func(ret *AttackResult) {
			defer rc.wg.Done()
			for _, f := range rc.handles {
				f(ret)
			}
		}(msg)
	}
}

func (rc *resultHandleChain) safeClose() {
	rc.wg.Wait()
	close(rc.ch)
}

func (re *reportHandleChain) AddHandle(fn HandleReport) {
	re.handles = append(re.handles, fn)
}

func (re *reportHandleChain) channel() chan map[string]*StatsReport {
	return re.ch
}

func (re *reportHandleChain) listening() {
	for msg := range re.ch {
		re.wg.Add(1)
		go func(s map[string]*StatsReport) {
			defer re.wg.Done()
			for _, h := range re.handles {
				h(s)
			}
		}(msg)
	}
}

func (re *reportHandleChain) safeClose() {
	re.wg.Wait()
	close(re.ch)
}

func printReportToConsole(report map[string]*StatsReport) {
	data, err := json.Marshal(report)
	if err == nil {
		fmt.Println(string(data))
		return
	}
}

func init() {
	ResultHandleChain = &resultHandleChain{
		handles: []HandleResult{defaultStatsCollector.log},
		ch:      make(chan *AttackResult),
	}
	ReportHandleChain = &reportHandleChain{
		handles: []HandleReport{printReportToConsole},
		ch:      make(chan map[string]*StatsReport),
	}
}
<commit_msg>change print format<commit_after>package ultron

import (
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

type (
	// HandleResult type of result handle function
	HandleResult func(*AttackResult)
	// HandleReport type of report handle function
	HandleReport func(map[string]*StatsReport)

	resultHandleChain struct {
		handles []HandleResult
		ch      chan *AttackResult
		wg      sync.WaitGroup
	}

	reportHandleChain struct {
		handles []HandleReport
		ch      chan map[string]*StatsReport
		wg      sync.WaitGroup
	}
)

const (
	// StatsReportInterval time interval
	StatsReportInterval = time.Second * 5
)

var (
	// ResultHandleChain handlers for each request
	ResultHandleChain *resultHandleChain
	// ReportHandleChain handlers for each timing stats report
	ReportHandleChain *reportHandleChain
)

func (rc *resultHandleChain) AddHandle(fn HandleResult) {
	rc.handles = append(rc.handles, fn)
}

func (rc *resultHandleChain) channel() chan *AttackResult {
	return rc.ch
}

func (rc *resultHandleChain) listening() {
	for msg := range rc.ch {
		rc.wg.Add(1)
		go func(ret *AttackResult) {
			defer rc.wg.Done()
			for _, f := range rc.handles {
				f(ret)
			}
		}(msg)
	}
}

func (rc *resultHandleChain) safeClose() {
	rc.wg.Wait()
	close(rc.ch)
}

func (re *reportHandleChain) AddHandle(fn HandleReport) {
	re.handles = append(re.handles, fn)
}

func (re *reportHandleChain) channel() chan map[string]*StatsReport {
	return re.ch
}

func (re *reportHandleChain) listening() {
	for msg := range re.ch {
		re.wg.Add(1)
		go func(s map[string]*StatsReport) {
			defer re.wg.Done()
			for _, h := range re.handles {
				h(s)
			}
		}(msg)
	}
}

func (re *reportHandleChain) safeClose() {
	re.wg.Wait()
	close(re.ch)
}

func printReportToConsole(report map[string]*StatsReport) {
	var full bool
	for _, r := range report {
		if r.FullHistory {
			full = true
		}
		break
	}

	if !full {
		fmt.Printf("|%-24s|%6s|%10s|%10s|%8s|%8s|%8s|%8s|%8s|%8s|%8s|%8s|%8s|\n", "Name", "QPS", "Requests", "Failures", "Min", "Max", "Avg", "Median", "70%", "80%", "90%", "95%", "99%")
		for _, r := range report {
			fmt.Printf("|%-24s|%6d|%10d|%10d|%8d|%8d|%8d|%8d|%8d|%8d|%8d|%8d|%8d|\n", r.Name, r.QPS, r.Requests, r.Failures, r.Min, r.Max, r.Average, r.Median, r.Distributions["0.70"], r.Distributions["0.80"], r.Distributions["0.90"], r.Distributions["0.95"], r.Distributions["0.99"])
		}
		fmt.Printf("\n")
	} else {
		data, err := json.Marshal(report)
		if err == nil {
			fmt.Println(string(data))
		}
	}
}

func init() {
	ResultHandleChain = &resultHandleChain{
		handles: []HandleResult{defaultStatsCollector.log},
		ch:      make(chan *AttackResult),
	}
	ReportHandleChain = &reportHandleChain{
		handles: []HandleReport{printReportToConsole},
		ch:      make(chan map[string]*StatsReport),
	}
}
<|endoftext|>


<commit_before>package main
<commit_msg>Remove empty file handle.go<commit_after><|endoftext|>


<commit_before>package receiver

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	gw "github.com/cvmfs/gateway/internal/gateway"
	"github.com/pkg/errors"
)

// Error is returned by the various receiver commands in case of error
type Error string

func (e Error) Error() string {
	return string(e)
}

// receiverOp is used to identify the different operation performed
// by the cvmfs_receiver process
type receiverOp int32

// The different operations are defined as constants. The numbering
// must match (enum receiver::Request from "cvmfs.git/cvmfs/receiver/reactor.h")
const (
	receiverQuit receiverOp = iota
	receiverEcho
	receiverGenerateToken // Unused
	receiverGetTokenID    // Unused
	receiverCheckToken    // Unused
	receiverSubmitPayload
	receiverCommit
	receiverError // Unused
)

// Receiver contains the operations that "receiver" worker processes perform
type Receiver interface {
	Quit() error
	Echo() error
	SubmitPayload(leasePath string, payload io.Reader, digest string, headerSize int) error
	Commit(leasePath, oldRootHash, newRootHash string, tag gw.RepositoryTag) error
}

// NewReceiver is the factory method for Receiver types
func NewReceiver(execPath string, mock bool) (Receiver, error) {
	if mock {
		return NewMockReceiver()
	}

	return NewCvmfsReceiver(execPath)
}

// CvmfsReceiver spawns an external cvmfs_receiver worker process
type CvmfsReceiver struct {
	worker *exec.Cmd
	stdin  io.Writer
	stdout io.Reader
}

// NewCvmfsReceiver will spawn an external cvmfs_receiver worker process and wait for a command
func NewCvmfsReceiver(execPath string) (*CvmfsReceiver, error) {
	if _, err := os.Stat(execPath); os.IsNotExist(err) {
		return nil, errors.Wrap(err, "worker process executable not found")
	}

	cmd := exec.Command(execPath)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, errors.Wrap(err, "could not get writable pipe")
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, errors.Wrap(err, "could not get readable pipe")
	}

	if err := cmd.Start(); err != nil {
		return nil, errors.Wrap(err, "could not start worker process")
	}

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "start").
		Msgf("worker process ready")

	return &CvmfsReceiver{worker: cmd, stdin: stdin, stdout: stdout}, nil
}

// Quit command is sent to the worker
func (r *CvmfsReceiver) Quit() error {
	if _, err := r.call(receiverQuit, []byte{}, nil); err != nil {
		return errors.Wrap(err, "worker 'quit' call failed")
	}

	if err := r.worker.Wait(); err != nil {
		return errors.Wrap(err, "waiting for worker process failed")
	}

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "quit").
		Msgf("worker process has stopped")

	return nil
}

// Echo command is sent to the worker
func (r *CvmfsReceiver) Echo() error {
	rep, err := r.call(receiverEcho, []byte("Ping"), nil)
	if err != nil {
		return errors.Wrap(err, "worker 'echo' call failed")
	}
	reply := string(rep)

	if !strings.HasPrefix(reply, "PID: ") {
		return fmt.Errorf("invalid 'echo' reply received: %v", reply)
	}

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "echo").
		Msgf("reply: %v", reply)

	return nil
}

// SubmitPayload command is sent to the worker
func (r *CvmfsReceiver) SubmitPayload(leasePath string, payload io.Reader, digest string, headerSize int) error {
	req := map[string]interface{}{"path": leasePath, "digest": digest, "header_size": headerSize}
	buf, err := json.Marshal(&req)
	if err != nil {
		return errors.Wrap(err, "request encoding failed")
	}
	reply, err := r.call(receiverSubmitPayload, buf, payload)
	if err != nil {
		return errors.Wrap(err, "worker 'payload submission' call failed")
	}

	result := toReceiverError(reply)

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "submit payload").
		Msgf("result: %v", result)

	return result
}

// Commit command is sent to the worker
func (r *CvmfsReceiver) Commit(leasePath, oldRootHash, newRootHash string, tag gw.RepositoryTag) error {
	req := map[string]interface{}{
		"lease_path":      leasePath,
		"old_root_hash":   oldRootHash,
		"new_root_hash":   newRootHash,
		"tag_name":        tag.Name,
		"tag_channel":     tag.Channel,
		"tag_description": tag.Description,
	}
	buf, err := json.Marshal(&req)
	if err != nil {
		return errors.Wrap(err, "request encoding failed")
	}

	reply, err := r.call(receiverCommit, buf, nil)
	if err != nil {
		return errors.Wrap(err, "worker 'commit' call failed")
	}

	result := toReceiverError(reply)

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "commit").
		Msgf("result: %v", result)

	return result
}

func (r *CvmfsReceiver) call(reqID receiverOp, msg []byte, payload io.Reader) ([]byte, error) {
	if err := r.request(reqID, msg, payload); err != nil {
		return nil, err
	}
	return r.reply()
}

func (r *CvmfsReceiver) request(reqID receiverOp, msg []byte, payload io.Reader) error {
	if err := binary.Write(r.stdin, binary.LittleEndian, reqID); err != nil {
		return errors.Wrap(err, "could not write request id")
	}
	if err := binary.Write(r.stdin, binary.LittleEndian, int32(len(msg))); err != nil {
		return errors.Wrap(err, "could not write request size")
	}
	if _, err := r.stdin.Write(msg); err != nil {
		return errors.Wrap(err, "could not write request body")
	}
	if payload != nil {
		if _, err := io.Copy(r.stdin, payload); err != nil {
			return errors.Wrap(err, "could not write request payload")
		}
	}
	return nil
}

func (r *CvmfsReceiver) reply() ([]byte, error) {
	var repSize int32
	if err := binary.Read(r.stdout, binary.LittleEndian, &repSize); err != nil {
		return nil, errors.Wrap(err, "could not read reply size")
	}

	reply := make([]byte, repSize)
	reply, err := ioutil.ReadAll(r.stdout)
	if err != nil {
		return nil, errors.Wrap(err, "could not read reply body")
	}

	return reply, nil
}

func toReceiverError(reply []byte) error {
	res := make(map[string]string)
	if err := json.Unmarshal(reply, &res); err != nil {
		return errors.Wrap(err, "could not decode reply")
	}

	if status, ok := res["status"]; ok {
		if status == "ok" {
			return nil
		}

		if reason, ok := res["reason"]; ok {
			return Error(reason)
		}

		return fmt.Errorf("invalid reply")
	}

	return fmt.Errorf("invalid reply")
}
<commit_msg>Fix I/O with cvmfs_receiver process<commit_after>package receiver

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	gw "github.com/cvmfs/gateway/internal/gateway"
	"github.com/pkg/errors"
)

// Error is returned by the various receiver commands in case of error
type Error string

func (e Error) Error() string {
	return string(e)
}

// receiverOp is used to identify the different operation performed
// by the cvmfs_receiver process
type receiverOp int32

// The different operations are defined as constants. The numbering
// must match (enum receiver::Request from "cvmfs.git/cvmfs/receiver/reactor.h")
const (
	receiverQuit receiverOp = iota
	receiverEcho
	receiverGenerateToken // Unused
	receiverGetTokenID    // Unused
	receiverCheckToken    // Unused
	receiverSubmitPayload
	receiverCommit
	receiverError // Unused
)

// Receiver contains the operations that "receiver" worker processes perform
type Receiver interface {
	Quit() error
	Echo() error
	SubmitPayload(leasePath string, payload io.Reader, digest string, headerSize int) error
	Commit(leasePath, oldRootHash, newRootHash string, tag gw.RepositoryTag) error
}

// NewReceiver is the factory method for Receiver types
func NewReceiver(execPath string, mock bool) (Receiver, error) {
	if mock {
		return NewMockReceiver()
	}

	return NewCvmfsReceiver(execPath)
}

// CvmfsReceiver spawns an external cvmfs_receiver worker process
type CvmfsReceiver struct {
	worker *exec.Cmd
	stdin  io.Writer
	stdout io.Reader
}

// NewCvmfsReceiver will spawn an external cvmfs_receiver worker process and wait for a command
func NewCvmfsReceiver(execPath string) (*CvmfsReceiver, error) {
	if _, err := os.Stat(execPath); os.IsNotExist(err) {
		return nil, errors.Wrap(err, "worker process executable not found")
	}

	cmd := exec.Command(execPath)

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return nil, errors.Wrap(err, "could not get writable pipe")
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return nil, errors.Wrap(err, "could not get readable pipe")
	}

	if err := cmd.Start(); err != nil {
		return nil, errors.Wrap(err, "could not start worker process")
	}

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "start").
		Msgf("worker process ready")

	return &CvmfsReceiver{worker: cmd, stdin: stdin, stdout: stdout}, nil
}

// Quit command is sent to the worker
func (r *CvmfsReceiver) Quit() error {
	if _, err := r.call(receiverQuit, []byte{}, nil); err != nil {
		return errors.Wrap(err, "worker 'quit' call failed")
	}

	if err := r.worker.Wait(); err != nil {
		return errors.Wrap(err, "waiting for worker process failed")
	}

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "quit").
		Msgf("worker process has stopped")

	return nil
}

// Echo command is sent to the worker
func (r *CvmfsReceiver) Echo() error {
	rep, err := r.call(receiverEcho, []byte("Ping"), nil)
	if err != nil {
		return errors.Wrap(err, "worker 'echo' call failed")
	}
	reply := string(rep)

	if !strings.HasPrefix(reply, "PID: ") {
		return fmt.Errorf("invalid 'echo' reply received: %v", reply)
	}

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "echo").
		Msgf("reply: %v", reply)

	return nil
}

// SubmitPayload command is sent to the worker
func (r *CvmfsReceiver) SubmitPayload(leasePath string, payload io.Reader, digest string, headerSize int) error {
	req := map[string]interface{}{"path": leasePath, "digest": digest, "header_size": headerSize}
	buf, err := json.Marshal(&req)
	if err != nil {
		return errors.Wrap(err, "request encoding failed")
	}
	reply, err := r.call(receiverSubmitPayload, buf, payload)
	if err != nil {
		return errors.Wrap(err, "worker 'payload submission' call failed")
	}

	result := toReceiverError(reply)

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "submit payload").
		Msgf("result: %v", result)

	return result
}

// Commit command is sent to the worker
func (r *CvmfsReceiver) Commit(leasePath, oldRootHash, newRootHash string, tag gw.RepositoryTag) error {
	req := map[string]interface{}{
		"lease_path":      leasePath,
		"old_root_hash":   oldRootHash,
		"new_root_hash":   newRootHash,
		"tag_name":        tag.Name,
		"tag_channel":     tag.Channel,
		"tag_description": tag.Description,
	}
	buf, err := json.Marshal(&req)
	if err != nil {
		return errors.Wrap(err, "request encoding failed")
	}

	reply, err := r.call(receiverCommit, buf, nil)
	if err != nil {
		return errors.Wrap(err, "worker 'commit' call failed")
	}

	result := toReceiverError(reply)

	gw.Log.Debug().
		Str("component", "receiver").
		Str("command", "commit").
		Msgf("result: %v", result)

	return result
}

func (r *CvmfsReceiver) call(reqID receiverOp, msg []byte, payload io.Reader) ([]byte, error) {
	if err := r.request(reqID, msg, payload); err != nil {
		return nil, err
	}
	return r.reply()
}

func (r *CvmfsReceiver) request(reqID receiverOp, msg []byte, payload io.Reader) error {
	if err := binary.Write(r.stdin, binary.LittleEndian, reqID); err != nil {
		return errors.Wrap(err, "could not write request id")
	}
	if err := binary.Write(r.stdin, binary.LittleEndian, int32(len(msg))); err != nil {
		return errors.Wrap(err, "could not write request size")
	}
	if _, err := r.stdin.Write(msg); err != nil {
		return errors.Wrap(err, "could not write request body")
	}
	if payload != nil {
		if _, err := io.Copy(r.stdin, payload); err != nil {
			return errors.Wrap(err, "could not write request payload")
		}
	}
	return nil
}

func (r *CvmfsReceiver) reply() ([]byte, error) {
	var repSize int32
	if err := binary.Read(r.stdout, binary.LittleEndian, &repSize); err != nil {
		return nil, errors.Wrap(err, "could not read reply size")
	}

	reply := make([]byte, repSize)
	reply, err := ioutil.ReadAll(io.LimitReader(r.stdout, int64(repSize)))
	if err != nil {
		return nil, errors.Wrap(err, "could not read reply body")
	}

	return reply, nil
}

func toReceiverError(reply []byte) error {
	res := make(map[string]string)
	if err := json.Unmarshal(reply, &res); err != nil {
		return errors.Wrap(err, "could not decode reply")
	}

	if status, ok := res["status"]; ok {
		if status == "ok" {
			return nil
		}

		if reason, ok := res["reason"]; ok {
			return Error(reason)
		}

		return fmt.Errorf("invalid reply")
	}

	return fmt.Errorf("invalid reply")
}
<|endoftext|>


<commit_before>package store

import (
	"database/sql"
	"reflect"
)

type SortBy struct {
	Column string
	Asc    bool
}

type Filter interface {
	SQL() string
	Vars() []interface{}
}

type Search struct {
	store   *Store
	i       interface{}
	Sort    []SortBy
	Filters []Filter
}

func (s *Store) Search(i interface{}) *Search {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if _, ok := s.types[typeName(i)]; !ok {
		return nil
	}
	return &Search{
		store: s,
		i:     i,
	}
}

type PreparedSearch struct {
	countStmt *sql.Stmt
	getStmt   *sql.Stmt
	vars      []interface{}
	store     *Store
}

func (s *Search) Prepare() (*PreparedSearch, error) {
	var (
		doneFirst    bool
		sql, sqlVars string
		vars         []interface{}
		name         = typeName(s.i)
	)
	if len(s.Filters) > 0 {
		sql += "WHERE "
		for n, f := range s.Filters {
			if n > 0 {
				sql += "AND "
			}
			sql += f.SQL() + " "
			for _, i := range f.Vars() {
				if v := reflect.ValueOf(i); v.Kind() != reflect.Ptr {
					i = v.Addr().Interface()
				}
				if !isValidType(i) {
					return nil, ErrInvalidType
				}
				vars = append(vars, i)
			}
		}
	}
	s.store.mutex.Lock()
	defer s.store.mutex.Unlock()
	count, err := s.store.db.Prepare("SELECT COUNT(1) FROM [" + name + "] " + sql)
	if err != nil {
		return nil, err
	}
	t := s.store.types[name]
	for pos, f := range t.fields {
		if pos == t.primary {
			continue
		}
		if doneFirst {
			sqlVars += ", "
		} else {
			doneFirst = true
		}
		sqlVars += "[" + f.name + "]"
	}
	if len(s.Sort) > 0 {
		sql += "SORT BY "
		for n, f := range s.Sort {
			if n > 0 {
				sql += ", "
			}
			sql += "[" + f.Column + "]"
			if f.Asc {
				sql += " ASC "
			} else {
				sql += " DESC "
			}
		}
	}
	get, err := s.store.db.Prepare("SELECT [" + t.fields[t.primary].name + "] FROM [" + name + "] " + sql)
	if err != nil {
		return nil, err
	}
	return &PreparedSearch{
		count,
		get,
		vars,
		s.store,
	}, nil
}

func (p *PreparedSearch) Count() (int, error) {
	p.store.mutex.Lock()
	defer p.store.mutex.Unlock()
	row := p.countStmt.QueryRow(p.getVars()...)
	var count int
	err := row.Scan(&count)
	return count, err
}

func (p *PreparedSearch) GetPage(is []interface{}, offset int) (int, error) {
	if len(is) == 0 {
		return 0, nil
	}
	p.store.mutex.Lock()
	defer p.store.mutex.Unlock()
	rows, err := p.getStmt.Query(p.getVars()...)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	return p.store.getPage(is, rows)
}

func (p *PreparedSearch) getVars() []interface{} {
	vars := make([]interface{}, len(p.vars))
	for n, v := range p.vars {
		vars[n] = reflect.ValueOf(v).Elem().Interface()
	}
	return vars
}
<commit_msg>Renames Store.Search to Store.NewSearch<commit_after>package store

import (
	"database/sql"
	"reflect"
)

type SortBy struct {
	Column string
	Asc    bool
}

type Filter interface {
	SQL() string
	Vars() []interface{}
}

type Search struct {
	store   *Store
	i       interface{}
	Sort    []SortBy
	Filters []Filter
}

func (s *Store) NewSearch(i interface{}) *Search {
	s.mutex.Lock()
	defer s.mutex.Unlock()
	if _, ok := s.types[typeName(i)]; !ok {
		return nil
	}
	return &Search{
		store: s,
		i:     i,
	}
}

type PreparedSearch struct {
	countStmt *sql.Stmt
	getStmt   *sql.Stmt
	vars      []interface{}
	store     *Store
}

func (s *Search) Prepare() (*PreparedSearch, error) {
	var (
		doneFirst    bool
		sql, sqlVars string
		vars         []interface{}
		name         = typeName(s.i)
	)
	if len(s.Filters) > 0 {
		sql += "WHERE "
		for n, f := range s.Filters {
			if n > 0 {
				sql += "AND "
			}
			sql += f.SQL() + " "
			for _, i := range f.Vars() {
				if v := reflect.ValueOf(i); v.Kind() != reflect.Ptr {
					i = v.Addr().Interface()
				}
				if !isValidType(i) {
					return nil, ErrInvalidType
				}
				vars = append(vars, i)
			}
		}
	}
	s.store.mutex.Lock()
	defer s.store.mutex.Unlock()
	count, err := s.store.db.Prepare("SELECT COUNT(1) FROM [" + name + "] " + sql)
	if err != nil {
		return nil, err
	}
	t := s.store.types[name]
	for pos, f := range t.fields {
		if pos == t.primary {
			continue
		}
		if doneFirst {
			sqlVars += ", "
		} else {
			doneFirst = true
		}
		sqlVars += "[" + f.name + "]"
	}
	if len(s.Sort) > 0 {
		sql += "SORT BY "
		for n, f := range s.Sort {
			if n > 0 {
				sql += ", "
			}
			sql += "[" + f.Column + "]"
			if f.Asc {
				sql += " ASC "
			} else {
				sql += " DESC "
			}
		}
	}
	get, err := s.store.db.Prepare("SELECT [" + t.fields[t.primary].name + "] FROM [" + name + "] " + sql)
	if err != nil {
		return nil, err
	}
	return &PreparedSearch{
		count,
		get,
		vars,
		s.store,
	}, nil
}

func (p *PreparedSearch) Count() (int, error) {
	p.store.mutex.Lock()
	defer p.store.mutex.Unlock()
	row := p.countStmt.QueryRow(p.getVars()...)
	var count int
	err := row.Scan(&count)
	return count, err
}

func (p *PreparedSearch) GetPage(is []interface{}, offset int) (int, error) {
	if len(is) == 0 {
		return 0, nil
	}
	p.store.mutex.Lock()
	defer p.store.mutex.Unlock()
	rows, err := p.getStmt.Query(p.getVars()...)
	if err != nil {
		return 0, err
	}
	defer rows.Close()
	return p.store.getPage(is, rows)
}

func (p *PreparedSearch) getVars() []interface{} {
	vars := make([]interface{}, len(p.vars))
	for n, v := range p.vars {
		vars[n] = reflect.ValueOf(v).Elem().Interface()
	}
	return vars
}
<|endoftext|>


<commit_before>/* Search functions */

/*
 * Copyright (c) 2013-2016, Jeremy Bingham (<jeremy@goiardi.gl>)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"fmt"
	"github.com/ctdk/goiardi/actor"
	"github.com/ctdk/goiardi/client"
	"github.com/ctdk/goiardi/config"
	"github.com/ctdk/goiardi/databag"
	"github.com/ctdk/goiardi/environment"
	"github.com/ctdk/goiardi/indexer"
	"github.com/ctdk/goiardi/node"
	"github.com/ctdk/goiardi/role"
	"github.com/ctdk/goiardi/search"
	"github.com/ctdk/goiardi/util"
	"github.com/tideland/golib/logger"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
)

func searchHandler(w http.ResponseWriter, r *http.Request) {
	/* ... and we need search to run the environment tests, so here we
	 * go. */
	w.Header().Set("Content-Type", "application/json")
	searchResponse := make(map[string]interface{})
	pathArray := splitPath(r.URL.Path)
	pathArrayLen := len(pathArray)

	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}

	/* set up query params for searching */
	var (
		paramQuery string
		paramsRows int
		sortOrder  string
		start      int
	)
	r.ParseForm()
	if q, found := r.Form["q"]; found {
		if len(q) < 0 {
			jsonErrorReport(w, r, "No query string specified for search", http.StatusBadRequest)
			return
		}
		paramQuery = q[0]
	} else if pathArrayLen != 1 {
		/* default to "*:*" for a search term */
		paramQuery = "*:*"
	}
	if pr, found := r.Form["rows"]; found {
		if len(pr) > 0 {
			paramsRows, _ = strconv.Atoi(pr[0])
		}
	} else {
		paramsRows = 1000
	}
	sortOrder = "id ASC"
	if s, found := r.Form["sort"]; found {
		if len(s) > 0 {
			if s[0] != "" {
				sortOrder = s[0]
			}
		} else {
			sortOrder = "id ASC"
		}
	}
	if st, found := r.Form["start"]; found {
		if len(st) > 0 {
			start, _ = strconv.Atoi(st[0])
		}
	} else {
		start = 0
	}

	var searcher search.Searcher
	if config.Config.PgSearch {
		searcher = &search.PostgresSearch{}
	} else {
		searcher = &search.TrieSearch{}
	}

	if pathArrayLen == 1 {
		/* base end points */
		switch r.Method {
		case "GET":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			searchEndpoints := searcher.GetEndpoints()
			for _, s := range searchEndpoints {
				searchResponse[s] = util.CustomURL(fmt.Sprintf("/search/%s", s))
			}
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else if pathArrayLen == 2 {
		switch r.Method {
		case "GET", "POST":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			var qerr error
			paramQuery, qerr = url.QueryUnescape(paramQuery)
			if qerr != nil {
				jsonErrorReport(w, r, qerr.Error(), http.StatusBadRequest)
				return
			}
			/* start figuring out what comes in POSTS now,
			 * so the partial search tests don't complain
			 * anymore. */
			var partialData map[string]interface{}
			if r.Method == "POST" {
				var perr error
				partialData, perr = parseObjJSON(r.Body)
				if perr != nil {
					jsonErrorReport(w, r, perr.Error(), http.StatusBadRequest)
					return
				}
			}

			idx := pathArray[1]
			res, err := searcher.Search(idx, paramQuery, paramsRows, sortOrder, start, partialData)

			if err != nil {
				statusCode := http.StatusBadRequest
				re := regexp.MustCompile(`^I don't know how to search for .*? data objects.`)
				if re.MatchString(err.Error()) {
					statusCode = http.StatusNotFound
				}
				jsonErrorReport(w, r, err.Error(), statusCode)
				return
			}

			searchResponse["total"] = len(res)
			searchResponse["start"] = start
			searchResponse["rows"] = res
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else {
		/* Say what? Bad request. */
		jsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
		return
	}

	enc := json.NewEncoder(w)
	if err := enc.Encode(&searchResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	reindexResponse := make(map[string]interface{})
	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}
	switch r.Method {
	case "POST":
		if !opUser.IsAdmin() {
			jsonErrorReport(w, r, "You are not allowed to perform that action.", http.StatusForbidden)
			return
		}
		reindexAll()
		reindexResponse["reindex"] = "OK"
	default:
		jsonErrorReport(w, r, "Method not allowed. If you're trying to do something with a data bag named 'reindex', it's not going to work I'm afraid.", http.StatusMethodNotAllowed)
		return
	}
	enc := json.NewEncoder(w)
	if err := enc.Encode(&reindexResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexAll() {
	reindexObjs := make([]indexer.Indexable, 0, 100)
	// We clear the index, *then* do the fetch because if
	// something comes in between the time we fetch the
	// objects to reindex and when it gets done, they'll
	// just be added naturally
	indexer.ClearIndex()

	for _, v := range client.AllClients() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range node.AllNodes() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range role.AllRoles() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range environment.AllEnvironments() {
		reindexObjs = append(reindexObjs, v)
	}
	defaultEnv, _ := environment.Get("_default")
	reindexObjs = append(reindexObjs, defaultEnv)
	// data bags have to be done separately
	dbags := databag.GetList()
	for _, db := range dbags {
		dbag, err := databag.Get(db)
		if err != nil {
			continue
		}
		dbis := make([]indexer.Indexable, dbag.NumDBItems())
		i := 0
		allDBItems, derr := dbag.AllDBItems()
		if derr != nil {
			logger.Errorf(derr.Error())
			continue
		}
		for _, k := range allDBItems {
			n := k
			dbis[i] = n
			i++
		}
		reindexObjs = append(reindexObjs, dbis...)
	}
	indexer.ReIndex(reindexObjs)
	return
}
<commit_msg>add a mutex to reindexing jobs so there won't be multiple reindexing jobs running simultaneously<commit_after>/* Search functions */

/*
 * Copyright (c) 2013-2016, Jeremy Bingham (<jeremy@goiardi.gl>)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package main

import (
	"encoding/json"
	"fmt"
	"github.com/ctdk/goiardi/actor"
	"github.com/ctdk/goiardi/client"
	"github.com/ctdk/goiardi/config"
	"github.com/ctdk/goiardi/databag"
	"github.com/ctdk/goiardi/environment"
	"github.com/ctdk/goiardi/indexer"
	"github.com/ctdk/goiardi/node"
	"github.com/ctdk/goiardi/role"
	"github.com/ctdk/goiardi/search"
	"github.com/ctdk/goiardi/util"
	"github.com/tideland/golib/logger"
	"net/http"
	"net/url"
	"regexp"
	"strconv"
	"sync"
)

var riM *sync.Mutex

func init() {
	riM = new(sync.Mutex)
}

func searchHandler(w http.ResponseWriter, r *http.Request) {
	/* ... and we need search to run the environment tests, so here we
	 * go. */
	w.Header().Set("Content-Type", "application/json")
	searchResponse := make(map[string]interface{})
	pathArray := splitPath(r.URL.Path)
	pathArrayLen := len(pathArray)

	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}

	/* set up query params for searching */
	var (
		paramQuery string
		paramsRows int
		sortOrder  string
		start      int
	)
	r.ParseForm()
	if q, found := r.Form["q"]; found {
		if len(q) < 0 {
			jsonErrorReport(w, r, "No query string specified for search", http.StatusBadRequest)
			return
		}
		paramQuery = q[0]
	} else if pathArrayLen != 1 {
		/* default to "*:*" for a search term */
		paramQuery = "*:*"
	}
	if pr, found := r.Form["rows"]; found {
		if len(pr) > 0 {
			paramsRows, _ = strconv.Atoi(pr[0])
		}
	} else {
		paramsRows = 1000
	}
	sortOrder = "id ASC"
	if s, found := r.Form["sort"]; found {
		if len(s) > 0 {
			if s[0] != "" {
				sortOrder = s[0]
			}
		} else {
			sortOrder = "id ASC"
		}
	}
	if st, found := r.Form["start"]; found {
		if len(st) > 0 {
			start, _ = strconv.Atoi(st[0])
		}
	} else {
		start = 0
	}

	var searcher search.Searcher
	if config.Config.PgSearch {
		searcher = &search.PostgresSearch{}
	} else {
		searcher = &search.TrieSearch{}
	}

	if pathArrayLen == 1 {
		/* base end points */
		switch r.Method {
		case "GET":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			searchEndpoints := searcher.GetEndpoints()
			for _, s := range searchEndpoints {
				searchResponse[s] = util.CustomURL(fmt.Sprintf("/search/%s", s))
			}
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else if pathArrayLen == 2 {
		switch r.Method {
		case "GET", "POST":
			if opUser.IsValidator() {
				jsonErrorReport(w, r, "You are not allowed to perform this action", http.StatusForbidden)
				return
			}
			var qerr error
			paramQuery, qerr = url.QueryUnescape(paramQuery)
			if qerr != nil {
				jsonErrorReport(w, r, qerr.Error(), http.StatusBadRequest)
				return
			}
			/* start figuring out what comes in POSTS now,
			 * so the partial search tests don't complain
			 * anymore. */
			var partialData map[string]interface{}
			if r.Method == "POST" {
				var perr error
				partialData, perr = parseObjJSON(r.Body)
				if perr != nil {
					jsonErrorReport(w, r, perr.Error(), http.StatusBadRequest)
					return
				}
			}

			idx := pathArray[1]
			res, err := searcher.Search(idx, paramQuery, paramsRows, sortOrder, start, partialData)

			if err != nil {
				statusCode := http.StatusBadRequest
				re := regexp.MustCompile(`^I don't know how to search for .*? data objects.`)
				if re.MatchString(err.Error()) {
					statusCode = http.StatusNotFound
				}
				jsonErrorReport(w, r, err.Error(), statusCode)
				return
			}

			searchResponse["total"] = len(res)
			searchResponse["start"] = start
			searchResponse["rows"] = res
		default:
			jsonErrorReport(w, r, "Method not allowed", http.StatusMethodNotAllowed)
			return
		}
	} else {
		/* Say what? Bad request. */
		jsonErrorReport(w, r, "Bad request", http.StatusBadRequest)
		return
	}

	enc := json.NewEncoder(w)
	if err := enc.Encode(&searchResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")
	reindexResponse := make(map[string]interface{})
	opUser, oerr := actor.GetReqUser(r.Header.Get("X-OPS-USERID"))
	if oerr != nil {
		jsonErrorReport(w, r, oerr.Error(), oerr.Status())
		return
	}
	switch r.Method {
	case "POST":
		if !opUser.IsAdmin() {
			jsonErrorReport(w, r, "You are not allowed to perform that action.", http.StatusForbidden)
			return
		}
		reindexAll()
		reindexResponse["reindex"] = "OK"
	default:
		jsonErrorReport(w, r, "Method not allowed. If you're trying to do something with a data bag named 'reindex', it's not going to work I'm afraid.", http.StatusMethodNotAllowed)
		return
	}
	enc := json.NewEncoder(w)
	if err := enc.Encode(&reindexResponse); err != nil {
		jsonErrorReport(w, r, err.Error(), http.StatusInternalServerError)
	}
}

func reindexAll() {
	// Take the mutex before starting to reindex everything. This way at
	// least reindexing jobs won't pile up on top of each other all trying
	// to execute simultaneously.
	riM.Lock()
	defer riM.Unlock()

	reindexObjs := make([]indexer.Indexable, 0, 100)
	// We clear the index, *then* do the fetch because if
	// something comes in between the time we fetch the
	// objects to reindex and when it gets done, they'll
	// just be added naturally
	indexer.ClearIndex()

	for _, v := range client.AllClients() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range node.AllNodes() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range role.AllRoles() {
		reindexObjs = append(reindexObjs, v)
	}
	for _, v := range environment.AllEnvironments() {
		reindexObjs = append(reindexObjs, v)
	}
	defaultEnv, _ := environment.Get("_default")
	reindexObjs = append(reindexObjs, defaultEnv)
	// data bags have to be done separately
	dbags := databag.GetList()
	for _, db := range dbags {
		dbag, err := databag.Get(db)
		if err != nil {
			continue
		}
		dbis := make([]indexer.Indexable, dbag.NumDBItems())
		i := 0
		allDBItems, derr := dbag.AllDBItems()
		if derr != nil {
			logger.Errorf(derr.Error())
			continue
		}
		for _, k := range allDBItems {
			n := k
			dbis[i] = n
			i++
		}
		reindexObjs = append(reindexObjs, dbis...)
	}
	indexer.ReIndex(reindexObjs)
	return
}
<|endoftext|>


<commit_before>package chevalier

import (
	"github.com/anchor/elastigo/api"
	es "github.com/anchor/elastigo/core"

	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"
)

// QueryEngine presents an interface for running queries for sources
// against Elasticsearch.
type QueryEngine struct {
	indexName      string
	metaIndex      string
	originType     string
	dataType       string
	nSources       int64
	updateInterval time.Duration
}

// NewQueryEngine initializes a QueryEngine with the supplied
// Elasticsearch metadata. indexName and dataType can be anything as
// long as they're consistent.
func NewQueryEngine(host, indexName, dataType, metaIndex string) *QueryEngine {
	e := new(QueryEngine)
	e.indexName = indexName
	e.metaIndex = metaIndex
	e.dataType = dataType
	e.originType = "chevalier_origin"
	api.Domain = host
	e.updateSourceCount()
	e.updateInterval = time.Second * 10
	go e.updateForever()
	return e
}

// sanitizeField takes an input field specifier (from a
// SourceRequest_Tag) and munges it so Elasticsearch likes it more. In
// particular, it makes single-wildcard queries work correctly.
//
// FIXME: mid-field wildcards still aren't working.
func (e *QueryEngine) sanitizeField(f string) string {
	f = strings.TrimSpace(f)
	if f == "*" {
		return fmt.Sprintf("%s._all", e.dataType)
	}
	f = fmt.Sprintf("%s.%s", e.dataType, f)
	return f
}

// sanitizeTag escapes assorted Elasticsearch wildcards.
//
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters
func (e *QueryEngine) sanitizeTag(field, value string) (string, string) {
	// * is normally in this list, but is not included here because
	// we want it to act as a wildcard.
	// Also, this can be made a lot faster.
	reservedChars := `\ + - && || ! ( ) { } [ ] ^ " ~ ? : /`
	for _, char := range strings.Split(reservedChars, " ") {
		escapedChar := fmt.Sprintf(`\%s`, char)
		field = strings.Replace(field, char, escapedChar, -1)
		value = strings.Replace(value, char, escapedChar, -1)
	}
	field = e.sanitizeField(field)
	return field, value
}

// buildTagQuery constructs an elastigo query object (search.QueryDsl)
// from a SourceRequest_Tag, designed to be plugged into a
// query-string-type[0] query later on. Returns error on empty query.
//
// [0]: http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html
func (e *QueryEngine) buildTagQuery(tag *SourceRequest_Tag) (map[string]interface{}, error) {
	field, value := e.sanitizeTag(*tag.Field, *tag.Value)
	// Don't bother running empty queries.
	if value == "" {
		return nil, errors.New("empty query string")
	}
	fields := make([]string, 1)
	fields[0] = field
	qs := map[string]interface{} {
		"query_string" : map[string]interface{} {
			"query" : value,
			"fields" : fields,
			"analyzer" : "keyword",
		},
	}
	return qs, nil
}

// buildOriginQuery returns the origin-matching part of the complete
// query.
func (e *QueryEngine) buildOriginQuery(origin string) map[string]interface{} {
	fields := make([]string, 1)
	fields[0] = "Origin"
	qs := map[string]interface{} {
		"query_string" : map[string]interface{} {
			"query" : origin,
			"fields" : fields,
			"analyzer" : "keyword",
		},
	}
	return qs
}

// SourceQuery is a multi-level map type representing an Elasticsearch
// query-string-type query. Suitable for marshalling as JSON and feeding
// to Elasticsearch.
type SourceQuery map[string]interface{}

func (e *QueryEngine) getStartResult(req *SourceRequest) int64 {
	startPage := req.GetStartPage()
	pageSize := req.GetSourcesPerPage()
	if pageSize == 0 {
		return int64(0)
	}
	return startPage * pageSize
}

// getResultCount returns the number of results to return for a
// request.
func (e *QueryEngine) getResultCount(req *SourceRequest) int64 {
	pageSize := req.GetSourcesPerPage()
	if pageSize <= 0 {
		return e.nSources
	}
	return pageSize
}

// BuildQuery takes a SourceRequest and turns it into a multi-level
// map suitable for marshalling to JSON and sending to Elasticsearch.
func (e *QueryEngine) BuildQuery(origin string, req *SourceRequest) (SourceQuery, error) {
	fromResult := e.getStartResult(req)
	resultCount := e.getResultCount(req)
	// First, we check if the Address field is set; if it is,
	// we can ignore the rest of the request.
	if req.Address != nil {
		query := map[string]interface{}{
			"query": map[string]interface{}{
				"term": map[string]uint64 {
					"Address" : *(req.Address),
				},
			},
			"from": fromResult,
			"size": resultCount,
		}
		return SourceQuery(query), nil
	}
	// Next, we check if the query_string field is set; if it is,
	// we can ignore the rest of the request.
	qs := req.GetQueryString()
	if qs != "" {
		query := map[string]interface{}{
			"query": map[string]interface{}{
				"query_string": map[string]interface{}{
					"query": qs,
					"analyzer": "keyword",
				},
			},
			"from": fromResult,
			"size": resultCount,
		}
		return SourceQuery(query), nil
	}
	tags := req.GetTags()
	tagQueries := make([]*map[string]interface{}, 0)
	for _, tag := range tags {
		q, err := e.buildTagQuery(tag)
		if q != nil && err == nil {
			tagQueries = append(tagQueries, &q)
		}
	}
	originTag := e.buildOriginQuery(origin)
	tagQueries = append(tagQueries, &originTag)
	if len(tagQueries) == 0 {
		return nil, errors.New("No valid query strings found.")
	}
	query := map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": tagQueries,
			},
		},
		"from": fromResult,
		"size": resultCount,
	}
	return SourceQuery(query), nil
}

// runSourceRequest takes a request object and returns an elastigo-type
// (i.e., intermediate) result.
func (e *QueryEngine) runSourceRequest(origin string, req *SourceRequest) (*es.SearchResult, error) {
	q, err := e.BuildQuery(origin, req)
	if err != nil {
		return nil, err
	}
	var args map[string]interface{}
	res, err := es.SearchRequest(e.indexName, e.dataType, args, q)
	return &res, err
}

// GetSources takes a request object and returns the DataSourceBurst of
// the sources it gets back from Elasticsearch. If error is not nil,
// then a valid DataSourceBurst will still be returned, with the Error
// field set.
func (e *QueryEngine) GetSources(origin string, req *SourceRequest) (*DataSourceBurst, error) {
	res, err := e.runSourceRequest(origin, req)
	if err != nil {
		msg := fmt.Sprintf("Request error: %v", err)
		burst := new(DataSourceBurst)
		burst.Error = &msg
		return burst, nil
	}
	sources := make([]*DataSource, len(res.Hits.Hits))
	for i, hit := range res.Hits.Hits {
		source := new(ElasticsearchSource)
		err = json.Unmarshal(*hit.Source, source)
		if err != nil {
			msg := fmt.Sprintf("Response decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, nil
		}
		sources[i], err = source.Unmarshal()
		if err != nil {
			msg := fmt.Sprintf("Elasticsearch object decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, nil
		}
	}
	burst := BuildSourceBurst(sources)
	return burst, nil
}

// FmtResult returns a string from a SearchResult by interpreting it in
// the most naive manner possible. For debugging.
func FmtResult(result *es.SearchResult) []string {
	results := make([]string, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		results[i] = string((*hit.Source)[:])
	}
	return results
}

// GetSourceCount returns the number of sources we currently think exist
// in the index.
func (e *QueryEngine) GetSourceCount() int64 {
	return e.nSources
}

// updateSourceCount updates our running total of documents-in-index
// (by asking Elasticsearch).
func (e *QueryEngine) updateSourceCount() error {
	var args map[string]interface{}
	resp, err := es.Count(e.indexName, e.dataType, args)
	e.nSources = int64(resp.Count)
	return err
}

// updateForever updates the source counter on a regular basis.
func (e *QueryEngine) updateForever() {
	for {
		time.Sleep(e.updateInterval)
		err := e.updateSourceCount()
		if err != nil {
			log.Printf("Error updating source count: %v", err)
		}
	}
}
<commit_msg>This should be a field query<commit_after>package chevalier

import (
	"github.com/anchor/elastigo/api"
	es "github.com/anchor/elastigo/core"

	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"
	"time"
)

// QueryEngine presents an interface for running queries for sources
// against Elasticsearch.
type QueryEngine struct {
	indexName      string
	metaIndex      string
	originType     string
	dataType       string
	nSources       int64
	updateInterval time.Duration
}

// NewQueryEngine initializes a QueryEngine with the supplied
// Elasticsearch metadata. indexName and dataType can be anything as
// long as they're consistent.
func NewQueryEngine(host, indexName, dataType, metaIndex string) *QueryEngine {
	e := new(QueryEngine)
	e.indexName = indexName
	e.metaIndex = metaIndex
	e.dataType = dataType
	e.originType = "chevalier_origin"
	api.Domain = host
	e.updateSourceCount()
	e.updateInterval = time.Second * 10
	go e.updateForever()
	return e
}

// sanitizeField takes an input field specifier (from a
// SourceRequest_Tag) and munges it so Elasticsearch likes it more. In
// particular, it makes single-wildcard queries work correctly.
//
// FIXME: mid-field wildcards still aren't working.
func (e *QueryEngine) sanitizeField(f string) string {
	f = strings.TrimSpace(f)
	if f == "*" {
		return fmt.Sprintf("%s._all", e.dataType)
	}
	f = fmt.Sprintf("%s.%s", e.dataType, f)
	return f
}

// sanitizeTag escapes assorted Elasticsearch wildcards.
//
// http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html#_reserved_characters
func (e *QueryEngine) sanitizeTag(field, value string) (string, string) {
	// * is normally in this list, but is not included here because
	// we want it to act as a wildcard.
	// Also, this can be made a lot faster.
	reservedChars := `\ + - && || ! ( ) { } [ ] ^ " ~ ? : /`
	for _, char := range strings.Split(reservedChars, " ") {
		escapedChar := fmt.Sprintf(`\%s`, char)
		field = strings.Replace(field, char, escapedChar, -1)
		value = strings.Replace(value, char, escapedChar, -1)
	}
	field = e.sanitizeField(field)
	return field, value
}

// buildTagQuery constructs an elastigo query object (search.QueryDsl)
// from a SourceRequest_Tag, designed to be plugged into a
// query-string-type[0] query later on. Returns error on empty query.
//
// [0]: http://www.elasticsearch.org/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html
func (e *QueryEngine) buildTagQuery(tag *SourceRequest_Tag) (map[string]interface{}, error) {
	field, value := e.sanitizeTag(*tag.Field, *tag.Value)
	// Don't bother running empty queries.
	if value == "" {
		return nil, errors.New("empty query string")
	}
	fields := make([]string, 1)
	fields[0] = field
	qs := map[string]interface{} {
		"query_string" : map[string]interface{} {
			"query" : value,
			"fields" : fields,
			"analyzer" : "keyword",
		},
	}
	return qs, nil
}

// buildOriginQuery returns the origin-matching part of the complete
// query.
func (e *QueryEngine) buildOriginQuery(origin string) map[string]interface{} {
	fields := make([]string, 1)
	fields[0] = "Origin"
	qs := map[string]interface{} {
		"query_string" : map[string]interface{} {
			"query" : origin,
			"fields" : fields,
			"analyzer" : "keyword",
		},
	}
	return qs
}

// SourceQuery is a multi-level map type representing an Elasticsearch
// query-string-type query. Suitable for marshalling as JSON and feeding
// to Elasticsearch.
type SourceQuery map[string]interface{}

func (e *QueryEngine) getStartResult(req *SourceRequest) int64 {
	startPage := req.GetStartPage()
	pageSize := req.GetSourcesPerPage()
	if pageSize == 0 {
		return int64(0)
	}
	return startPage * pageSize
}

// getResultCount returns the number of results to return for a
// request.
func (e *QueryEngine) getResultCount(req *SourceRequest) int64 {
	pageSize := req.GetSourcesPerPage()
	if pageSize <= 0 {
		return e.nSources
	}
	return pageSize
}

// BuildQuery takes a SourceRequest and turns it into a multi-level
// map suitable for marshalling to JSON and sending to Elasticsearch.
func (e *QueryEngine) BuildQuery(origin string, req *SourceRequest) (SourceQuery, error) {
	fromResult := e.getStartResult(req)
	resultCount := e.getResultCount(req)
	// First, we check if the Address field is set; if it is,
	// we can ignore the rest of the request.
	if req.Address != nil {
		query := map[string]interface{}{
			"query": map[string]interface{}{
				"field": map[string]interface{} {
					"Address" : *(req.Address),
					"analyzer" : "keyword",
				},
			},
			"from": fromResult,
			"size": resultCount,
		}
		return SourceQuery(query), nil
	}
	// Next, we check if the query_string field is set; if it is,
	// we can ignore the rest of the request.
	qs := req.GetQueryString()
	if qs != "" {
		query := map[string]interface{}{
			"query": map[string]interface{}{
				"query_string": map[string]interface{}{
					"query": qs,
					"analyzer": "keyword",
				},
			},
			"from": fromResult,
			"size": resultCount,
		}
		return SourceQuery(query), nil
	}
	tags := req.GetTags()
	tagQueries := make([]*map[string]interface{}, 0)
	for _, tag := range tags {
		q, err := e.buildTagQuery(tag)
		if q != nil && err == nil {
			tagQueries = append(tagQueries, &q)
		}
	}
	originTag := e.buildOriginQuery(origin)
	tagQueries = append(tagQueries, &originTag)
	if len(tagQueries) == 0 {
		return nil, errors.New("No valid query strings found.")
	}
	query := map[string]interface{}{
		"query": map[string]interface{}{
			"bool": map[string]interface{}{
				"must": tagQueries,
			},
		},
		"from": fromResult,
		"size": resultCount,
	}
	return SourceQuery(query), nil
}

// runSourceRequest takes a request object and returns an elastigo-type
// (i.e., intermediate) result.
func (e *QueryEngine) runSourceRequest(origin string, req *SourceRequest) (*es.SearchResult, error) {
	q, err := e.BuildQuery(origin, req)
	if err != nil {
		return nil, err
	}
	var args map[string]interface{}
	res, err := es.SearchRequest(e.indexName, e.dataType, args, q)
	return &res, err
}

// GetSources takes a request object and returns the DataSourceBurst of
// the sources it gets back from Elasticsearch. If error is not nil,
// then a valid DataSourceBurst will still be returned, with the Error
// field set.
func (e *QueryEngine) GetSources(origin string, req *SourceRequest) (*DataSourceBurst, error) {
	res, err := e.runSourceRequest(origin, req)
	if err != nil {
		msg := fmt.Sprintf("Request error: %v", err)
		burst := new(DataSourceBurst)
		burst.Error = &msg
		return burst, nil
	}
	sources := make([]*DataSource, len(res.Hits.Hits))
	for i, hit := range res.Hits.Hits {
		source := new(ElasticsearchSource)
		err = json.Unmarshal(*hit.Source, source)
		if err != nil {
			msg := fmt.Sprintf("Response decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, nil
		}
		sources[i], err = source.Unmarshal()
		if err != nil {
			msg := fmt.Sprintf("Elasticsearch object decoding error: %v", err)
			burst := new(DataSourceBurst)
			burst.Error = &msg
			return burst, nil
		}
	}
	burst := BuildSourceBurst(sources)
	return burst, nil
}

// FmtResult returns a string from a SearchResult by interpreting it in
// the most naive manner possible. For debugging.
func FmtResult(result *es.SearchResult) []string {
	results := make([]string, len(result.Hits.Hits))
	for i, hit := range result.Hits.Hits {
		results[i] = string((*hit.Source)[:])
	}
	return results
}

// GetSourceCount returns the number of sources we currently think exist
// in the index.
func (e *QueryEngine) GetSourceCount() int64 {
	return e.nSources
}

// updateSourceCount updates our running total of documents-in-index
// (by asking Elasticsearch).
func (e *QueryEngine) updateSourceCount() error {
	var args map[string]interface{}
	resp, err := es.Count(e.indexName, e.dataType, args)
	e.nSources = int64(resp.Count)
	return err
}

// updateForever updates the source counter on a regular basis.
func (e *QueryEngine) updateForever() {
	for {
		time.Sleep(e.updateInterval)
		err := e.updateSourceCount()
		if err != nil {
			log.Printf("Error updating source count: %v", err)
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
)

const (
	evalInf = 50000
)

var (
	errCheckmate checkmateError
	errStalemate = errors.New("stalemate")
	errFiftyMove = errors.New("fifty-move rule")
	// TODO: threefold repetition
)

// A checkmateError indicates the number of plies until a forced checkmate can be delivered.
// The zero value of type checkmateError indicates that the current position is checkmate.
type checkmateError int

func (n checkmateError) Error() string {
	if n&1 == 0 {
		return fmt.Sprintf("-#%d", n/2)
	}
	return fmt.Sprintf("#%d", (n+1)/2)
}

// Prev returns the checkmateError corresponding to the previous ply.
func (n checkmateError) Prev() checkmateError { return n + 1 }

// SearchPosition searches a Position to the specified depth via iterative deepening
// and returns the search results.
func SearchPosition(ctx context.Context, pos Position, depth int) Results {
	var results Results
	s := Search{
		allowCutoff: true,
		counters:    make([]int, depth+1),
	}
	for d := 1; d <= depth; d++ {
		_, results, _ = s.negamax(ctx, pos, results, Window{-evalInf, evalInf}, d)
		if ctx.Err() != nil {
			break
		}
	}
	return results
}

type Search struct {
	allowCutoff bool
	counters    []int
}

// negamax recursively searches a Position to the specified depth and returns the evaluation score
// relative to the side to move and the search results. It employs alpha-beta pruning outside of
// the specified Window. If recommended is zero length, negamax will generate and search all legal
// moves; if recommended moves are provided, they must all be legal, and only they will be searched.
func (s *Search) negamax(
	ctx context.Context,
	pos Position,
	recommended Results,
	w Window,
	depth int,
) (bestScore RelScore, results Results, err error) {
	s.counters[len(s.counters)-1-depth]++

	score, recommended, err := checkDone(pos, recommended)
	if err != nil {
		return score, nil, err
	}
	// Invariant: len(recommended) > 0 and recommended contains only legal moves
	if depth == 0 {
		score, err := Eval(pos)
		return score.Rel(pos.ToMove), nil, err
	}

	results = recommended
	defer func() { results.SortFor(pos.ToMove) }()

	for _, r := range recommended {
		if r.err != nil {
			// The move is already known not to avoid a game-ending state; no need to search it further
			results = results.Update(r)
			continue
		}

		score, cont, err := s.negamax(ctx, Make(pos, r.move), r.cont, w.Neg(), depth-1)
		score *= -1

		results = results.Update(Result{move: r.move, score: score.Abs(pos.ToMove), depth: depth - 1, cont: cont, err: err})

		if depth >= 3 && ctx.Err() != nil {
			break
		}

		var constrained, ok bool
		w, constrained, ok = w.Constrain(score)
		if constrained {
			// improved lower bound
		}
		if !ok && s.allowCutoff {
			break
		}
	}
	if err, ok := results[0].err.(checkmateError); ok {
		return w.alpha, results, err.Prev()
	}
	return w.alpha, results, results[0].err
}

// IsPseudoLegal returns whether a Move is pseudo-legal in a Position.
// A move is pseudo-legal if the square to be moved from contains the specified piece
// and the piece is capable of moving to the target square if doing so would not put the king in check.
func IsPseudoLegal(pos Position, move Move) bool {
	for _, m := range PieceMoves[move.Piece](pos) {
		if m == move {
			return true
		}
	}
	return false
}

// IsLegal returns whether a Position is legal.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool { return !anyLegal(pos, Candidates(pos)) }

// anyLegal returns whether any of the given Moves are legal in the given Position.
func anyLegal(pos Position, moves []Move) bool {
	for _, m := range moves {
		if IsLegal(Make(pos, m)) {
			return true
		}
	}
	return false
}

func checkDone(pos Position, recommended Results) (RelScore, Results, error) {
	if len(recommended) == 0 {
		moves := Candidates(pos) // pseudo-legal
		recommended = make(Results, 0, len(moves))
		for _, m := range moves {
			if !IsLegal(Make(pos, m)) {
				continue
			}
			recommended = append(recommended, Result{move: m})
		}
	}
	if len(recommended) == 0 {
		// no legal moves
		if IsCheck(pos) {
			return -evalInf, nil, errCheckmate
		}
		return 0, nil, errStalemate
	}
	if pos.HalfMove == 100 {
		// fifty-move rule
		return 0, nil, errFiftyMove
	}
	return 0, recommended, nil
}

// A Result holds a searched Move along with its evaluated Score,
// the search depth, and the continuation Results of the search.
type Result struct {
	move  Move
	score Score
	depth int
	cont  Results
	err   error
}

// String returns a string representation of r, including its principal variation.
func (r Result) String() string {
	s := fmt.Sprintf("%v (%v) %v", r.score, r.depth, r.PV())
	if r.err != nil {
		return s + " " + r.err.Error()
	}
	return s
}

// PV returns a string representation of r's principal variation.
func (r Result) PV() string {
	if r.depth == 0 || len(r.cont) == 0 {
		return LongAlgebraic(r.move)
	}
	return LongAlgebraic(r.move) + " " + r.cont[0].PV()
}

// A Results contains the results of a search.
// Functions returning a Results should sort it before returning.
type Results []Result

// SortFor sorts rs beginning with the best move for c.
// The sort is not guaranteed to be stable.
func (rs Results) SortFor(c Color) {
	if c == White {
		// Sort first by mate condition, then by depth decreasing, and then by Score decreasing
		sort.Slice(rs, func(i, j int) bool {
			if less, ok := rs.mateSort(i, j); ok {
				return less
			}
			return rs[i].depth > rs[j].depth || rs[i].depth == rs[j].depth && rs[i].score > rs[j].score
		})
	} else {
		// Sort first by mate condition, then by depth decreasing, and then by Score increasing
		sort.Slice(rs, func(i, j int) bool {
			if less, ok := rs.mateSort(i, j); ok {
				return less
			}
			return rs[i].depth > rs[j].depth || rs[i].depth == rs[j].depth && rs[i].score < rs[j].score
		})
	}
}

// mateSort reports how two Result elements should be sorted if at least one of them leads to checkmate.
// ok reports whether either Result contains an error of type checkmateError.
// If ok is true, less reports whether the Result with index i should sort before the Result with index j.
// If ok is false, the value of less is meaningless.
func (rs Results) mateSort(i, j int) (less bool, ok bool) {
	ich, iok := rs[i].err.(checkmateError)
	jch, jok := rs[j].err.(checkmateError)
	switch {
	case iok && jok:
		switch iwin, jwin := ich&1 == 0, jch&1 == 0; {
		case iwin && jwin:
			return ich < jch, true
		case !iwin && !jwin:
			return ich > jch, true
		default:
			return iwin, true
		}
	case iok:
		return ich&1 == 0, true
	case jok:
		return jch&1 != 0, true
	}
	return false, false
}

// Update looks for a Result in rs with the same Move as r and replaces it with r,
// or appends r if there is none, and returns the possibly updated slice.
func (rs Results) Update(r Result) Results {
	for i := range rs {
		if rs[i].move == r.move {
			rs[i] = r
			return rs
		}
	}
	return append(rs, r)
}

// String returns a string representation of all Result values in r.
func (rs Results) String() string {
	var s string
	for _, r := range rs {
		s += fmt.Sprintf("%v\n", r.String())
	}
	return s
}

// A Window represents the bounds of a position's evaluation.
type Window struct{ alpha, beta RelScore }

// Constrain updates the lower bound of w, if applicable, and returns the updated window,
// whether the lower bound was changed, and whether the returned Window remains valid.
func (w Window) Constrain(n RelScore) (c Window, constrained bool, ok bool) {
	if n <= w.alpha {
		return w, false, true
	}
	return Window{n, w.beta}, true, n <= w.beta
}

// Neg returns the additive inverse of w.
func (w Window) Neg() Window { return Window{-w.beta, -w.alpha} }
<commit_msg>refine SortFor by Squares<commit_after>package main

import (
	"context"
	"errors"
	"fmt"
	"sort"
)

const (
	evalInf = 50000
)

var (
	errCheckmate checkmateError
	errStalemate = errors.New("stalemate")
	errFiftyMove = errors.New("fifty-move rule")
	// TODO: threefold repetition
)

// A checkmateError indicates the number of plies until a forced checkmate can be delivered.
// The zero value of type checkmateError indicates that the current position is checkmate.
type checkmateError int

func (n checkmateError) Error() string {
	if n&1 == 0 {
		return fmt.Sprintf("-#%d", n/2)
	}
	return fmt.Sprintf("#%d", (n+1)/2)
}

// Prev returns the checkmateError corresponding to the previous ply.
func (n checkmateError) Prev() checkmateError { return n + 1 }

// SearchPosition searches a Position to the specified depth via iterative deepening
// and returns the search results.
func SearchPosition(ctx context.Context, pos Position, depth int) Results {
	var results Results
	s := Search{
		allowCutoff: true,
		counters:    make([]int, depth+1),
	}
	for d := 1; d <= depth; d++ {
		_, results, _ = s.negamax(ctx, pos, results, Window{-evalInf, evalInf}, d)
		if ctx.Err() != nil {
			break
		}
	}
	return results
}

type Search struct {
	allowCutoff bool
	counters    []int
}

// negamax recursively searches a Position to the specified depth and returns the evaluation score
// relative to the side to move and the search results. It employs alpha-beta pruning outside of
// the specified Window. If recommended is zero length, negamax will generate and search all legal
// moves; if recommended moves are provided, they must all be legal, and only they will be searched.
func (s *Search) negamax(
	ctx context.Context,
	pos Position,
	recommended Results,
	w Window,
	depth int,
) (bestScore RelScore, results Results, err error) {
	s.counters[len(s.counters)-1-depth]++

	score, recommended, err := checkDone(pos, recommended)
	if err != nil {
		return score, nil, err
	}
	// Invariant: len(recommended) > 0 and recommended contains only legal moves
	if depth == 0 {
		score, err := Eval(pos)
		return score.Rel(pos.ToMove), nil, err
	}

	results = recommended
	defer func() { results.SortFor(pos.ToMove) }()

	for _, r := range recommended {
		if r.err != nil {
			// The move is already known not to avoid a game-ending state; no need to search it further
			results = results.Update(r)
			continue
		}

		score, cont, err := s.negamax(ctx, Make(pos, r.move), r.cont, w.Neg(), depth-1)
		score *= -1

		results = results.Update(Result{move: r.move, score: score.Abs(pos.ToMove), depth: depth - 1, cont: cont, err: err})

		if depth >= 3 && ctx.Err() != nil {
			break
		}

		var constrained, ok bool
		w, constrained, ok = w.Constrain(score)
		if constrained {
			// improved lower bound
		}
		if !ok && s.allowCutoff {
			break
		}
	}
	if err, ok := results[0].err.(checkmateError); ok {
		return w.alpha, results, err.Prev()
	}
	return w.alpha, results, results[0].err
}

// IsPseudoLegal returns whether a Move is pseudo-legal in a Position.
// A move is pseudo-legal if the square to be moved from contains the specified piece
// and the piece is capable of moving to the target square if doing so would not put the king in check.
func IsPseudoLegal(pos Position, move Move) bool {
	for _, m := range PieceMoves[move.Piece](pos) {
		if m == move {
			return true
		}
	}
	return false
}

// IsLegal returns whether a Position is legal.
// A position is illegal if the king of the side that just moved is in check.
func IsLegal(pos Position) bool {
	return !IsAttacked(pos, pos.KingSquare[pos.Opp()], pos.ToMove)
}

// IsCheck returns whether the king of the side to move is in check.
func IsCheck(pos Position) bool {
	return IsAttacked(pos, pos.KingSquare[pos.ToMove], pos.Opp())
}

// IsTerminal returns whether or not a Position is checkmate or stalemate.
// A position is checkmate or stalemate if the side to move has no legal moves.
func IsTerminal(pos Position) bool { return !anyLegal(pos, Candidates(pos)) }

// anyLegal returns whether any of the given Moves are legal in the given Position.
func anyLegal(pos Position, moves []Move) bool {
	for _, m := range moves {
		if IsLegal(Make(pos, m)) {
			return true
		}
	}
	return false
}

func checkDone(pos Position, recommended Results) (RelScore, Results, error) {
	if len(recommended) == 0 {
		moves := Candidates(pos) // pseudo-legal
		recommended = make(Results, 0, len(moves))
		for _, m := range moves {
			if !IsLegal(Make(pos, m)) {
				continue
			}
			recommended = append(recommended, Result{move: m})
		}
	}
	if len(recommended) == 0 {
		// no legal moves
		if IsCheck(pos) {
			return -evalInf, nil, errCheckmate
		}
		return 0, nil, errStalemate
	}
	if pos.HalfMove == 100 {
		// fifty-move rule
		return 0, nil, errFiftyMove
	}
	return 0, recommended, nil
}

// A Result holds a searched Move along with its evaluated Score,
// the search depth, and the continuation Results of the search.
type Result struct {
	move  Move
	score Score
	depth int
	cont  Results
	err   error
}

// String returns a string representation of r, including its principal variation.
func (r Result) String() string {
	s := fmt.Sprintf("%v (%v) %v", r.score, r.depth, r.PV())
	if r.err != nil {
		return s + " " + r.err.Error()
	}
	return s
}

// PV returns a string representation of r's principal variation.
func (r Result) PV() string {
	if r.depth == 0 || len(r.cont) == 0 {
		return LongAlgebraic(r.move)
	}
	return LongAlgebraic(r.move) + " " + r.cont[0].PV()
}

// A Results contains the results of a search.
// Functions returning a Results should sort it before returning.
type Results []Result

// SortFor sorts rs beginning with the best move for c.
// The sort is not guaranteed to be stable.
func (rs Results) SortFor(c Color) {
	// Sort first by mate condition, then by depth decreasing,
	// then by Score increasing/decreasing for White/Black,
	// and then by origin and destination Square increasing
	sort.Slice(rs, func(i, j int) bool {
		if less, ok := rs.mateSort(i, j); ok {
			return less
		}
		if rs[i].depth != rs[j].depth {
			return rs[i].depth > rs[j].depth
		}
		if rs[i].score != rs[j].score {
			return (rs[i].score > rs[j].score) == (c == White)
		}
		if rs[i].move.From != rs[j].move.From {
			return rs[i].move.From < rs[j].move.From
		}
		return rs[i].move.To < rs[j].move.To
	})
}

// mateSort reports how two Result elements should be sorted if at least one of them leads to checkmate.
// ok reports whether either Result contains an error of type checkmateError.
// If ok is true, less reports whether the Result with index i should sort before the Result with index j.
// If ok is false, the value of less is meaningless.
func (rs Results) mateSort(i, j int) (less bool, ok bool) {
	ich, iok := rs[i].err.(checkmateError)
	jch, jok := rs[j].err.(checkmateError)
	switch {
	case iok && jok:
		switch iwin, jwin := ich&1 == 0, jch&1 == 0; {
		case iwin && jwin:
			return ich < jch, true
		case !iwin && !jwin:
			return ich > jch, true
		default:
			return iwin, true
		}
	case iok:
		return ich&1 == 0, true
	case jok:
		return jch&1 != 0, true
	}
	return false, false
}

// Update looks for a Result in rs with the same Move as r and replaces it with r,
// or appends r if there is none, and returns the possibly updated slice.
func (rs Results) Update(r Result) Results {
	for i := range rs {
		if rs[i].move == r.move {
			rs[i] = r
			return rs
		}
	}
	return append(rs, r)
}

// String returns a string representation of all Result values in r.
func (rs Results) String() string {
	var s string
	for _, r := range rs {
		s += fmt.Sprintf("%v\n", r.String())
	}
	return s
}

// A Window represents the bounds of a position's evaluation.
type Window struct{ alpha, beta RelScore }

// Constrain updates the lower bound of w, if applicable, and returns the updated window,
// whether the lower bound was changed, and whether the returned Window remains valid.
func (w Window) Constrain(n RelScore) (c Window, constrained bool, ok bool) {
	if n <= w.alpha {
		return w, false, true
	}
	return Window{n, w.beta}, true, n <= w.beta
}

// Neg returns the additive inverse of w.
func (w Window) Neg() Window { return Window{-w.beta, -w.alpha} }
<|endoftext|>