<commit_before>/*
Copyright 2014 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 scdeny

import (
	"fmt"
	"io"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apiserver/pkg/admission"
	api "k8s.io/kubernetes/pkg/apis/core"
)

// PluginName indicates name of admission plugin.
const PluginName = "SecurityContextDeny"

// Register registers a plugin
func Register(plugins *admission.Plugins) {
	plugins.Register(PluginName, func(config io.Reader) (admission.Interface, error) {
		return NewSecurityContextDeny(), nil
	})
}

// Plugin implements admission.Interface.
type Plugin struct {
	*admission.Handler
}

var _ admission.ValidationInterface = &Plugin{}

// NewSecurityContextDeny creates a new instance of the SecurityContextDeny admission controller
func NewSecurityContextDeny() *Plugin {
	return &Plugin{
		Handler: admission.NewHandler(admission.Create, admission.Update),
	}
}

// Validate will deny any pod that defines SELinuxOptions or RunAsUser.
func (p *Plugin) Validate(a admission.Attributes) (err error) {
	if a.GetSubresource() != "" || a.GetResource().GroupResource() != api.Resource("pods") {
		return nil
	}

	pod, ok := a.GetObject().(*api.Pod)
	if !ok {
		return apierrors.NewBadRequest("Resource was marked with kind Pod but was unable to be converted")
	}

	if pod.Spec.SecurityContext != nil && pod.Spec.SecurityContext.SupplementalGroups != nil {
		return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.SupplementalGroups is forbidden"))
	}
	if pod.Spec.SecurityContext != nil {
		if pod.Spec.SecurityContext.SELinuxOptions != nil {
			return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("pod.Spec.SecurityContext.SELinuxOptions is forbidden"))
		}
		if pod.Spec.SecurityContext.RunAsUser != nil {
			return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("pod.Spec.SecurityContext.RunAsUser is forbidden"))
		}
	}

	if pod.Spec.SecurityContext != nil && pod.Spec.SecurityContext.FSGroup != nil {
		return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.FSGroup is forbidden"))
	}

	for _, v := range pod.Spec.InitContainers {
		if v.SecurityContext != nil {
			if v.SecurityContext.SELinuxOptions != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.SELinuxOptions is forbidden"))
			}
			if v.SecurityContext.RunAsUser != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.RunAsUser is forbidden"))
			}
		}
	}

	for _, v := range pod.Spec.Containers {
		if v.SecurityContext != nil {
			if v.SecurityContext.SELinuxOptions != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.SELinuxOptions is forbidden"))
			}
			if v.SecurityContext.RunAsUser != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.RunAsUser is forbidden"))
			}
		}
	}
	return nil
}
<commit_msg>simplify the if logic<commit_after>/*
Copyright 2014 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 scdeny

import (
	"fmt"
	"io"

	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apiserver/pkg/admission"
	api "k8s.io/kubernetes/pkg/apis/core"
)

// PluginName indicates name of admission plugin.
const PluginName = "SecurityContextDeny"

// Register registers a plugin
func Register(plugins *admission.Plugins) {
	plugins.Register(PluginName, func(config io.Reader) (admission.Interface, error) {
		return NewSecurityContextDeny(), nil
	})
}

// Plugin implements admission.Interface.
type Plugin struct {
	*admission.Handler
}

var _ admission.ValidationInterface = &Plugin{}

// NewSecurityContextDeny creates a new instance of the SecurityContextDeny admission controller
func NewSecurityContextDeny() *Plugin {
	return &Plugin{
		Handler: admission.NewHandler(admission.Create, admission.Update),
	}
}

// Validate will deny any pod that defines SupplementalGroups, SELinuxOptions, RunAsUser or FSGroup
func (p *Plugin) Validate(a admission.Attributes) (err error) {
	if a.GetSubresource() != "" || a.GetResource().GroupResource() != api.Resource("pods") {
		return nil
	}

	pod, ok := a.GetObject().(*api.Pod)
	if !ok {
		return apierrors.NewBadRequest("Resource was marked with kind Pod but was unable to be converted")
	}

	if pod.Spec.SecurityContext != nil {
		if pod.Spec.SecurityContext.SupplementalGroups != nil {
			return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("pod.Spec.SecurityContext.SupplementalGroups is forbidden"))
		}
		if pod.Spec.SecurityContext.SELinuxOptions != nil {
			return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("pod.Spec.SecurityContext.SELinuxOptions is forbidden"))
		}
		if pod.Spec.SecurityContext.RunAsUser != nil {
			return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("pod.Spec.SecurityContext.RunAsUser is forbidden"))
		}
		if pod.Spec.SecurityContext.FSGroup != nil {
			return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("pod.Spec.SecurityContext.FSGroup is forbidden"))
		}
	}

	for _, v := range pod.Spec.InitContainers {
		if v.SecurityContext != nil {
			if v.SecurityContext.SELinuxOptions != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.SELinuxOptions is forbidden"))
			}
			if v.SecurityContext.RunAsUser != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.RunAsUser is forbidden"))
			}
		}
	}

	for _, v := range pod.Spec.Containers {
		if v.SecurityContext != nil {
			if v.SecurityContext.SELinuxOptions != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.SELinuxOptions is forbidden"))
			}
			if v.SecurityContext.RunAsUser != nil {
				return apierrors.NewForbidden(a.GetResource().GroupResource(), pod.Name, fmt.Errorf("SecurityContext.RunAsUser is forbidden"))
			}
		}
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2018 The Kubeflow 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 controller provides a Kubernetes controller for a MXJob resource.
package mxnet

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"

	mxv1beta1 "github.com/kubeflow/mxnet-operator/pkg/apis/mxnet/v1beta1"
	"github.com/kubeflow/tf-operator/pkg/common/jobcontroller"
	mxlogger "github.com/kubeflow/tf-operator/pkg/logger"
	train_util "github.com/kubeflow/tf-operator/pkg/util/train"
)

const (
	// mxConfig is the environment variable name of MXNet cluster spec.
	mxConfig = "MX_CONFIG"

	// podTemplateRestartPolicyReason is the warning reason when the restart
	// policy is set in pod template.
	podTemplateRestartPolicyReason = "SettedPodTemplateRestartPolicy"
	// exitedWithCodeReason is the normal reason when the pod is exited because of the exit code.
	exitedWithCodeReason = "ExitedWithCode"
)

// reconcilePods checks and updates pods for each given MXReplicaSpec.
// It will requeue the mxjob in case of an error while creating/deleting pods.
func (tc *MXController) reconcilePods(
	mxjob *mxv1beta1.MXJob,
	pods []*v1.Pod,
	rtype mxv1beta1.MXReplicaType,
	spec *mxv1beta1.MXReplicaSpec, rstatus map[string]v1.PodPhase) error {

	// Convert MXReplicaType to lower string.
	rt := strings.ToLower(string(rtype))
	logger := mxlogger.LoggerForReplica(mxjob, rt)
	// Get all pods for the type rt.
	pods, err := tc.FilterPodsForReplicaType(pods, rt)
	if err != nil {
		return err
	}
	replicas := int(*spec.Replicas)
	restart := false
	schedulerCompleted := false

	initializeMXReplicaStatuses(mxjob, rtype)

	podSlices := tc.GetPodSlices(pods, replicas, logger)
	for index, podSlice := range podSlices {
		if len(podSlice) > 1 {
			logger.Warningf("We have too many pods for %s %d", rt, index)
			// TODO(gaocegege): Kill some pods.
		} else if len(podSlice) == 0 {
			logger.Infof("Need to create new pod: %s-%d", rt, index)
			err = tc.createNewPod(mxjob, rt, strconv.Itoa(index), spec)
			if err != nil {
				return err
			}
		} else {
			// Check the status of the current pod.
			pod := podSlice[0]
			// Get the exit code of the mxnet container.
			var exitCode int32 = 0xbeef // magic number
			for _, status := range pod.Status.ContainerStatuses {
				state := status.State
				if status.Name == mxv1beta1.DefaultContainerName && state.Terminated != nil {
					exitCode = state.Terminated.ExitCode
					logger.Infof("Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
					tc.Recorder.Eventf(mxjob, v1.EventTypeNormal, exitedWithCodeReason, "Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
				}
			}
			// Check if the pod is retryable.
			if spec.RestartPolicy == mxv1beta1.RestartPolicyExitCode {
				if pod.Status.Phase == v1.PodFailed && train_util.IsRetryableExitCode(exitCode) {
					logger.Infof("Need to restart the pod: %v.%v", pod.Namespace, pod.Name)
					if err := tc.PodControl.DeletePod(pod.Namespace, pod.Name, mxjob); err != nil {
						return err
					}
					restart = true
				}
			}

			// Check whether scheduler is exited without error.
			if rtype == mxv1beta1.MXReplicaTypeScheduler && exitCode == 0 {
				schedulerCompleted = true
			}
			updateMXJobReplicaStatuses(mxjob, rtype, pod)
		}
	}

	return updateStatusSingle(mxjob, rtype, replicas, restart, schedulerCompleted)
}

// createNewPod creates a new pod for the given index and type.
func (tc *MXController) createNewPod(mxjob *mxv1beta1.MXJob, rt, index string, spec *mxv1beta1.MXReplicaSpec) error {
	mxjobKey, err := KeyFunc(mxjob)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("Couldn't get key for mxjob object %#v: %v", mxjob, err))
		return err
	}
	expectationPodsKey := jobcontroller.GenExpectationPodsKey(mxjobKey, rt)
	err = tc.Expectations.ExpectCreations(expectationPodsKey, 1)
	if err != nil {
		return err
	}
	logger := mxlogger.LoggerForReplica(mxjob, rt)
	// Create OwnerReference.
	controllerRef := tc.GenOwnerReference(mxjob)

	// Set type and index for the worker.
	labels := tc.GenLabels(mxjob.Name)
	labels[mxReplicaTypeLabel] = rt
	labels[mxReplicaIndexLabel] = index

	podTemplate := spec.Template.DeepCopy()

	// Set name for the template.
	podTemplate.Name = jobcontroller.GenGeneralName(mxjob.Name, rt, index)

	if podTemplate.Labels == nil {
		podTemplate.Labels = make(map[string]string)
	}

	for key, value := range labels {
		podTemplate.Labels[key] = value
	}

	if err := setClusterSpec(podTemplate, mxjob, rt, index); err != nil {
		return err
	}

	// Submit a warning event if the user specifies restart policy for
	// the pod template. We recommend to set it from the replica level.
	if podTemplate.Spec.RestartPolicy != v1.RestartPolicy("") {
		errMsg := "Restart policy in pod template will be overwritten by restart policy in replica spec"
		logger.Warning(errMsg)
		tc.Recorder.Event(mxjob, v1.EventTypeWarning, podTemplateRestartPolicyReason, errMsg)
	}
	setRestartPolicy(podTemplate, spec)

	err = tc.PodControl.CreatePodsWithControllerRef(mxjob.Namespace, podTemplate, mxjob, controllerRef)
	if err != nil && errors.IsTimeout(err) {
		// Pod is created but its initialization has timed out.
		// If the initialization is successful eventually, the
		// controller will observe the creation via the informer.
		// If the initialization fails, or if the pod keeps
		// uninitialized for a long time, the informer will not
		// receive any update, and the controller will create a new
		// pod when the expectation expires.
		return nil
	} else if err != nil {
		return err
	}
	return nil
}

func setClusterSpec(podTemplateSpec *v1.PodTemplateSpec, mxjob *mxv1beta1.MXJob, rt, index string) error {

	// Generate MX_CONFIG JSON.
	mxConfigData, err := genMXConfig(mxjob, rt, index)
	if err != nil {
		return err
	}

	// Generate MX_CONFIG JSON Str.
	mxConfigJson, err := json.Marshal(mxConfigData)
	if err != nil {
		return err
	}

	// Add MX_CONFIG environment variable.
	for i := range podTemplateSpec.Spec.Containers {

		c := &podTemplateSpec.Spec.Containers[i]

		// Set environment variable MX_CONFIG
		c.Env = append(c.Env, v1.EnvVar{
			Name:  mxConfig,
			Value: string(mxConfigJson),
		})

		// Set Mxnet Distributed Training environment variable
		// We get these envs from MX_COFING to make them stay identical
		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_PS_ROOT_PORT",
			Value: strconv.Itoa(getConfigAddr(&mxConfigData, mxv1beta1.MXReplicaTypeScheduler, 0).Port),
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_PS_ROOT_URI",
			Value: getConfigAddr(&mxConfigData, mxv1beta1.MXReplicaTypeScheduler, 0).Url,
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_NUM_SERVER",
			Value: strconv.Itoa(getConfigReplica(&mxConfigData, mxv1beta1.MXReplicaTypeServer)),
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_NUM_WORKER",
			Value: strconv.Itoa(getConfigReplica(&mxConfigData, mxv1beta1.MXReplicaTypeWorker)),
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_ROLE",
			Value: mxConfigData.Task.Type,
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_USE_KUBERNETES",
			Value: strconv.Itoa(1),
		})
	}
	return nil
}

func setRestartPolicy(podTemplateSpec *v1.PodTemplateSpec, spec *mxv1beta1.MXReplicaSpec) {
	if spec.RestartPolicy == mxv1beta1.RestartPolicyExitCode {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicyNever
	} else {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicy(spec.RestartPolicy)
	}
}

func getConfigAddr(mxConfigData *MXConfig, rtype mxv1beta1.MXReplicaType, index int) Url_Port {
	rt := strings.ToLower(string(rtype))
	var url_port Url_Port
	if len(mxConfigData.Cluster[rt]) <= index {
		// index out of range, maybe this url doen't exist
		url_port = Url_Port{
			Url:  "",
			Port: 0,
		}
	} else {
		url_port = mxConfigData.Cluster[rt][index]
	}
	return url_port
}

func getConfigReplica(mxConfigData *MXConfig, rtype mxv1beta1.MXReplicaType) int {
	rt := strings.ToLower(string(rtype))
	return len(mxConfigData.Cluster[rt])
}
<commit_msg>Use kube-batch as scheduler by default when gang-scheduling is enabled (#35)<commit_after>// Copyright 2018 The Kubeflow 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 controller provides a Kubernetes controller for a MXJob resource.
package mxnet

import (
	"encoding/json"
	"fmt"
	"strconv"
	"strings"

	"k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/errors"
	utilruntime "k8s.io/apimachinery/pkg/util/runtime"

	mxv1beta1 "github.com/kubeflow/mxnet-operator/pkg/apis/mxnet/v1beta1"
	"github.com/kubeflow/tf-operator/pkg/common/jobcontroller"
	mxlogger "github.com/kubeflow/tf-operator/pkg/logger"
	train_util "github.com/kubeflow/tf-operator/pkg/util/train"
)

const (
	// gang scheduler name.
	gangSchedulerName = "kube-batch"

	// mxConfig is the environment variable name of MXNet cluster spec.
	mxConfig = "MX_CONFIG"

	// podTemplateRestartPolicyReason is the warning reason when the restart
	// policy is set in pod template.
	podTemplateRestartPolicyReason = "SettedPodTemplateRestartPolicy"
	// exitedWithCodeReason is the normal reason when the pod is exited because of the exit code.
	exitedWithCodeReason = "ExitedWithCode"
	// podTemplateSchedulerNameReason is the warning reason when other scheduler name is set
	// in pod templates with gang-scheduling enabled
	podTemplateSchedulerNameReason = "SettedPodTemplateSchedulerName"
)

// reconcilePods checks and updates pods for each given MXReplicaSpec.
// It will requeue the mxjob in case of an error while creating/deleting pods.
func (tc *MXController) reconcilePods(
	mxjob *mxv1beta1.MXJob,
	pods []*v1.Pod,
	rtype mxv1beta1.MXReplicaType,
	spec *mxv1beta1.MXReplicaSpec, rstatus map[string]v1.PodPhase) error {

	// Convert MXReplicaType to lower string.
	rt := strings.ToLower(string(rtype))
	logger := mxlogger.LoggerForReplica(mxjob, rt)
	// Get all pods for the type rt.
	pods, err := tc.FilterPodsForReplicaType(pods, rt)
	if err != nil {
		return err
	}
	replicas := int(*spec.Replicas)
	restart := false
	schedulerCompleted := false

	initializeMXReplicaStatuses(mxjob, rtype)

	podSlices := tc.GetPodSlices(pods, replicas, logger)
	for index, podSlice := range podSlices {
		if len(podSlice) > 1 {
			logger.Warningf("We have too many pods for %s %d", rt, index)
			// TODO(gaocegege): Kill some pods.
		} else if len(podSlice) == 0 {
			logger.Infof("Need to create new pod: %s-%d", rt, index)
			err = tc.createNewPod(mxjob, rt, strconv.Itoa(index), spec)
			if err != nil {
				return err
			}
		} else {
			// Check the status of the current pod.
			pod := podSlice[0]
			// Get the exit code of the mxnet container.
			var exitCode int32 = 0xbeef // magic number
			for _, status := range pod.Status.ContainerStatuses {
				state := status.State
				if status.Name == mxv1beta1.DefaultContainerName && state.Terminated != nil {
					exitCode = state.Terminated.ExitCode
					logger.Infof("Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
					tc.Recorder.Eventf(mxjob, v1.EventTypeNormal, exitedWithCodeReason, "Pod: %v.%v exited with code %v", pod.Namespace, pod.Name, exitCode)
				}
			}
			// Check if the pod is retryable.
			if spec.RestartPolicy == mxv1beta1.RestartPolicyExitCode {
				if pod.Status.Phase == v1.PodFailed && train_util.IsRetryableExitCode(exitCode) {
					logger.Infof("Need to restart the pod: %v.%v", pod.Namespace, pod.Name)
					if err := tc.PodControl.DeletePod(pod.Namespace, pod.Name, mxjob); err != nil {
						return err
					}
					restart = true
				}
			}

			// Check whether scheduler is exited without error.
			if rtype == mxv1beta1.MXReplicaTypeScheduler && exitCode == 0 {
				schedulerCompleted = true
			}
			updateMXJobReplicaStatuses(mxjob, rtype, pod)
		}
	}

	return updateStatusSingle(mxjob, rtype, replicas, restart, schedulerCompleted)
}

// createNewPod creates a new pod for the given index and type.
func (tc *MXController) createNewPod(mxjob *mxv1beta1.MXJob, rt, index string, spec *mxv1beta1.MXReplicaSpec) error {
	mxjobKey, err := KeyFunc(mxjob)
	if err != nil {
		utilruntime.HandleError(fmt.Errorf("Couldn't get key for mxjob object %#v: %v", mxjob, err))
		return err
	}
	expectationPodsKey := jobcontroller.GenExpectationPodsKey(mxjobKey, rt)
	err = tc.Expectations.ExpectCreations(expectationPodsKey, 1)
	if err != nil {
		return err
	}
	logger := mxlogger.LoggerForReplica(mxjob, rt)
	// Create OwnerReference.
	controllerRef := tc.GenOwnerReference(mxjob)

	// Set type and index for the worker.
	labels := tc.GenLabels(mxjob.Name)
	labels[mxReplicaTypeLabel] = rt
	labels[mxReplicaIndexLabel] = index

	podTemplate := spec.Template.DeepCopy()

	// Set name for the template.
	podTemplate.Name = jobcontroller.GenGeneralName(mxjob.Name, rt, index)

	if podTemplate.Labels == nil {
		podTemplate.Labels = make(map[string]string)
	}

	for key, value := range labels {
		podTemplate.Labels[key] = value
	}

	if err := setClusterSpec(podTemplate, mxjob, rt, index); err != nil {
		return err
	}

	// Submit a warning event if the user specifies restart policy for
	// the pod template. We recommend to set it from the replica level.
	if podTemplate.Spec.RestartPolicy != v1.RestartPolicy("") {
		errMsg := "Restart policy in pod template will be overwritten by restart policy in replica spec"
		logger.Warning(errMsg)
		tc.Recorder.Event(mxjob, v1.EventTypeWarning, podTemplateRestartPolicyReason, errMsg)
	}
	setRestartPolicy(podTemplate, spec)

	// if gang-scheduling is enabled:
	// 1. if user has specified other scheduler, we report a warning without overriding any fields.
	// 2. if no SchedulerName is set for pods, then we set the SchedulerName to "kube-batch".
	if tc.Config.EnableGangScheduling {
		if isNonGangSchedulerSet(mxjob) {
			errMsg := "Another scheduler is specified when gang-scheduling is enabled and it will not be overwritten"
			logger.Warning(errMsg)
			tc.Recorder.Event(mxjob, v1.EventTypeWarning, podTemplateSchedulerNameReason, errMsg)
		} else {
			podTemplate.Spec.SchedulerName = gangSchedulerName
		}
	}

	err = tc.PodControl.CreatePodsWithControllerRef(mxjob.Namespace, podTemplate, mxjob, controllerRef)
	if err != nil && errors.IsTimeout(err) {
		// Pod is created but its initialization has timed out.
		// If the initialization is successful eventually, the
		// controller will observe the creation via the informer.
		// If the initialization fails, or if the pod keeps
		// uninitialized for a long time, the informer will not
		// receive any update, and the controller will create a new
		// pod when the expectation expires.
		return nil
	} else if err != nil {
		return err
	}
	return nil
}

func setClusterSpec(podTemplateSpec *v1.PodTemplateSpec, mxjob *mxv1beta1.MXJob, rt, index string) error {

	// Generate MX_CONFIG JSON.
	mxConfigData, err := genMXConfig(mxjob, rt, index)
	if err != nil {
		return err
	}

	// Generate MX_CONFIG JSON Str.
	mxConfigJson, err := json.Marshal(mxConfigData)
	if err != nil {
		return err
	}

	// Add MX_CONFIG environment variable.
	for i := range podTemplateSpec.Spec.Containers {

		c := &podTemplateSpec.Spec.Containers[i]

		// Set environment variable MX_CONFIG
		c.Env = append(c.Env, v1.EnvVar{
			Name:  mxConfig,
			Value: string(mxConfigJson),
		})

		// Set Mxnet Distributed Training environment variable
		// We get these envs from MX_COFING to make them stay identical
		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_PS_ROOT_PORT",
			Value: strconv.Itoa(getConfigAddr(&mxConfigData, mxv1beta1.MXReplicaTypeScheduler, 0).Port),
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_PS_ROOT_URI",
			Value: getConfigAddr(&mxConfigData, mxv1beta1.MXReplicaTypeScheduler, 0).Url,
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_NUM_SERVER",
			Value: strconv.Itoa(getConfigReplica(&mxConfigData, mxv1beta1.MXReplicaTypeServer)),
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_NUM_WORKER",
			Value: strconv.Itoa(getConfigReplica(&mxConfigData, mxv1beta1.MXReplicaTypeWorker)),
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_ROLE",
			Value: mxConfigData.Task.Type,
		})

		c.Env = append(c.Env, v1.EnvVar{
			Name:  "DMLC_USE_KUBERNETES",
			Value: strconv.Itoa(1),
		})
	}
	return nil
}

func setRestartPolicy(podTemplateSpec *v1.PodTemplateSpec, spec *mxv1beta1.MXReplicaSpec) {
	if spec.RestartPolicy == mxv1beta1.RestartPolicyExitCode {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicyNever
	} else {
		podTemplateSpec.Spec.RestartPolicy = v1.RestartPolicy(spec.RestartPolicy)
	}
}

func getConfigAddr(mxConfigData *MXConfig, rtype mxv1beta1.MXReplicaType, index int) Url_Port {
	rt := strings.ToLower(string(rtype))
	var url_port Url_Port
	if len(mxConfigData.Cluster[rt]) <= index {
		// index out of range, maybe this url doen't exist
		url_port = Url_Port{
			Url:  "",
			Port: 0,
		}
	} else {
		url_port = mxConfigData.Cluster[rt][index]
	}
	return url_port
}

func getConfigReplica(mxConfigData *MXConfig, rtype mxv1beta1.MXReplicaType) int {
	rt := strings.ToLower(string(rtype))
	return len(mxConfigData.Cluster[rt])
}

func isNonGangSchedulerSet(job *mxv1beta1.MXJob) bool {
	for _, spec := range job.Spec.MXReplicaSpecs {
		if spec.Template.Spec.SchedulerName != "" && spec.Template.Spec.SchedulerName != gangSchedulerName {
			return true
		}
	}
	return false
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"math/rand"
	"reflect"
	"sync"
	"time"

	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/log"
	"golang.org/x/net/context"
)

const (
	initialSessionFailureBackoff = time.Second
	maxSessionFailureBackoff     = 8 * time.Second
)

// Agent implements the primary node functionality for a member of a swarm
// cluster. The primary functionality id to run and report on the status of
// tasks assigned to the node.
type Agent struct {
	config *Config

	// The latest node object state from manager
	// for this node known to the agent.
	node *api.Node

	keys []*api.EncryptionKey

	sessionq chan sessionOperation
	worker   Worker

	started chan struct{}
	ready   chan struct{}
	stopped chan struct{} // requests shutdown
	closed  chan struct{} // only closed in run
	err     error         // read only after closed is closed
	mu      sync.Mutex
}

// New returns a new agent, ready for task dispatch.
func New(config *Config) (*Agent, error) {
	if err := config.validate(); err != nil {
		return nil, err
	}

	a := &Agent{
		config:   config,
		worker:   newWorker(config.DB, config.Executor),
		sessionq: make(chan sessionOperation),
		started:  make(chan struct{}),
		stopped:  make(chan struct{}),
		closed:   make(chan struct{}),
	}

	return a, nil
}

// Start begins execution of the agent in the provided context, if not already
// started.
func (a *Agent) Start(ctx context.Context) error {
	select {
	case <-a.started:
		select {
		case <-a.closed:
			return a.err
		case <-a.stopped:
			return errAgentStopped
		case <-ctx.Done():
			return ctx.Err()
		default:
			return errAgentStarted
		}
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	close(a.started)
	go a.run(ctx)

	return nil
}

// Stop shuts down the agent, blocking until full shutdown. If the agent is not
// started, Stop will block until Started.
func (a *Agent) Stop(ctx context.Context) error {
	select {
	case <-a.started:
		select {
		case <-a.closed:
			return a.err
		case <-a.stopped:
			select {
			case <-a.closed:
				return a.err
			case <-ctx.Done():
				return ctx.Err()
			}
		case <-ctx.Done():
			return ctx.Err()
		default:
			close(a.stopped)
			// recurse and wait for closure
			return a.Stop(ctx)
		}
	case <-ctx.Done():
		return ctx.Err()
	default:
		return errAgentNotStarted
	}
}

// Err returns the error that caused the agent to shutdown or nil. Err blocks
// until the agent is fully shutdown.
func (a *Agent) Err(ctx context.Context) error {
	select {
	case <-a.closed:
		return a.err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Ready returns a channel that will be closed when agent first becomes ready.
func (a *Agent) Ready() <-chan struct{} {
	return a.ready
}

func (a *Agent) run(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer close(a.closed) // full shutdown.

	log.G(ctx).Debugf("(*Agent).run")
	defer log.G(ctx).Debugf("(*Agent).run exited")

	var (
		backoff    time.Duration
		session    = newSession(ctx, a, backoff) // start the initial session
		registered = session.registered
		ready      = a.ready // first session ready
		sessionq   = a.sessionq
	)

	if err := a.worker.Init(ctx); err != nil {
		log.G(ctx).WithError(err).Error("worker initialization failed")
		a.err = err
		return // fatal?
	}

	// setup a reliable reporter to call back to us.
	reporter := newStatusReporter(ctx, a)
	defer reporter.Close()

	a.worker.Listen(ctx, reporter)

	for {
		select {
		case operation := <-sessionq:
			operation.response <- operation.fn(session)
		case msg := <-session.tasks:
			go func() {
				if err := a.worker.Assign(ctx, msg.Tasks); err != nil {
					log.G(ctx).WithError(err).Error("task assignment failed")
				}
			}()
		case msg := <-session.messages:
			if err := a.handleSessionMessage(ctx, msg); err != nil {
				log.G(ctx).WithError(err).Error("session message handler failed")
			}
		case <-registered:
			if ready != nil {
				close(ready)
			}
			ready = nil
			log.G(ctx).Debugln("agent: registered")
			registered = nil // we only care about this once per session
			backoff = 0      // reset backoff
		case err := <-session.errs:
			// TODO(stevvooe): This may actually block if a session is closed
			// but no error was sent. Session.close must only be called here
			// for this to work.
			if err != nil {
				log.G(ctx).WithError(err).Error("agent: session failed")
				backoff = initialSessionFailureBackoff + 2*backoff
				if backoff > maxSessionFailureBackoff {
					backoff = maxSessionFailureBackoff
				}
			}

			if err := session.close(); err != nil {
				log.G(ctx).WithError(err).Error("agent: closing session failed")
			}
			sessionq = nil
		case <-session.closed:
			log.G(ctx).Debugf("agent: rebuild session")

			// select a session registration delay from backoff range.
			delay := time.Duration(rand.Int63n(int64(backoff)))
			session = newSession(ctx, a, delay)
			registered = session.registered
			sessionq = a.sessionq
		case <-a.stopped:
			// TODO(stevvooe): Wait on shutdown and cleanup. May need to pump
			// this loop a few times.
			return
		case <-ctx.Done():
			if a.err == nil {
				a.err = ctx.Err()
			}

			return
		}
	}
}

func (a *Agent) handleSessionMessage(ctx context.Context, message *api.SessionMessage) error {
	seen := map[api.Peer]struct{}{}
	for _, manager := range message.Managers {
		if manager.Peer.Addr == "" {
			log.G(ctx).WithField("manager.addr", manager.Peer.Addr).
				Warnf("skipping bad manager address")
			continue
		}

		a.config.Managers.Observe(*manager.Peer, int(manager.Weight))
		seen[*manager.Peer] = struct{}{}
	}

	if message.Node != nil {
		if a.node == nil || !nodesEqual(a.node, message.Node) {
			if a.config.NotifyRoleChange != nil {
				a.config.NotifyRoleChange <- message.Node.Spec.Role
			}
			a.node = message.Node.Copy()
			if err := a.config.Executor.Configure(ctx, a.node); err != nil {
				log.G(ctx).WithError(err).Error("node configure failed")
			}
		}
	}

	// prune managers not in list.
	for peer := range a.config.Managers.Weights() {
		if _, ok := seen[peer]; !ok {
			a.config.Managers.Remove(peer)
		}
	}

	if message.NetworkBootstrapKeys == nil {
		return nil
	}

	for _, key := range message.NetworkBootstrapKeys {
		same := false
		for _, agentKey := range a.keys {
			if agentKey.LamportTime == key.LamportTime {
				same = true
			}
		}
		if !same {
			a.keys = message.NetworkBootstrapKeys
			if err := a.config.Executor.SetNetworkBootstrapKeys(a.keys); err != nil {
				panic(fmt.Errorf("configuring network key failed"))
			}
		}
	}

	return nil
}

type sessionOperation struct {
	fn       func(session *session) error
	response chan error
}

// withSession runs fn with the current session.
func (a *Agent) withSession(ctx context.Context, fn func(session *session) error) error {
	response := make(chan error, 1)
	select {
	case a.sessionq <- sessionOperation{
		fn:       fn,
		response: response,
	}:
		select {
		case err := <-response:
			return err
		case <-a.closed:
			return ErrClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	case <-a.closed:
		return ErrClosed
	case <-ctx.Done():
		return ctx.Err()
	}
}

// UpdateTaskStatus attempts to send a task status update over the current session,
// blocking until the operation is completed.
//
// If an error is returned, the operation should be retried.
func (a *Agent) UpdateTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {
	log.G(ctx).Debugf("(*Agent).UpdateTaskStatus")
	defer log.G(ctx).Debugf("(*Agent).UpdateTaskStatus leave")
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	errs := make(chan error, 1)
	if err := a.withSession(ctx, func(session *session) error {
		log.G(ctx).Debugf("(*Agent).withSession")
		defer log.G(ctx).Debugf("(*Agent).withSession leave")
		go func() {
			err := session.sendTaskStatus(ctx, taskID, status)
			if err != nil {
				if err == errTaskUnknown {
					err = nil // dispatcher no longer cares about this task.
				} else {
					log.G(ctx).WithError(err).Error("sending task status update failed")
				}
			} else {
				log.G(ctx).Debug("task status reported")
			}

			errs <- err
		}()

		return nil
	}); err != nil {
		return err
	}

	select {
	case err := <-errs:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// nodesEqual returns true if the node states are functionaly equal, ignoring status,
// version and other superfluous fields.
//
// This used to decide whether or not to propagate a node update to executor.
func nodesEqual(a, b *api.Node) bool {
	a, b = a.Copy(), b.Copy()

	a.Status, b.Status = api.NodeStatus{}, api.NodeStatus{}
	a.Meta, b.Meta = api.Meta{}, api.Meta{}

	return reflect.DeepEqual(a, b)
}
<commit_msg>Fix agent ready event<commit_after>package agent

import (
	"fmt"
	"math/rand"
	"reflect"
	"sync"
	"time"

	"github.com/docker/swarmkit/api"
	"github.com/docker/swarmkit/log"
	"golang.org/x/net/context"
)

const (
	initialSessionFailureBackoff = time.Second
	maxSessionFailureBackoff     = 8 * time.Second
)

// Agent implements the primary node functionality for a member of a swarm
// cluster. The primary functionality id to run and report on the status of
// tasks assigned to the node.
type Agent struct {
	config *Config

	// The latest node object state from manager
	// for this node known to the agent.
	node *api.Node

	keys []*api.EncryptionKey

	sessionq chan sessionOperation
	worker   Worker

	started chan struct{}
	ready   chan struct{}
	stopped chan struct{} // requests shutdown
	closed  chan struct{} // only closed in run
	err     error         // read only after closed is closed
	mu      sync.Mutex
}

// New returns a new agent, ready for task dispatch.
func New(config *Config) (*Agent, error) {
	if err := config.validate(); err != nil {
		return nil, err
	}

	a := &Agent{
		config:   config,
		worker:   newWorker(config.DB, config.Executor),
		sessionq: make(chan sessionOperation),
		started:  make(chan struct{}),
		stopped:  make(chan struct{}),
		closed:   make(chan struct{}),
		ready:    make(chan struct{}),
	}

	return a, nil
}

// Start begins execution of the agent in the provided context, if not already
// started.
func (a *Agent) Start(ctx context.Context) error {
	select {
	case <-a.started:
		select {
		case <-a.closed:
			return a.err
		case <-a.stopped:
			return errAgentStopped
		case <-ctx.Done():
			return ctx.Err()
		default:
			return errAgentStarted
		}
	case <-ctx.Done():
		return ctx.Err()
	default:
	}

	close(a.started)
	go a.run(ctx)

	return nil
}

// Stop shuts down the agent, blocking until full shutdown. If the agent is not
// started, Stop will block until Started.
func (a *Agent) Stop(ctx context.Context) error {
	select {
	case <-a.started:
		select {
		case <-a.closed:
			return a.err
		case <-a.stopped:
			select {
			case <-a.closed:
				return a.err
			case <-ctx.Done():
				return ctx.Err()
			}
		case <-ctx.Done():
			return ctx.Err()
		default:
			close(a.stopped)
			// recurse and wait for closure
			return a.Stop(ctx)
		}
	case <-ctx.Done():
		return ctx.Err()
	default:
		return errAgentNotStarted
	}
}

// Err returns the error that caused the agent to shutdown or nil. Err blocks
// until the agent is fully shutdown.
func (a *Agent) Err(ctx context.Context) error {
	select {
	case <-a.closed:
		return a.err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// Ready returns a channel that will be closed when agent first becomes ready.
func (a *Agent) Ready() <-chan struct{} {
	return a.ready
}

func (a *Agent) run(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	defer close(a.closed) // full shutdown.

	log.G(ctx).Debugf("(*Agent).run")
	defer log.G(ctx).Debugf("(*Agent).run exited")

	var (
		backoff    time.Duration
		session    = newSession(ctx, a, backoff) // start the initial session
		registered = session.registered
		ready      = a.ready // first session ready
		sessionq   = a.sessionq
	)

	if err := a.worker.Init(ctx); err != nil {
		log.G(ctx).WithError(err).Error("worker initialization failed")
		a.err = err
		return // fatal?
	}

	// setup a reliable reporter to call back to us.
	reporter := newStatusReporter(ctx, a)
	defer reporter.Close()

	a.worker.Listen(ctx, reporter)

	for {
		select {
		case operation := <-sessionq:
			operation.response <- operation.fn(session)
		case msg := <-session.tasks:
			go func() {
				if err := a.worker.Assign(ctx, msg.Tasks); err != nil {
					log.G(ctx).WithError(err).Error("task assignment failed")
				}
			}()
		case msg := <-session.messages:
			if err := a.handleSessionMessage(ctx, msg); err != nil {
				log.G(ctx).WithError(err).Error("session message handler failed")
			}
		case <-registered:
			if ready != nil {
				close(ready)
			}
			ready = nil
			log.G(ctx).Debugln("agent: registered")
			registered = nil // we only care about this once per session
			backoff = 0      // reset backoff
		case err := <-session.errs:
			// TODO(stevvooe): This may actually block if a session is closed
			// but no error was sent. Session.close must only be called here
			// for this to work.
			if err != nil {
				log.G(ctx).WithError(err).Error("agent: session failed")
				backoff = initialSessionFailureBackoff + 2*backoff
				if backoff > maxSessionFailureBackoff {
					backoff = maxSessionFailureBackoff
				}
			}

			if err := session.close(); err != nil {
				log.G(ctx).WithError(err).Error("agent: closing session failed")
			}
			sessionq = nil
		case <-session.closed:
			log.G(ctx).Debugf("agent: rebuild session")

			// select a session registration delay from backoff range.
			delay := time.Duration(rand.Int63n(int64(backoff)))
			session = newSession(ctx, a, delay)
			registered = session.registered
			sessionq = a.sessionq
		case <-a.stopped:
			// TODO(stevvooe): Wait on shutdown and cleanup. May need to pump
			// this loop a few times.
			return
		case <-ctx.Done():
			if a.err == nil {
				a.err = ctx.Err()
			}

			return
		}
	}
}

func (a *Agent) handleSessionMessage(ctx context.Context, message *api.SessionMessage) error {
	seen := map[api.Peer]struct{}{}
	for _, manager := range message.Managers {
		if manager.Peer.Addr == "" {
			log.G(ctx).WithField("manager.addr", manager.Peer.Addr).
				Warnf("skipping bad manager address")
			continue
		}

		a.config.Managers.Observe(*manager.Peer, int(manager.Weight))
		seen[*manager.Peer] = struct{}{}
	}

	if message.Node != nil {
		if a.node == nil || !nodesEqual(a.node, message.Node) {
			if a.config.NotifyRoleChange != nil {
				a.config.NotifyRoleChange <- message.Node.Spec.Role
			}
			a.node = message.Node.Copy()
			if err := a.config.Executor.Configure(ctx, a.node); err != nil {
				log.G(ctx).WithError(err).Error("node configure failed")
			}
		}
	}

	// prune managers not in list.
	for peer := range a.config.Managers.Weights() {
		if _, ok := seen[peer]; !ok {
			a.config.Managers.Remove(peer)
		}
	}

	if message.NetworkBootstrapKeys == nil {
		return nil
	}

	for _, key := range message.NetworkBootstrapKeys {
		same := false
		for _, agentKey := range a.keys {
			if agentKey.LamportTime == key.LamportTime {
				same = true
			}
		}
		if !same {
			a.keys = message.NetworkBootstrapKeys
			if err := a.config.Executor.SetNetworkBootstrapKeys(a.keys); err != nil {
				panic(fmt.Errorf("configuring network key failed"))
			}
		}
	}

	return nil
}

type sessionOperation struct {
	fn       func(session *session) error
	response chan error
}

// withSession runs fn with the current session.
func (a *Agent) withSession(ctx context.Context, fn func(session *session) error) error {
	response := make(chan error, 1)
	select {
	case a.sessionq <- sessionOperation{
		fn:       fn,
		response: response,
	}:
		select {
		case err := <-response:
			return err
		case <-a.closed:
			return ErrClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	case <-a.closed:
		return ErrClosed
	case <-ctx.Done():
		return ctx.Err()
	}
}

// UpdateTaskStatus attempts to send a task status update over the current session,
// blocking until the operation is completed.
//
// If an error is returned, the operation should be retried.
func (a *Agent) UpdateTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {
	log.G(ctx).Debugf("(*Agent).UpdateTaskStatus")
	defer log.G(ctx).Debugf("(*Agent).UpdateTaskStatus leave")
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	errs := make(chan error, 1)
	if err := a.withSession(ctx, func(session *session) error {
		log.G(ctx).Debugf("(*Agent).withSession")
		defer log.G(ctx).Debugf("(*Agent).withSession leave")
		go func() {
			err := session.sendTaskStatus(ctx, taskID, status)
			if err != nil {
				if err == errTaskUnknown {
					err = nil // dispatcher no longer cares about this task.
				} else {
					log.G(ctx).WithError(err).Error("sending task status update failed")
				}
			} else {
				log.G(ctx).Debug("task status reported")
			}

			errs <- err
		}()

		return nil
	}); err != nil {
		return err
	}

	select {
	case err := <-errs:
		return err
	case <-ctx.Done():
		return ctx.Err()
	}
}

// nodesEqual returns true if the node states are functionaly equal, ignoring status,
// version and other superfluous fields.
//
// This used to decide whether or not to propagate a node update to executor.
func nodesEqual(a, b *api.Node) bool {
	a, b = a.Copy(), b.Copy()

	a.Status, b.Status = api.NodeStatus{}, api.NodeStatus{}
	a.Meta, b.Meta = api.Meta{}, api.Meta{}

	return reflect.DeepEqual(a, b)
}
<|endoftext|>


<commit_before>package sqlstore

import (
	"testing"
	"time"

	"github.com/go-xorm/xorm"
	. "github.com/smartystreets/goconvey/convey"

	"github.com/grafana/grafana/pkg/components/simplejson"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
)

func TestDashboardSnapshotDBAccess(t *testing.T) {

	Convey("Testing DashboardSnapshot data access", t, func() {
		InitTestDB(t)

		Convey("Given saved snapshot", func() {
			cmd := m.CreateDashboardSnapshotCommand{
				Key: "hej",
				Dashboard: simplejson.NewFromAny(map[string]interface{}{
					"hello": "mupp",
				}),
				UserId: 1000,
				OrgId:  1,
			}
			err := CreateDashboardSnapshot(&cmd)
			So(err, ShouldBeNil)

			Convey("Should be able to get snapshot by key", func() {
				query := m.GetDashboardSnapshotQuery{Key: "hej"}
				err = GetDashboardSnapshot(&query)
				So(err, ShouldBeNil)

				So(query.Result, ShouldNotBeNil)
				So(query.Result.Dashboard.Get("hello").MustString(), ShouldEqual, "mupp")
			})

			Convey("And the user has the admin role", func() {
				Convey("Should return all the snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_ADMIN},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 1)
				})
			})

			Convey("And the user has the editor role and has created a snapshot", func() {
				Convey("Should return all the snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR, UserId: 1000},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 1)
				})
			})

			Convey("And the user has the editor role and has not created any snapshot", func() {
				Convey("Should not return any snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR, UserId: 2},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 0)
				})
			})

			Convey("And the user is anonymous", func() {
				cmd := m.CreateDashboardSnapshotCommand{
					Key:       "strangesnapshotwithuserid0",
					DeleteKey: "adeletekey",
					Dashboard: simplejson.NewFromAny(map[string]interface{}{
						"hello": "mupp",
					}),
					UserId: 0,
					OrgId:  1,
				}
				err := CreateDashboardSnapshot(&cmd)
				So(err, ShouldBeNil)

				Convey("Should not return any snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR, IsAnonymous: true, UserId: 0},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 0)
				})
			})
		})
	})
}

func TestDeleteExpiredSnapshots(t *testing.T) {
	x := InitTestDB(t)

	Convey("Testing dashboard snapshots clean up", t, func() {
		setting.SnapShotRemoveExpired = true

		notExpiredsnapshot := createTestSnapshot(x, "key1", 1200)
		createTestSnapshot(x, "key2", -1200)
		createTestSnapshot(x, "key3", -1200)

		err := DeleteExpiredSnapshots(&m.DeleteExpiredSnapshotsCommand{})
		So(err, ShouldBeNil)

		query := m.GetDashboardSnapshotsQuery{
			OrgId:        1,
			SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_ADMIN},
		}
		err = SearchDashboardSnapshots(&query)
		So(err, ShouldBeNil)

		So(len(query.Result), ShouldEqual, 1)
		So(query.Result[0].Key, ShouldEqual, notExpiredsnapshot.Key)

		err = DeleteExpiredSnapshots(&m.DeleteExpiredSnapshotsCommand{})
		So(err, ShouldBeNil)

		query = m.GetDashboardSnapshotsQuery{
			OrgId:        1,
			SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_ADMIN},
		}
		SearchDashboardSnapshots(&query)

		So(len(query.Result), ShouldEqual, 1)
		So(query.Result[0].Key, ShouldEqual, notExpiredsnapshot.Key)
	})
}

func createTestSnapshot(x *xorm.Engine, key string, expires int64) *m.DashboardSnapshot {
	cmd := m.CreateDashboardSnapshotCommand{
		Key:       key,
		DeleteKey: "delete" + key,
		Dashboard: simplejson.NewFromAny(map[string]interface{}{
			"hello": "mupp",
		}),
		UserId:  1000,
		OrgId:   1,
		Expires: expires,
	}
	err := CreateDashboardSnapshot(&cmd)
	So(err, ShouldBeNil)

	// Set expiry date manually - to be able to create expired snapshots
	if expires < 0 {
		expireDate := time.Now().Add(time.Second * time.Duration(expires))
		_, err = x.Exec("UPDATE dashboard_snapshot SET expires = ? WHERE id = ?", expireDate, cmd.Result.Id)
		So(err, ShouldBeNil)
	}

	return cmd.Result
}
<commit_msg>test: increase expire time to avoid tz issues in tests<commit_after>package sqlstore

import (
	"testing"
	"time"

	"github.com/go-xorm/xorm"
	. "github.com/smartystreets/goconvey/convey"

	"github.com/grafana/grafana/pkg/components/simplejson"
	m "github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/setting"
)

func TestDashboardSnapshotDBAccess(t *testing.T) {

	Convey("Testing DashboardSnapshot data access", t, func() {
		InitTestDB(t)

		Convey("Given saved snapshot", func() {
			cmd := m.CreateDashboardSnapshotCommand{
				Key: "hej",
				Dashboard: simplejson.NewFromAny(map[string]interface{}{
					"hello": "mupp",
				}),
				UserId: 1000,
				OrgId:  1,
			}
			err := CreateDashboardSnapshot(&cmd)
			So(err, ShouldBeNil)

			Convey("Should be able to get snapshot by key", func() {
				query := m.GetDashboardSnapshotQuery{Key: "hej"}
				err = GetDashboardSnapshot(&query)
				So(err, ShouldBeNil)

				So(query.Result, ShouldNotBeNil)
				So(query.Result.Dashboard.Get("hello").MustString(), ShouldEqual, "mupp")
			})

			Convey("And the user has the admin role", func() {
				Convey("Should return all the snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_ADMIN},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 1)
				})
			})

			Convey("And the user has the editor role and has created a snapshot", func() {
				Convey("Should return all the snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR, UserId: 1000},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 1)
				})
			})

			Convey("And the user has the editor role and has not created any snapshot", func() {
				Convey("Should not return any snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR, UserId: 2},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 0)
				})
			})

			Convey("And the user is anonymous", func() {
				cmd := m.CreateDashboardSnapshotCommand{
					Key:       "strangesnapshotwithuserid0",
					DeleteKey: "adeletekey",
					Dashboard: simplejson.NewFromAny(map[string]interface{}{
						"hello": "mupp",
					}),
					UserId: 0,
					OrgId:  1,
				}
				err := CreateDashboardSnapshot(&cmd)
				So(err, ShouldBeNil)

				Convey("Should not return any snapshots", func() {
					query := m.GetDashboardSnapshotsQuery{
						OrgId:        1,
						SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_EDITOR, IsAnonymous: true, UserId: 0},
					}
					err := SearchDashboardSnapshots(&query)
					So(err, ShouldBeNil)

					So(query.Result, ShouldNotBeNil)
					So(len(query.Result), ShouldEqual, 0)
				})
			})
		})
	})
}

func TestDeleteExpiredSnapshots(t *testing.T) {
	x := InitTestDB(t)

	Convey("Testing dashboard snapshots clean up", t, func() {
		setting.SnapShotRemoveExpired = true

		notExpiredsnapshot := createTestSnapshot(x, "key1", 48000)
		createTestSnapshot(x, "key2", -1200)
		createTestSnapshot(x, "key3", -1200)

		err := DeleteExpiredSnapshots(&m.DeleteExpiredSnapshotsCommand{})
		So(err, ShouldBeNil)

		query := m.GetDashboardSnapshotsQuery{
			OrgId:        1,
			SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_ADMIN},
		}
		err = SearchDashboardSnapshots(&query)
		So(err, ShouldBeNil)

		So(len(query.Result), ShouldEqual, 1)
		So(query.Result[0].Key, ShouldEqual, notExpiredsnapshot.Key)

		err = DeleteExpiredSnapshots(&m.DeleteExpiredSnapshotsCommand{})
		So(err, ShouldBeNil)

		query = m.GetDashboardSnapshotsQuery{
			OrgId:        1,
			SignedInUser: &m.SignedInUser{OrgRole: m.ROLE_ADMIN},
		}
		SearchDashboardSnapshots(&query)

		So(len(query.Result), ShouldEqual, 1)
		So(query.Result[0].Key, ShouldEqual, notExpiredsnapshot.Key)
	})
}

func createTestSnapshot(x *xorm.Engine, key string, expires int64) *m.DashboardSnapshot {
	cmd := m.CreateDashboardSnapshotCommand{
		Key:       key,
		DeleteKey: "delete" + key,
		Dashboard: simplejson.NewFromAny(map[string]interface{}{
			"hello": "mupp",
		}),
		UserId:  1000,
		OrgId:   1,
		Expires: expires,
	}
	err := CreateDashboardSnapshot(&cmd)
	So(err, ShouldBeNil)

	// Set expiry date manually - to be able to create expired snapshots
	if expires < 0 {
		expireDate := time.Now().Add(time.Second * time.Duration(expires))
		_, err = x.Exec("UPDATE dashboard_snapshot SET expires = ? WHERE id = ?", expireDate, cmd.Result.Id)
		So(err, ShouldBeNil)
	}

	return cmd.Result
}
<|endoftext|>


<commit_before>package slinga

import (
	"bufio"
	"github.com/golang/glog"
	"os/exec"
	"strings"
)

// HelmCodeExecutor is an executor that uses Helm for deployment of apps on kubernetes
type HelmCodeExecutor struct {
	Code *Code
}

func HelmName(str string) string {
	r := strings.NewReplacer("#", "-", "_", "-")
	return r.Replace(str)
}

// Install for HelmCodeExecutor runs "helm install" for the corresponding helm chart
func (executor HelmCodeExecutor) Install(key string, content map[string]map[string]string) error {
	uid := HelmName(key)

	chartName := content["chart"]["name"]

	// TODO(slukjanov): Replace with marshalling all params to temp file (YAML)
	setValues := ""
	if params, ok := content["params"]; ok {
		for key, value := range params {
			setValues += key + "=" + value + ","
		}
	}

	helmArgs := []string{"install", "--name", uid}
	if len(setValues) > 0 {
		helmArgs = append(helmArgs, "--set", setValues)
	}
	if version, ok := content["chart"]["version"]; ok {
		helmArgs = append(helmArgs, "--version", version)
	}
	if namespace, ok := content["chart"]["namespace"]; ok {
		helmArgs = append(helmArgs, "--namespace", namespace)
	} else {
		helmArgs = append(helmArgs, "--namespace", "aptomi")
	}
	helmArgs = append(helmArgs, chartName)

	return runHelmCmd(helmArgs...)
}

// Update for HelmCodeExecutor runs "helm update" for the corresponding helm chart
func (executor HelmCodeExecutor) Update(key string, labels LabelSet) error {
	// TODO: implement update method
	return nil
}

// Destroy for HelmCodeExecutor runs "helm delete" for the corresponding helm chart
func (executor HelmCodeExecutor) Destroy(key string) error {
	uid := HelmName(key)

	return runHelmCmd("delete", "--purge", uid)
}

func runHelmCmd(helmArgs ...string) error {
	return runCmd("helm", helmArgs...)
}

func runCmd(cmdName string, cmdArgs ...string) error {
	cmd := exec.Command(cmdName, cmdArgs...)
	glog.Infof("Running command '%s' with args: %s", cmdName, cmdArgs)

	cmdStdoutReader, err := cmd.StdoutPipe()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	cmdStdoutScanner := bufio.NewScanner(cmdStdoutReader)
	go func() {
		for cmdStdoutScanner.Scan() {
			glog.Infof("%s out | %s\n", cmdName, cmdStdoutScanner.Text())
		}
	}()

	cmdStderrReader, err := cmd.StderrPipe()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	cmdStderrScanner := bufio.NewScanner(cmdStderrReader)
	go func() {
		for cmdStderrScanner.Scan() {
			glog.Infof("%s err | %s\n", cmdName, cmdStderrScanner.Text())
		}
	}()

	err = cmd.Start()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	err = cmd.Wait()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	return nil
}
<commit_msg>Use tmp file to pass helm params<commit_after>package slinga

import (
	"bufio"
	"github.com/golang/glog"
	"os/exec"
	"strings"
	"io/ioutil"
)

// HelmCodeExecutor is an executor that uses Helm for deployment of apps on kubernetes
type HelmCodeExecutor struct {
	Code *Code
}

func HelmName(str string) string {
	r := strings.NewReplacer("#", "-", "_", "-")
	return r.Replace(str)
}

// Install for HelmCodeExecutor runs "helm install" for the corresponding helm chart
func (executor HelmCodeExecutor) Install(key string, content map[string]map[string]string) error {
	uid := HelmName(key)

	chartName := content["chart"]["name"]

	// TODO(slukjanov): Replace with marshalling all params to temp file (YAML)
	setValues := ""
	if params, ok := content["params"]; ok {
		for key, value := range params {
			setValues += key + "=" + value + ","
		}
	}

	tmpFile, _ /*err*/ := ioutil.TempFile("", "aptomi-helm-params")
	//TODO: slukjanov: defer os.Remove(tmpFile.Name())
	glog.Info("Temp helm params: ", tmpFile.Name())
	if params, ok := content["params"]; ok {
		content := []byte(serializeObject(params))
		if _, err := tmpFile.Write(content); err != nil {
			glog.Info(err)
		}
		if err := tmpFile.Close(); err != nil {
			glog.Info(err)
		}
	}

	helmArgs := []string{"install", "--name", uid}
	if len(setValues) > 0 {
		helmArgs = append(helmArgs, "--set", setValues)
	}
	if version, ok := content["chart"]["version"]; ok {
		helmArgs = append(helmArgs, "--version", version)
	}
	if namespace, ok := content["chart"]["namespace"]; ok {
		helmArgs = append(helmArgs, "--namespace", namespace)
	} else {
		helmArgs = append(helmArgs, "--namespace", "aptomi")
	}
	helmArgs = append(helmArgs, chartName)

	glog.Infof("Running Helm with args: %s", helmArgs)
	return runHelmCmd(helmArgs...)
	//return nil
}

// Update for HelmCodeExecutor runs "helm update" for the corresponding helm chart
func (executor HelmCodeExecutor) Update(key string, labels LabelSet) error {
	// TODO: implement update method
	return nil
}

// Destroy for HelmCodeExecutor runs "helm delete" for the corresponding helm chart
func (executor HelmCodeExecutor) Destroy(key string) error {
	uid := HelmName(key)

	return runHelmCmd("delete", "--purge", uid)
}

func runHelmCmd(helmArgs ...string) error {
	return runCmd("helm", helmArgs...)
}

func runCmd(cmdName string, cmdArgs ...string) error {
	cmd := exec.Command(cmdName, cmdArgs...)
	glog.Infof("Running command '%s' with args: %s", cmdName, cmdArgs)

	cmdStdoutReader, err := cmd.StdoutPipe()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	cmdStdoutScanner := bufio.NewScanner(cmdStdoutReader)
	go func() {
		for cmdStdoutScanner.Scan() {
			glog.Infof("%s out | %s\n", cmdName, cmdStdoutScanner.Text())
		}
	}()

	cmdStderrReader, err := cmd.StderrPipe()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	cmdStderrScanner := bufio.NewScanner(cmdStderrReader)
	go func() {
		for cmdStderrScanner.Scan() {
			glog.Infof("%s err | %s\n", cmdName, cmdStderrScanner.Text())
		}
	}()

	err = cmd.Start()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	err = cmd.Wait()
	if err != nil {
		glog.Errorf("Failed running command '%s' (with args: %s): %s", cmdName, cmdArgs, err)
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Ceph-CSI 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 util

import (
	"fmt"
	"io/ioutil"
	"os"
)

const (
	credUserID           = "userID"
	credUserKey          = "userKey"
	credAdminID          = "adminID"
	credAdminKey         = "adminKey"
	credMonitors         = "monitors"
	tmpKeyFileLocation   = "/tmp/csi/keys"
	tmpKeyFileNamePrefix = "keyfile-"
)

type Credentials struct {
	ID      string
	KeyFile string
}

func storeKey(key string) (string, error) {
	tmpfile, err := ioutil.TempFile(tmpKeyFileLocation, tmpKeyFileNamePrefix)
	if err != nil {
		return "", fmt.Errorf("error creating a temporary keyfile (%s)", err)
	}
	defer func() {
		if err != nil {
			// don't complain about unhandled error
			_ = os.Remove(tmpfile.Name())
		}
	}()

	if _, err = tmpfile.Write([]byte(key)); err != nil {
		return "", fmt.Errorf("error writing key to temporary keyfile (%s)", err)
	}

	keyFile := tmpfile.Name()
	if keyFile == "" {
		err = fmt.Errorf("error reading temporary filename for key (%s)", err)
		return "", err
	}

	if err = tmpfile.Close(); err != nil {
		return "", fmt.Errorf("error closing temporary filename (%s)", err)
	}

	return keyFile, nil
}

func newCredentialsFromSecret(idField, keyField string, secrets map[string]string) (*Credentials, error) {
	var (
		c  = &Credentials{}
		ok bool
	)

	if c.ID, ok = secrets[idField]; !ok {
		return nil, fmt.Errorf("missing ID field '%s' in secrets", idField)
	}

	key := secrets[keyField]
	if key == "" {
		return nil, fmt.Errorf("missing key field '%s' in secrets", keyField)
	}

	keyFile, err := storeKey(key)
	if err == nil {
		c.KeyFile = keyFile
	}

	return c, err
}

func (cr *Credentials) DeleteCredentials() {
	// don't complain about unhandled error
	_ = os.Remove(cr.KeyFile)
}

func NewUserCredentials(secrets map[string]string) (*Credentials, error) {
	return newCredentialsFromSecret(credUserID, credUserKey, secrets)
}

func NewAdminCredentials(secrets map[string]string) (*Credentials, error) {
	return newCredentialsFromSecret(credAdminID, credAdminKey, secrets)
}

func NewCredentials(id, key string) (*Credentials, error) {
	var c = &Credentials{}

	c.ID = id
	keyFile, err := storeKey(key)
	if err == nil {
		c.KeyFile = keyFile
	}

	return c, err
}

func GetMonValFromSecret(secrets map[string]string) (string, error) {
	if mons, ok := secrets[credMonitors]; ok {
		return mons, nil
	}
	return "", fmt.Errorf("missing %q", credMonitors)
}
<commit_msg>Add a check for nil secrets<commit_after>/*
Copyright 2018 The Ceph-CSI 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 util

import (
	"errors"
	"fmt"
	"io/ioutil"
	"os"
)

const (
	credUserID           = "userID"
	credUserKey          = "userKey"
	credAdminID          = "adminID"
	credAdminKey         = "adminKey"
	credMonitors         = "monitors"
	tmpKeyFileLocation   = "/tmp/csi/keys"
	tmpKeyFileNamePrefix = "keyfile-"
)

type Credentials struct {
	ID      string
	KeyFile string
}

func storeKey(key string) (string, error) {
	tmpfile, err := ioutil.TempFile(tmpKeyFileLocation, tmpKeyFileNamePrefix)
	if err != nil {
		return "", fmt.Errorf("error creating a temporary keyfile (%s)", err)
	}
	defer func() {
		if err != nil {
			// don't complain about unhandled error
			_ = os.Remove(tmpfile.Name())
		}
	}()

	if _, err = tmpfile.Write([]byte(key)); err != nil {
		return "", fmt.Errorf("error writing key to temporary keyfile (%s)", err)
	}

	keyFile := tmpfile.Name()
	if keyFile == "" {
		err = fmt.Errorf("error reading temporary filename for key (%s)", err)
		return "", err
	}

	if err = tmpfile.Close(); err != nil {
		return "", fmt.Errorf("error closing temporary filename (%s)", err)
	}

	return keyFile, nil
}

func newCredentialsFromSecret(idField, keyField string, secrets map[string]string) (*Credentials, error) {
	var (
		c  = &Credentials{}
		ok bool
	)

	if len(secrets) == 0 {
		return nil, errors.New("provided secret is empty")
	}
	if c.ID, ok = secrets[idField]; !ok {
		return nil, fmt.Errorf("missing ID field '%s' in secrets", idField)
	}

	key := secrets[keyField]
	if key == "" {
		return nil, fmt.Errorf("missing key field '%s' in secrets", keyField)
	}

	keyFile, err := storeKey(key)
	if err == nil {
		c.KeyFile = keyFile
	}

	return c, err
}

func (cr *Credentials) DeleteCredentials() {
	// don't complain about unhandled error
	_ = os.Remove(cr.KeyFile)
}

func NewUserCredentials(secrets map[string]string) (*Credentials, error) {
	return newCredentialsFromSecret(credUserID, credUserKey, secrets)
}

func NewAdminCredentials(secrets map[string]string) (*Credentials, error) {
	return newCredentialsFromSecret(credAdminID, credAdminKey, secrets)
}

func NewCredentials(id, key string) (*Credentials, error) {
	var c = &Credentials{}

	c.ID = id
	keyFile, err := storeKey(key)
	if err == nil {
		c.KeyFile = keyFile
	}

	return c, err
}

func GetMonValFromSecret(secrets map[string]string) (string, error) {
	if mons, ok := secrets[credMonitors]; ok {
		return mons, nil
	}
	return "", fmt.Errorf("missing %q", credMonitors)
}
<|endoftext|>


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

// Binary local_server provides a local 9P2000.L server for the p9 package.
//
// To use, first start the server:
//     local_server /tmp/my_bind_addr
//
// Then, connect using the Linux 9P filesystem:
//     mount -t 9p -o trans=unix,version=9P2000.L /tmp/my_bind_addr /mnt
//
// This package also serves as an examplar.
package main

import (
	"os"
	"path"
	"syscall"

	"gvisor.googlesource.com/gvisor/pkg/fd"
	"gvisor.googlesource.com/gvisor/pkg/log"
	"gvisor.googlesource.com/gvisor/pkg/p9"
	"gvisor.googlesource.com/gvisor/pkg/unet"
)

// local wraps a local file.
type local struct {
	p9.DefaultWalkGetAttr

	path string
	file *os.File
}

// info constructs a QID for this file.
func (l *local) info() (p9.QID, os.FileInfo, error) {
	var (
		qid p9.QID
		fi  os.FileInfo
		err error
	)

	// Stat the file.
	if l.file != nil {
		fi, err = l.file.Stat()
	} else {
		fi, err = os.Lstat(l.path)
	}
	if err != nil {
		log.Warningf("error stating %#v: %v", l, err)
		return qid, nil, err
	}

	// Construct the QID type.
	qid.Type = p9.ModeFromOS(fi.Mode()).QIDType()

	// Save the path from the Ino.
	qid.Path = fi.Sys().(*syscall.Stat_t).Ino
	return qid, fi, nil
}

// Attach implements p9.Attacher.Attach.
func (l *local) Attach(name string) (p9.File, error) {
	return &local{path: path.Clean(name)}, nil
}

// Walk implements p9.File.Walk.
func (l *local) Walk(names []string) ([]p9.QID, p9.File, error) {
	var qids []p9.QID
	last := &local{path: l.path}
	for _, name := range names {
		c := &local{path: path.Join(last.path, name)}
		qid, _, err := c.info()
		if err != nil {
			return nil, nil, err
		}
		qids = append(qids, qid)
		last = c
	}
	return qids, last, nil
}

// StatFS implements p9.File.StatFS.
//
// Not implemented.
func (l *local) StatFS() (p9.FSStat, error) {
	return p9.FSStat{}, syscall.ENOSYS
}

// FSync implements p9.File.FSync.
func (l *local) FSync() error {
	return l.file.Sync()
}

// GetAttr implements p9.File.GetAttr.
//
// Not fully implemented.
func (l *local) GetAttr(req p9.AttrMask) (p9.QID, p9.AttrMask, p9.Attr, error) {
	qid, fi, err := l.info()
	if err != nil {
		return qid, p9.AttrMask{}, p9.Attr{}, err
	}

	stat := fi.Sys().(*syscall.Stat_t)
	attr := p9.Attr{
		Mode:             p9.FileMode(stat.Mode),
		UID:              p9.UID(stat.Uid),
		GID:              p9.GID(stat.Gid),
		NLink:            stat.Nlink,
		RDev:             stat.Rdev,
		Size:             uint64(stat.Size),
		BlockSize:        uint64(stat.Blksize),
		Blocks:           uint64(stat.Blocks),
		ATimeSeconds:     uint64(stat.Atim.Sec),
		ATimeNanoSeconds: uint64(stat.Atim.Nsec),
		MTimeSeconds:     uint64(stat.Mtim.Sec),
		MTimeNanoSeconds: uint64(stat.Mtim.Nsec),
		CTimeSeconds:     uint64(stat.Ctim.Sec),
		CTimeNanoSeconds: uint64(stat.Ctim.Nsec),
	}
	valid := p9.AttrMask{
		Mode:   true,
		UID:    true,
		GID:    true,
		NLink:  true,
		RDev:   true,
		Size:   true,
		Blocks: true,
		ATime:  true,
		MTime:  true,
		CTime:  true,
	}

	return qid, valid, attr, nil
}

// SetAttr implements p9.File.SetAttr.
//
// Not implemented.
func (l *local) SetAttr(valid p9.SetAttrMask, attr p9.SetAttr) error {
	return syscall.ENOSYS
}

// Remove implements p9.File.Remove.
//
// Not implemented.
func (l *local) Remove() error {
	return syscall.ENOSYS
}

// Rename implements p9.File.Rename.
//
// Not implemented.
func (l *local) Rename(directory p9.File, name string) error {
	return syscall.ENOSYS
}

// Close implements p9.File.Close.
func (l *local) Close() error {
	if l.file != nil {
		return l.file.Close()
	}
	return nil
}

// Open implements p9.File.Open.
func (l *local) Open(mode p9.OpenFlags) (*fd.FD, p9.QID, uint32, error) {
	qid, _, err := l.info()
	if err != nil {
		return nil, qid, 0, err
	}

	// Do the actual open.
	f, err := os.OpenFile(l.path, int(mode), 0)
	if err != nil {
		return nil, qid, 0, err
	}
	l.file = f

	// Note: we don't send the local file for this server.
	return nil, qid, 4096, nil
}

// Read implements p9.File.Read.
func (l *local) ReadAt(p []byte, offset uint64) (int, error) {
	return l.file.ReadAt(p, int64(offset))
}

// Write implements p9.File.Write.
func (l *local) WriteAt(p []byte, offset uint64) (int, error) {
	return l.file.WriteAt(p, int64(offset))
}

// Create implements p9.File.Create.
func (l *local) Create(name string, mode p9.OpenFlags, permissions p9.FileMode, _ p9.UID, _ p9.GID) (*fd.FD, p9.File, p9.QID, uint32, error) {
	f, err := os.OpenFile(l.path, int(mode)|syscall.O_CREAT|syscall.O_EXCL, os.FileMode(permissions))
	if err != nil {
		return nil, nil, p9.QID{}, 0, err
	}

	l2 := &local{path: path.Join(l.path, name), file: f}
	qid, _, err := l2.info()
	if err != nil {
		l2.Close()
		return nil, nil, p9.QID{}, 0, err
	}

	return nil, l2, qid, 4096, nil
}

// Mkdir implements p9.File.Mkdir.
//
// Not properly implemented.
func (l *local) Mkdir(name string, permissions p9.FileMode, _ p9.UID, _ p9.GID) (p9.QID, error) {
	if err := os.Mkdir(path.Join(l.path, name), os.FileMode(permissions)); err != nil {
		return p9.QID{}, err
	}

	// Blank QID.
	return p9.QID{}, nil
}

// Symlink implements p9.File.Symlink.
//
// Not properly implemented.
func (l *local) Symlink(oldname string, newname string, _ p9.UID, _ p9.GID) (p9.QID, error) {
	if err := os.Symlink(oldname, path.Join(l.path, newname)); err != nil {
		return p9.QID{}, err
	}

	// Blank QID.
	return p9.QID{}, nil
}

// Link implements p9.File.Link.
//
// Not properly implemented.
func (l *local) Link(target p9.File, newname string) error {
	return os.Link(target.(*local).path, path.Join(l.path, newname))
}

// Mknod implements p9.File.Mknod.
//
// Not implemented.
func (l *local) Mknod(name string, permissions p9.FileMode, major uint32, minor uint32, _ p9.UID, _ p9.GID) (p9.QID, error) {
	return p9.QID{}, syscall.ENOSYS
}

// RenameAt implements p9.File.RenameAt.
//
// Not implemented.
func (l *local) RenameAt(oldname string, newdir p9.File, newname string) error {
	return syscall.ENOSYS
}

// UnlinkAt implements p9.File.UnlinkAt.
//
// Not implemented.
func (l *local) UnlinkAt(name string, flags uint32) error {
	return syscall.ENOSYS
}

// Readdir implements p9.File.Readdir.
func (l *local) Readdir(offset uint64, count uint32) ([]p9.Dirent, error) {
	// We only do *all* dirents in single shot.
	const maxDirentBuffer = 1024 * 1024
	buf := make([]byte, maxDirentBuffer)
	n, err := syscall.ReadDirent(int(l.file.Fd()), buf)
	if err != nil {
		// Return zero entries.
		return nil, nil
	}

	// Parse the entries; note that we read up to offset+count here.
	_, newCount, newNames := syscall.ParseDirent(buf[:n], int(offset)+int(count), nil)
	var dirents []p9.Dirent
	for i := int(offset); i >= 0 && i < newCount; i++ {
		entry := local{path: path.Join(l.path, newNames[i])}
		qid, _, err := entry.info()
		if err != nil {
			continue
		}
		dirents = append(dirents, p9.Dirent{
			QID:    qid,
			Type:   qid.Type,
			Name:   newNames[i],
			Offset: uint64(i + 1),
		})
	}

	return dirents, nil
}

// Readlink implements p9.File.Readlink.
//
// Not properly implemented.
func (l *local) Readlink() (string, error) {
	return os.Readlink(l.path)
}

// Flush implements p9.File.Flush.
func (l *local) Flush() error {
	return nil
}

// Connect implements p9.File.Connect.
func (l *local) Connect(p9.ConnectFlags) (*fd.FD, error) {
	return nil, syscall.ECONNREFUSED
}

func main() {
	log.SetLevel(log.Debug)

	if len(os.Args) != 2 {
		log.Warningf("usage: %s <bind-addr>", os.Args[0])
		os.Exit(1)
	}

	// Bind and listen on the socket.
	serverSocket, err := unet.BindAndListen(os.Args[1], false)
	if err != nil {
		log.Warningf("err binding: %v", err)
		os.Exit(1)
	}

	// Run the server.
	s := p9.NewServer(&local{})
	s.Serve(serverSocket)
}

var (
	_ p9.File = &local{}
)
<commit_msg>Drop version option from mount command<commit_after>// Copyright 2018 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Binary local_server provides a local 9P2000.L server for the p9 package.
//
// To use, first start the server:
//     local_server /tmp/my_bind_addr
//
// Then, connect using the Linux 9P filesystem:
//     mount -t 9p -o trans=unix /tmp/my_bind_addr /mnt
//
// This package also serves as an examplar.
package main

import (
	"os"
	"path"
	"syscall"

	"gvisor.googlesource.com/gvisor/pkg/fd"
	"gvisor.googlesource.com/gvisor/pkg/log"
	"gvisor.googlesource.com/gvisor/pkg/p9"
	"gvisor.googlesource.com/gvisor/pkg/unet"
)

// local wraps a local file.
type local struct {
	p9.DefaultWalkGetAttr

	path string
	file *os.File
}

// info constructs a QID for this file.
func (l *local) info() (p9.QID, os.FileInfo, error) {
	var (
		qid p9.QID
		fi  os.FileInfo
		err error
	)

	// Stat the file.
	if l.file != nil {
		fi, err = l.file.Stat()
	} else {
		fi, err = os.Lstat(l.path)
	}
	if err != nil {
		log.Warningf("error stating %#v: %v", l, err)
		return qid, nil, err
	}

	// Construct the QID type.
	qid.Type = p9.ModeFromOS(fi.Mode()).QIDType()

	// Save the path from the Ino.
	qid.Path = fi.Sys().(*syscall.Stat_t).Ino
	return qid, fi, nil
}

// Attach implements p9.Attacher.Attach.
func (l *local) Attach(name string) (p9.File, error) {
	return &local{path: path.Clean(name)}, nil
}

// Walk implements p9.File.Walk.
func (l *local) Walk(names []string) ([]p9.QID, p9.File, error) {
	var qids []p9.QID
	last := &local{path: l.path}
	for _, name := range names {
		c := &local{path: path.Join(last.path, name)}
		qid, _, err := c.info()
		if err != nil {
			return nil, nil, err
		}
		qids = append(qids, qid)
		last = c
	}
	return qids, last, nil
}

// StatFS implements p9.File.StatFS.
//
// Not implemented.
func (l *local) StatFS() (p9.FSStat, error) {
	return p9.FSStat{}, syscall.ENOSYS
}

// FSync implements p9.File.FSync.
func (l *local) FSync() error {
	return l.file.Sync()
}

// GetAttr implements p9.File.GetAttr.
//
// Not fully implemented.
func (l *local) GetAttr(req p9.AttrMask) (p9.QID, p9.AttrMask, p9.Attr, error) {
	qid, fi, err := l.info()
	if err != nil {
		return qid, p9.AttrMask{}, p9.Attr{}, err
	}

	stat := fi.Sys().(*syscall.Stat_t)
	attr := p9.Attr{
		Mode:             p9.FileMode(stat.Mode),
		UID:              p9.UID(stat.Uid),
		GID:              p9.GID(stat.Gid),
		NLink:            stat.Nlink,
		RDev:             stat.Rdev,
		Size:             uint64(stat.Size),
		BlockSize:        uint64(stat.Blksize),
		Blocks:           uint64(stat.Blocks),
		ATimeSeconds:     uint64(stat.Atim.Sec),
		ATimeNanoSeconds: uint64(stat.Atim.Nsec),
		MTimeSeconds:     uint64(stat.Mtim.Sec),
		MTimeNanoSeconds: uint64(stat.Mtim.Nsec),
		CTimeSeconds:     uint64(stat.Ctim.Sec),
		CTimeNanoSeconds: uint64(stat.Ctim.Nsec),
	}
	valid := p9.AttrMask{
		Mode:   true,
		UID:    true,
		GID:    true,
		NLink:  true,
		RDev:   true,
		Size:   true,
		Blocks: true,
		ATime:  true,
		MTime:  true,
		CTime:  true,
	}

	return qid, valid, attr, nil
}

// SetAttr implements p9.File.SetAttr.
//
// Not implemented.
func (l *local) SetAttr(valid p9.SetAttrMask, attr p9.SetAttr) error {
	return syscall.ENOSYS
}

// Remove implements p9.File.Remove.
//
// Not implemented.
func (l *local) Remove() error {
	return syscall.ENOSYS
}

// Rename implements p9.File.Rename.
//
// Not implemented.
func (l *local) Rename(directory p9.File, name string) error {
	return syscall.ENOSYS
}

// Close implements p9.File.Close.
func (l *local) Close() error {
	if l.file != nil {
		return l.file.Close()
	}
	return nil
}

// Open implements p9.File.Open.
func (l *local) Open(mode p9.OpenFlags) (*fd.FD, p9.QID, uint32, error) {
	qid, _, err := l.info()
	if err != nil {
		return nil, qid, 0, err
	}

	// Do the actual open.
	f, err := os.OpenFile(l.path, int(mode), 0)
	if err != nil {
		return nil, qid, 0, err
	}
	l.file = f

	// Note: we don't send the local file for this server.
	return nil, qid, 4096, nil
}

// Read implements p9.File.Read.
func (l *local) ReadAt(p []byte, offset uint64) (int, error) {
	return l.file.ReadAt(p, int64(offset))
}

// Write implements p9.File.Write.
func (l *local) WriteAt(p []byte, offset uint64) (int, error) {
	return l.file.WriteAt(p, int64(offset))
}

// Create implements p9.File.Create.
func (l *local) Create(name string, mode p9.OpenFlags, permissions p9.FileMode, _ p9.UID, _ p9.GID) (*fd.FD, p9.File, p9.QID, uint32, error) {
	f, err := os.OpenFile(l.path, int(mode)|syscall.O_CREAT|syscall.O_EXCL, os.FileMode(permissions))
	if err != nil {
		return nil, nil, p9.QID{}, 0, err
	}

	l2 := &local{path: path.Join(l.path, name), file: f}
	qid, _, err := l2.info()
	if err != nil {
		l2.Close()
		return nil, nil, p9.QID{}, 0, err
	}

	return nil, l2, qid, 4096, nil
}

// Mkdir implements p9.File.Mkdir.
//
// Not properly implemented.
func (l *local) Mkdir(name string, permissions p9.FileMode, _ p9.UID, _ p9.GID) (p9.QID, error) {
	if err := os.Mkdir(path.Join(l.path, name), os.FileMode(permissions)); err != nil {
		return p9.QID{}, err
	}

	// Blank QID.
	return p9.QID{}, nil
}

// Symlink implements p9.File.Symlink.
//
// Not properly implemented.
func (l *local) Symlink(oldname string, newname string, _ p9.UID, _ p9.GID) (p9.QID, error) {
	if err := os.Symlink(oldname, path.Join(l.path, newname)); err != nil {
		return p9.QID{}, err
	}

	// Blank QID.
	return p9.QID{}, nil
}

// Link implements p9.File.Link.
//
// Not properly implemented.
func (l *local) Link(target p9.File, newname string) error {
	return os.Link(target.(*local).path, path.Join(l.path, newname))
}

// Mknod implements p9.File.Mknod.
//
// Not implemented.
func (l *local) Mknod(name string, permissions p9.FileMode, major uint32, minor uint32, _ p9.UID, _ p9.GID) (p9.QID, error) {
	return p9.QID{}, syscall.ENOSYS
}

// RenameAt implements p9.File.RenameAt.
//
// Not implemented.
func (l *local) RenameAt(oldname string, newdir p9.File, newname string) error {
	return syscall.ENOSYS
}

// UnlinkAt implements p9.File.UnlinkAt.
//
// Not implemented.
func (l *local) UnlinkAt(name string, flags uint32) error {
	return syscall.ENOSYS
}

// Readdir implements p9.File.Readdir.
func (l *local) Readdir(offset uint64, count uint32) ([]p9.Dirent, error) {
	// We only do *all* dirents in single shot.
	const maxDirentBuffer = 1024 * 1024
	buf := make([]byte, maxDirentBuffer)
	n, err := syscall.ReadDirent(int(l.file.Fd()), buf)
	if err != nil {
		// Return zero entries.
		return nil, nil
	}

	// Parse the entries; note that we read up to offset+count here.
	_, newCount, newNames := syscall.ParseDirent(buf[:n], int(offset)+int(count), nil)
	var dirents []p9.Dirent
	for i := int(offset); i >= 0 && i < newCount; i++ {
		entry := local{path: path.Join(l.path, newNames[i])}
		qid, _, err := entry.info()
		if err != nil {
			continue
		}
		dirents = append(dirents, p9.Dirent{
			QID:    qid,
			Type:   qid.Type,
			Name:   newNames[i],
			Offset: uint64(i + 1),
		})
	}

	return dirents, nil
}

// Readlink implements p9.File.Readlink.
//
// Not properly implemented.
func (l *local) Readlink() (string, error) {
	return os.Readlink(l.path)
}

// Flush implements p9.File.Flush.
func (l *local) Flush() error {
	return nil
}

// Connect implements p9.File.Connect.
func (l *local) Connect(p9.ConnectFlags) (*fd.FD, error) {
	return nil, syscall.ECONNREFUSED
}

func main() {
	log.SetLevel(log.Debug)

	if len(os.Args) != 2 {
		log.Warningf("usage: %s <bind-addr>", os.Args[0])
		os.Exit(1)
	}

	// Bind and listen on the socket.
	serverSocket, err := unet.BindAndListen(os.Args[1], false)
	if err != nil {
		log.Warningf("err binding: %v", err)
		os.Exit(1)
	}

	// Run the server.
	s := p9.NewServer(&local{})
	s.Serve(serverSocket)
}

var (
	_ p9.File = &local{}
)
<|endoftext|>


<commit_before>package ai

import (
	"fmt"
	"io"
	"text/tabwriter"

	"github.com/nelhage/taktician/bitboard"
	"github.com/nelhage/taktician/tak"
)

type Weights struct {
	TopFlat  int
	Standing int
	Capstone int

	Flat     int
	Captured int

	Liberties int

	Tempo int

	Groups [8]int
}

var DefaultWeights = Weights{
	TopFlat:  300,
	Standing: 200,
	Capstone: 300,

	Flat:      100,
	Liberties: 25,

	Captured: 25,

	Tempo: 250,

	Groups: [8]int{
		0,   // 0
		0,   // 1
		0,   // 2
		100, // 3
		300, // 4
		500, // 5
	},
}

func MakeEvaluator(w *Weights) EvaluationFunc {
	return func(m *MinimaxAI, p *tak.Position) int64 {
		return evaluate(w, m, p)
	}
}

var DefaultEvaluate = MakeEvaluator(&DefaultWeights)

func evaluate(w *Weights, m *MinimaxAI, p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		var pieces int64
		if winner == tak.White {
			pieces = int64(p.WhiteStones())
		} else {
			pieces = int64(p.BlackStones())
		}
		switch winner {
		case tak.NoColor:
			return 0
		case p.ToMove():
			return maxEval - int64(p.MoveNumber()) + pieces
		default:
			return minEval + int64(p.MoveNumber()) - pieces
		}
	}
	mine, theirs := 0, 0
	me := p.ToMove()
	addw := func(c tak.Color, w int) {
		if c == me {
			mine += w
		} else {
			theirs += w
		}
	}
	addw(p.ToMove(), w.Tempo)
	analysis := p.Analysis()
	for x := 0; x < p.Size(); x++ {
		for y := 0; y < p.Size(); y++ {
			sq := p.At(x, y)
			if len(sq) == 0 {
				continue
			}
			switch sq[0].Kind() {
			case tak.Standing:
				addw(sq[0].Color(), w.Standing)
			case tak.Flat:
				addw(sq[0].Color(), w.TopFlat)
			case tak.Capstone:
				addw(sq[0].Color(), w.Capstone)
			}

			for i, stone := range sq {
				if i > 0 && i < p.Size() {
					addw(sq[0].Color(), w.Captured)
				}
				if stone.Kind() == tak.Flat {
					addw(stone.Color(), w.Flat)
				}
			}
		}
	}

	addw(tak.White, m.scoreGroups(analysis.WhiteGroups, w))
	addw(tak.Black, m.scoreGroups(analysis.BlackGroups, w))

	wr := p.White &^ p.Standing
	br := p.Black &^ p.Standing
	wl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.Black, wr) &^ wr)
	bl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.White, br) &^ br)
	addw(tak.White, w.Liberties*wl)
	addw(tak.Black, w.Liberties*bl)

	return int64(mine - theirs)
}

func (ai *MinimaxAI) scoreGroups(gs []uint64, ws *Weights) int {
	sc := 0
	for _, g := range gs {
		w, h := bitboard.Dimensions(&ai.c, g)

		sc += ws.Groups[w]
		sc += ws.Groups[h]
	}

	return sc
}

func ExplainScore(m *MinimaxAI, out io.Writer, p *tak.Position) {
	tw := tabwriter.NewWriter(out, 4, 8, 1, '\t', 0)
	fmt.Fprintf(tw, "\twhite\tblack\n")
	var scores [2]struct {
		flats    int
		standing int
		caps     int

		stones   int
		captured int
	}
	for x := 0; x < p.Size(); x++ {
		for y := 0; y < p.Size(); y++ {
			sq := p.At(x, y)
			if len(sq) == 0 {
				continue
			}
			switch sq[0].Kind() {
			case tak.Standing:
				if sq[0].Color() == tak.White {
					scores[0].standing++
				} else {
					scores[1].standing++
				}
			case tak.Flat:
				if sq[0].Color() == tak.White {
					scores[0].flats++
				} else {
					scores[1].flats++
				}
			case tak.Capstone:
				if sq[0].Color() == tak.White {
					scores[0].caps++
				} else {
					scores[1].caps++
				}
			}

			for i, stone := range sq {
				if i > 0 && i < p.Size() {
					if sq[0].Color() == tak.White {
						scores[0].captured++
					} else {
						scores[1].captured++
					}
				}
				if stone.Kind() == tak.Flat {
					if sq[0].Color() == tak.White {
						scores[0].stones++
					} else {
						scores[1].stones++
					}
				}
			}
		}
	}
	fmt.Fprintf(tw, "flats\t%d\t%d\n", scores[0].flats, scores[1].flats)
	fmt.Fprintf(tw, "standing\t%d\t%d\n", scores[0].standing, scores[1].standing)
	fmt.Fprintf(tw, "caps\t%d\t%d\n", scores[0].caps, scores[1].caps)
	fmt.Fprintf(tw, "capured\t%d\t%d\n", scores[0].captured, scores[1].captured)
	fmt.Fprintf(tw, "stones\t%d\t%d\n", scores[0].stones, scores[1].stones)

	analysis := p.Analysis()

	wr := p.White &^ p.Standing
	br := p.Black &^ p.Standing
	wl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.Black, wr) &^ wr)
	bl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.White, br) &^ br)

	fmt.Fprintf(tw, "liberties\t%d\t%d\n", wl, bl)

	for i, g := range analysis.WhiteGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t%dx%x\n", i, w, h)
	}
	for i, g := range analysis.BlackGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t\t%dx%x\n", i, w, h)
	}
	tw.Flush()
}
<commit_msg>direct bitboard evaluate<commit_after>package ai

import (
	"fmt"
	"io"
	"text/tabwriter"

	"github.com/nelhage/taktician/bitboard"
	"github.com/nelhage/taktician/tak"
)

type Weights struct {
	TopFlat  int
	Standing int
	Capstone int

	Flat     int
	Captured int

	Liberties int

	Tempo int

	Groups [8]int
}

var DefaultWeights = Weights{
	TopFlat:  400,
	Standing: 200,
	Capstone: 300,

	Flat:      100,
	Liberties: 25,

	Captured: 25,

	Tempo: 250,

	Groups: [8]int{
		0,   // 0
		0,   // 1
		0,   // 2
		100, // 3
		300, // 4
		500, // 5
	},
}

func MakeEvaluator(w *Weights) EvaluationFunc {
	return func(m *MinimaxAI, p *tak.Position) int64 {
		return evaluate(w, m, p)
	}
}

var DefaultEvaluate = MakeEvaluator(&DefaultWeights)

func evaluate(w *Weights, m *MinimaxAI, p *tak.Position) int64 {
	if over, winner := p.GameOver(); over {
		var pieces int64
		if winner == tak.White {
			pieces = int64(p.WhiteStones())
		} else {
			pieces = int64(p.BlackStones())
		}
		switch winner {
		case tak.NoColor:
			return 0
		case p.ToMove():
			return maxEval - int64(p.MoveNumber()) + pieces
		default:
			return minEval + int64(p.MoveNumber()) - pieces
		}
	}

	var ws, bs int64

	if p.ToMove() == tak.White {
		ws += int64(w.Tempo)
	} else {
		bs += int64(w.Tempo)
	}
	analysis := p.Analysis()

	ws += int64(bitboard.Popcount(p.White&^p.Caps&^p.Standing) * w.TopFlat)
	bs += int64(bitboard.Popcount(p.Black&^p.Caps&^p.Standing) * w.TopFlat)
	ws += int64(bitboard.Popcount(p.White&p.Standing) * w.Standing)
	bs += int64(bitboard.Popcount(p.Black&p.Standing) * w.Standing)
	ws += int64(bitboard.Popcount(p.White&p.Caps) * w.Capstone)
	bs += int64(bitboard.Popcount(p.Black&p.Caps) * w.Capstone)

	for i, h := range p.Height {
		if h == 0 {
			continue
		}
		s := p.Stacks[i] & ((1 << (h - 1)) - 1)
		bf := bitboard.Popcount(s)
		wf := int(h) - bf - 1
		ws += int64(wf * w.Flat)
		bs += int64(bf * w.Flat)
		if p.White&(1<<uint(i)) != 0 {
			ws += int64(int(h-1) * w.Captured)
		} else {
			bs += int64(int(h-1) * w.Captured)
		}
	}

	ws += int64(m.scoreGroups(analysis.WhiteGroups, w))
	bs += int64(m.scoreGroups(analysis.BlackGroups, w))

	wr := p.White &^ p.Standing
	br := p.Black &^ p.Standing
	wl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.Black, wr) &^ wr)
	bl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.White, br) &^ br)
	ws += int64(w.Liberties * wl)
	bs += int64(w.Liberties * bl)

	if p.ToMove() == tak.White {
		return ws - bs
	}
	return bs - ws
}

func (ai *MinimaxAI) scoreGroups(gs []uint64, ws *Weights) int {
	sc := 0
	for _, g := range gs {
		w, h := bitboard.Dimensions(&ai.c, g)

		sc += ws.Groups[w]
		sc += ws.Groups[h]
	}

	return sc
}

func ExplainScore(m *MinimaxAI, out io.Writer, p *tak.Position) {
	tw := tabwriter.NewWriter(out, 4, 8, 1, '\t', 0)
	fmt.Fprintf(tw, "\twhite\tblack\n")
	var scores [2]struct {
		flats    int
		standing int
		caps     int

		stones   int
		captured int
	}

	scores[0].flats = bitboard.Popcount(p.White &^ p.Caps &^ p.Standing)
	scores[1].flats = bitboard.Popcount(p.Black &^ p.Caps &^ p.Standing)
	scores[0].standing = bitboard.Popcount(p.White & p.Standing)
	scores[1].standing = bitboard.Popcount(p.Black & p.Standing)
	scores[0].caps = bitboard.Popcount(p.White & p.Caps)
	scores[1].caps = bitboard.Popcount(p.Black & p.Caps)

	for i, h := range p.Height {
		if h == 0 {
			continue
		}
		s := p.Stacks[i] & ((1 << (h - 1)) - 1)
		bf := bitboard.Popcount(s)
		wf := int(h) - bf - 1
		scores[0].stones += wf
		scores[1].stones += bf

		if p.White&(1<<uint(i)) != 0 {
			scores[0].captured += int(h - 1)
		} else {
			scores[1].captured += int(h - 1)
		}
	}

	fmt.Fprintf(tw, "flats\t%d\t%d\n", scores[0].flats, scores[1].flats)
	fmt.Fprintf(tw, "standing\t%d\t%d\n", scores[0].standing, scores[1].standing)
	fmt.Fprintf(tw, "caps\t%d\t%d\n", scores[0].caps, scores[1].caps)
	fmt.Fprintf(tw, "captured\t%d\t%d\n", scores[0].captured, scores[1].captured)
	fmt.Fprintf(tw, "stones\t%d\t%d\n", scores[0].stones, scores[1].stones)

	analysis := p.Analysis()

	wr := p.White &^ p.Standing
	br := p.Black &^ p.Standing
	wl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.Black, wr) &^ wr)
	bl := bitboard.Popcount(bitboard.Grow(&m.c, ^p.White, br) &^ br)

	fmt.Fprintf(tw, "liberties\t%d\t%d\n", wl, bl)

	for i, g := range analysis.WhiteGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t%dx%x\n", i, w, h)
	}
	for i, g := range analysis.BlackGroups {
		w, h := bitboard.Dimensions(&m.c, g)
		fmt.Fprintf(tw, "g%d\t\t%dx%x\n", i, w, h)
	}
	tw.Flush()
}
<|endoftext|>