<commit_before>// k8s_checker is an application that checks for the following and alerts if necessary:
// * Dirty images checked into K8s config files.
// * Dirty configs running in K8s.
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"time"

	yaml "gopkg.in/yaml.v2"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/gitauth"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
)

const (
	IMAGE_DIRTY_SUFFIX = "-dirty"

	// Metric names.
	DIRTY_COMMITTED_IMAGE_METRIC = "dirty_committed_image_metric"
	DIRTY_CONFIG_METRIC          = "dirty_config_metric"
	LIVENESS_METRIC              = "k8s_checker"
)

var (
	// Flags.
	k8sYamlRepo             = flag.String("k8s_yaml_repo", "https://skia.googlesource.com/skia-public-config", "The repository where K8s yaml files are stored (eg: https://skia.googlesource.com/skia-public-config)")
	kubeConfig              = flag.String("kube_config", "/var/secrets/kube-config/kube_config", "The kube config of the project kubectl will query against.")
	workdir                 = flag.String("workdir", "/tmp/", "Directory to use for scratch work.")
	promPort                = flag.String("prom_port", ":20000", "Metrics service address (e.g., ':20000')")
	serviceAccountKey       = flag.String("service_account_key", "", "Should be set when running in K8s.")
	dirtyConfigChecksPeriod = flag.Duration("dirty_config_checks_period", 2*time.Minute, "How often to check for dirty configs/images in K8s.")
)

type K8sPodsJson struct {
	Items []struct {
		Metadata struct {
			Labels struct {
				App string `json:"app"`
			} `json:"labels"`
		} `json:"metadata"`
		Spec struct {
			Containers []struct {
				Name  string `json:"name"`
				Image string `json:"image"`
			} `json:"containers"`
		} `json:"spec"`
	} `json:"items"`
}

// getLiveAppContainersToImages returns a map of app names to their containers to the images running on them.
func getLiveAppContainersToImages(ctx context.Context) (map[string]map[string]string, error) {
	// Get JSON output of pods running in K8s.
	getPodsCommand := fmt.Sprintf("kubectl get pods --kubeconfig=%s -o json --field-selector=status.phase=Running", *kubeConfig)
	output, err := exec.RunSimple(ctx, getPodsCommand)
	if err != nil {
		return nil, fmt.Errorf("Error when running \"%s\": %s", getPodsCommand, err)
	}

	liveAppContainersToImages := map[string]map[string]string{}
	var podsJson K8sPodsJson
	if err := json.Unmarshal([]byte(output), &podsJson); err != nil {
		return nil, fmt.Errorf("Error when unmarshalling JSON: %s", err)
	}
	for _, item := range podsJson.Items {
		app := item.Metadata.Labels.App
		liveAppContainersToImages[app] = map[string]string{}
		for _, container := range item.Spec.Containers {
			liveAppContainersToImages[app][container.Name] = container.Image
		}
	}
	return liveAppContainersToImages, nil
}

type K8sConfig struct {
	Spec struct {
		Template struct {
			Metadata struct {
				Labels struct {
					App string `yaml:"app"`
				} `yaml:"labels"`
			} `yaml:"metadata"`
			TemplateSpec struct {
				Containers []struct {
					Name  string `yaml:"name"`
					Image string `yaml:"image"`
				} `yaml:"containers"`
			} `yaml:"spec"`
		} `yaml:"template"`
	} `yaml:"spec"`
}

// checkForDirtyConfigs checks for:
// * Dirty images checked into K8s config files.
// * Dirty configs running in K8s.
// It takes in a map of oldMetrics, any metrics from that map that are not encountered during this
// invocation of the function are deleted. This is done to handle the case when metric tags
// change. Eg: liveImage in dirtyConfigMetricTags.
// It returns a map of newMetrics, which are all the metrics that were used during this
// invocation of the function.
func checkForDirtyConfigs(ctx context.Context, oldMetrics map[metrics2.Int64Metric]struct{}) (map[metrics2.Int64Metric]struct{}, error) {
	sklog.Info("\n\n---------- New round of checking k8 dirty configs ----------\n\n")

	// Get mapping from live apps to their containers and images.
	liveAppContainerToImages, err := getLiveAppContainersToImages(ctx)
	if err != nil {
		return nil, fmt.Errorf("Could not get live pods from kubectl: %s", err)
	}

	// Checkout the K8s config repo.
	// Use gitiles if this ends up giving us any problems.
	g, err := git.NewCheckout(ctx, *k8sYamlRepo, *workdir)
	if err != nil {
		return nil, fmt.Errorf("Error when checking out %s: %s", *k8sYamlRepo, err)
	}
	if err := g.Update(ctx); err != nil {
		return nil, fmt.Errorf("Error when updating %s: %s", *k8sYamlRepo, err)
	}

	files, err := ioutil.ReadDir(g.Dir())
	if err != nil {
		return nil, fmt.Errorf("Error when reading from %s: %s", g.Dir(), err)
	}

	newMetrics := map[metrics2.Int64Metric]struct{}{}
	for _, f := range files {
		if filepath.Ext(f.Name()) != ".yaml" {
			// Only interested in YAML configs.
			continue
		}
		b, err := ioutil.ReadFile(filepath.Join(g.Dir(), f.Name()))
		if err != nil {
			sklog.Fatal(err)
		}

		// There can be multiple YAML documents within a single YAML file.
		yamlDocs := strings.Split(string(b), "---")
		for _, yamlDoc := range yamlDocs {
			var config K8sConfig
			if err := yaml.Unmarshal([]byte(yamlDoc), &config); err != nil {
				sklog.Fatal(err)
			}
			app := config.Spec.Template.Metadata.Labels.App
			if app == "" {
				// This YAML config does not have an app. Continue.
				continue
			}
			for _, c := range config.Spec.Template.TemplateSpec.Containers {
				container := c.Name
				committedImage := c.Image

				// Check if the image in the config is dirty.
				dirtyCommittedMetricTags := map[string]string{
					"yaml":           f.Name(),
					"repo":           *k8sYamlRepo,
					"committedImage": committedImage,
				}
				dirtyCommittedMetric := metrics2.GetInt64Metric(DIRTY_COMMITTED_IMAGE_METRIC, dirtyCommittedMetricTags)
				newMetrics[dirtyCommittedMetric] = struct{}{}
				if strings.HasSuffix(committedImage, IMAGE_DIRTY_SUFFIX) {
					sklog.Infof("%s has a dirty committed image: %s\n\n", f.Name(), committedImage)
					dirtyCommittedMetric.Update(1)
				} else {
					dirtyCommittedMetric.Update(0)
				}

				// Check if the image running in k8s matches the checked in image.
				if liveContainersToImages, ok := liveAppContainerToImages[app]; ok {
					if liveImage, ok := liveContainersToImages[container]; ok {
						dirtyConfigMetricTags := map[string]string{
							"app":            app,
							"container":      container,
							"yaml":           f.Name(),
							"repo":           *k8sYamlRepo,
							"committedImage": committedImage,
							"liveImage":      liveImage,
						}
						dirtyConfigMetric := metrics2.GetInt64Metric(DIRTY_CONFIG_METRIC, dirtyConfigMetricTags)
						newMetrics[dirtyConfigMetric] = struct{}{}
						if liveImage != committedImage {
							dirtyConfigMetric.Update(1)
							sklog.Infof("For app %s and container %s the running image differs from the image in config: %s != %s\n\n", app, container, liveImage, committedImage)
						} else {
							dirtyConfigMetric.Update(0)
						}
					} else {
						sklog.Infof("There is no running container %s for the config file %s\n\n", container, f.Name())
					}
				} else {
					sklog.Infof("There is no running app %s for the config file %s\n\n", app, f.Name())
				}
			}
		}
	}

	// Delete unused old metrics.
	for m := range oldMetrics {
		if _, ok := newMetrics[m]; !ok {
			if err := m.Delete(); err != nil {
				sklog.Errorf("Failed to delete metric: %s", err)
				// Add the metric to newMetrics so that we'll
				// have the chance to delete it again on the
				// next cycle.
				newMetrics[m] = struct{}{}
			}
		}
	}
	return newMetrics, nil
}

func main() {
	common.InitWithMust("k8s_checker", common.PrometheusOpt(promPort))
	defer sklog.Flush()
	ctx := context.Background()

	if *serviceAccountKey != "" {
		activationCmd := fmt.Sprintf("gcloud auth activate-service-account --key-file %s", *serviceAccountKey)
		if _, err := exec.RunSimple(ctx, activationCmd); err != nil {
			sklog.Fatal(err)
		}

		// Use the gitcookie created by gitauth package.
		ts, err := auth.NewDefaultTokenSource(false, auth.SCOPE_USERINFO_EMAIL, auth.SCOPE_GERRIT)
		if err != nil {
			sklog.Fatal(err)
		}
		gitcookiesPath := filepath.Join(*workdir, ".gitcookies")
		if _, err := gitauth.New(ts, gitcookiesPath, true, ""); err != nil {
			sklog.Fatalf("Failed to create git cookie updater: %s", err)
		}
	}

	liveness := metrics2.NewLiveness(LIVENESS_METRIC)
	oldMetrics := map[metrics2.Int64Metric]struct{}{}
	for range time.Tick(*dirtyConfigChecksPeriod) {
		newMetrics, err := checkForDirtyConfigs(ctx, oldMetrics)
		if err != nil {
			sklog.Errorf("Error when checking for dirty configs: %s", err)
		} else {
			liveness.Reset()
			oldMetrics = newMetrics
		}
	}
}
<commit_msg>[k8s checker] Add metrics for k8s pod status<commit_after>// k8s_checker is an application that checks for the following and alerts if necessary:
// * Dirty images checked into K8s config files.
// * Dirty configs running in K8s.
package main

import (
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"path/filepath"
	"strings"
	"time"

	yaml "gopkg.in/yaml.v2"

	"go.skia.org/infra/go/auth"
	"go.skia.org/infra/go/common"
	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/gitauth"
	"go.skia.org/infra/go/metrics2"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	IMAGE_DIRTY_SUFFIX = "-dirty"

	// Metric names.
	DIRTY_COMMITTED_IMAGE_METRIC  = "dirty_committed_image_metric"
	DIRTY_CONFIG_METRIC           = "dirty_config_metric"
	LIVENESS_DIRTY_CONFIGS_METRIC = "k8s_checker"
	LIVENESS_POD_STATUS_METRIC    = "k8s_checker_pod_status"
	POD_STATUS_METRIC             = "k8s_pod_status"

	// Possible values for the Phase of a pod. More detail in the docs:
	// https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/#pod-phase
	POD_PHASE_PENDING   = "Pending"
	POD_PHASE_RUNNING   = "Running"
	POD_PHASE_SUCCEEDED = "Succeeded"
	POD_PHASE_FAILED    = "Failed"
	POD_PHASE_UNKNOWN   = "Unknown"
)

var (
	// Flags.
	k8sYamlRepo             = flag.String("k8s_yaml_repo", "https://skia.googlesource.com/skia-public-config", "The repository where K8s yaml files are stored (eg: https://skia.googlesource.com/skia-public-config)")
	kubeConfig              = flag.String("kube_config", "/var/secrets/kube-config/kube_config", "The kube config of the project kubectl will query against.")
	workdir                 = flag.String("workdir", "/tmp/", "Directory to use for scratch work.")
	promPort                = flag.String("prom_port", ":20000", "Metrics service address (e.g., ':20000')")
	serviceAccountKey       = flag.String("service_account_key", "", "Should be set when running in K8s.")
	dirtyConfigChecksPeriod = flag.Duration("dirty_config_checks_period", 2*time.Minute, "How often to check for dirty configs/images in K8s.")
	podStatusMetricsPeriod  = flag.Duration("pod_status_metrics_period", time.Minute, "How often to update pod status metrics.")
)

type containerState struct {
	Running *struct {
		StartedAt time.Time `json:"startedAt"`
	} `json:"running"`
	Terminated *struct {
		ContainerID string    `json:"containerID"`
		ExitCode    int       `json:"exitCode"`
		FinishedAt  time.Time `json:"finishedAt"`
		Message     string    `json:"message"`
		Reason      string    `json:"reason"`
		Signal      int       `json:"signal"`
		StartedAt   time.Time `json:"startedAt"`
	} `json:"terminated"`
	Waiting *struct {
		Message string `json:"message"`
		Reason  string `json:"reason"`
	} `json:"waiting"`
}

type K8sPodsJson struct {
	Items []struct {
		Metadata struct {
			Labels struct {
				App string `json:"app"`
			} `json:"labels"`
			Name string `json:"name"`
		} `json:"metadata"`
		Spec struct {
			Containers []struct {
				Name  string `json:"name"`
				Image string `json:"image"`
			} `json:"containers"`
		} `json:"spec"`
		Status struct {
			Conditions []struct {
				LastProbeTime      time.Time `json:"lastProbeTime"`
				LastTransitionTime time.Time `json:"lastTransitionTime"`
				Status             string    `json:"status"`
				Type               string    `json:"type"`
			} `json:"conditions"`
			ContainerStatuses []struct {
				ContainerID  string          `json:"containerID"`
				Image        string          `json:"image"`
				ImageID      string          `json:"imageID"`
				LastState    *containerState `json:"lastState"`
				Name         string          `json:"name"`
				Ready        bool            `json:"ready"`
				RestartCount int             `json:"restartCount"`
				State        *containerState `json:"state"`
			} `json:"containerStatuses"`
			Message   string    `json:"message"`
			Phase     string    `json:"phase"`
			Reason    string    `json:"reason"`
			StartTime time.Time `json:"startTime"`
		} `json:"status"`
	} `json:"items"`
}

// getLiveAppContainersToImages returns a map of app names to their containers to the images running on them.
func getLiveAppContainersToImages(ctx context.Context) (map[string]map[string]string, error) {
	// Get JSON output of pods running in K8s.
	getPodsCommand := fmt.Sprintf("kubectl get pods --kubeconfig=%s -o json --field-selector=status.phase=Running", *kubeConfig)
	output, err := exec.RunSimple(ctx, getPodsCommand)
	if err != nil {
		return nil, fmt.Errorf("Error when running \"%s\": %s", getPodsCommand, err)
	}

	liveAppContainersToImages := map[string]map[string]string{}
	var podsJson K8sPodsJson
	if err := json.Unmarshal([]byte(output), &podsJson); err != nil {
		return nil, fmt.Errorf("Error when unmarshalling JSON: %s", err)
	}
	for _, item := range podsJson.Items {
		app := item.Metadata.Labels.App
		liveAppContainersToImages[app] = map[string]string{}
		for _, container := range item.Spec.Containers {
			liveAppContainersToImages[app][container.Name] = container.Image
		}
	}
	return liveAppContainersToImages, nil
}

type K8sConfig struct {
	Spec struct {
		Template struct {
			Metadata struct {
				Labels struct {
					App string `yaml:"app"`
				} `yaml:"labels"`
			} `yaml:"metadata"`
			TemplateSpec struct {
				Containers []struct {
					Name  string `yaml:"name"`
					Image string `yaml:"image"`
				} `yaml:"containers"`
			} `yaml:"spec"`
		} `yaml:"template"`
	} `yaml:"spec"`
}

// checkForDirtyConfigs checks for:
// * Dirty images checked into K8s config files.
// * Dirty configs running in K8s.
// It takes in a map of oldMetrics, any metrics from that map that are not encountered during this
// invocation of the function are deleted. This is done to handle the case when metric tags
// change. Eg: liveImage in dirtyConfigMetricTags.
// It returns a map of newMetrics, which are all the metrics that were used during this
// invocation of the function.
func checkForDirtyConfigs(ctx context.Context, oldMetrics map[metrics2.Int64Metric]struct{}) (map[metrics2.Int64Metric]struct{}, error) {
	sklog.Info("\n\n---------- New round of checking k8 dirty configs ----------\n\n")

	// Get mapping from live apps to their containers and images.
	liveAppContainerToImages, err := getLiveAppContainersToImages(ctx)
	if err != nil {
		return nil, fmt.Errorf("Could not get live pods from kubectl: %s", err)
	}

	// Checkout the K8s config repo.
	// Use gitiles if this ends up giving us any problems.
	g, err := git.NewCheckout(ctx, *k8sYamlRepo, *workdir)
	if err != nil {
		return nil, fmt.Errorf("Error when checking out %s: %s", *k8sYamlRepo, err)
	}
	if err := g.Update(ctx); err != nil {
		return nil, fmt.Errorf("Error when updating %s: %s", *k8sYamlRepo, err)
	}

	files, err := ioutil.ReadDir(g.Dir())
	if err != nil {
		return nil, fmt.Errorf("Error when reading from %s: %s", g.Dir(), err)
	}

	newMetrics := map[metrics2.Int64Metric]struct{}{}
	for _, f := range files {
		if filepath.Ext(f.Name()) != ".yaml" {
			// Only interested in YAML configs.
			continue
		}
		b, err := ioutil.ReadFile(filepath.Join(g.Dir(), f.Name()))
		if err != nil {
			sklog.Fatal(err)
		}

		// There can be multiple YAML documents within a single YAML file.
		yamlDocs := strings.Split(string(b), "---")
		for _, yamlDoc := range yamlDocs {
			var config K8sConfig
			if err := yaml.Unmarshal([]byte(yamlDoc), &config); err != nil {
				sklog.Fatal(err)
			}
			app := config.Spec.Template.Metadata.Labels.App
			if app == "" {
				// This YAML config does not have an app. Continue.
				continue
			}
			for _, c := range config.Spec.Template.TemplateSpec.Containers {
				container := c.Name
				committedImage := c.Image

				// Check if the image in the config is dirty.
				dirtyCommittedMetricTags := map[string]string{
					"yaml":           f.Name(),
					"repo":           *k8sYamlRepo,
					"committedImage": committedImage,
				}
				dirtyCommittedMetric := metrics2.GetInt64Metric(DIRTY_COMMITTED_IMAGE_METRIC, dirtyCommittedMetricTags)
				newMetrics[dirtyCommittedMetric] = struct{}{}
				if strings.HasSuffix(committedImage, IMAGE_DIRTY_SUFFIX) {
					sklog.Infof("%s has a dirty committed image: %s\n\n", f.Name(), committedImage)
					dirtyCommittedMetric.Update(1)
				} else {
					dirtyCommittedMetric.Update(0)
				}

				// Check if the image running in k8s matches the checked in image.
				if liveContainersToImages, ok := liveAppContainerToImages[app]; ok {
					if liveImage, ok := liveContainersToImages[container]; ok {
						dirtyConfigMetricTags := map[string]string{
							"app":            app,
							"container":      container,
							"yaml":           f.Name(),
							"repo":           *k8sYamlRepo,
							"committedImage": committedImage,
							"liveImage":      liveImage,
						}
						dirtyConfigMetric := metrics2.GetInt64Metric(DIRTY_CONFIG_METRIC, dirtyConfigMetricTags)
						newMetrics[dirtyConfigMetric] = struct{}{}
						if liveImage != committedImage {
							dirtyConfigMetric.Update(1)
							sklog.Infof("For app %s and container %s the running image differs from the image in config: %s != %s\n\n", app, container, liveImage, committedImage)
						} else {
							dirtyConfigMetric.Update(0)
						}
					} else {
						sklog.Infof("There is no running container %s for the config file %s\n\n", container, f.Name())
					}
				} else {
					sklog.Infof("There is no running app %s for the config file %s\n\n", app, f.Name())
				}
			}
		}
	}

	// Delete unused old metrics.
	for m := range oldMetrics {
		if _, ok := newMetrics[m]; !ok {
			if err := m.Delete(); err != nil {
				sklog.Errorf("Failed to delete metric: %s", err)
				// Add the metric to newMetrics so that we'll
				// have the chance to delete it again on the
				// next cycle.
				newMetrics[m] = struct{}{}
			}
		}
	}
	return newMetrics, nil
}

func updatePodStatusMetrics(ctx context.Context, oldMetrics map[metrics2.Int64Metric]struct{}) (map[metrics2.Int64Metric]struct{}, error) {
	now := time.Now()
	// Get JSON output of pods running in K8s.
	getPodsCommand := fmt.Sprintf("kubectl get pods --kubeconfig=%s -o json", *kubeConfig)
	output, err := exec.RunSimple(ctx, getPodsCommand)
	if err != nil {
		return nil, fmt.Errorf("Error when running \"%s\": %s", getPodsCommand, err)
	}
	var podsJson K8sPodsJson
	if err := json.Unmarshal([]byte(output), &podsJson); err != nil {
		return nil, fmt.Errorf("Error when unmarshalling JSON: %s", err)
	}
	newMetrics := make(map[metrics2.Int64Metric]struct{}, len(podsJson.Items))
	for _, item := range podsJson.Items {
		// Attempt to find the time that the pod transitioned into its
		// current state.
		var lastTransitionTime time.Time
		for _, condition := range item.Status.Conditions {
			if condition.LastTransitionTime.After(lastTransitionTime) {
				lastTransitionTime = condition.LastTransitionTime
			}
		}
		if util.TimeIsZero(lastTransitionTime) && (item.Status.Phase == POD_PHASE_FAILED || item.Status.Phase == POD_PHASE_SUCCEEDED) {
			lastTransitionTime = item.Status.StartTime
		}

		if util.TimeIsZero(lastTransitionTime) {
			b, err := json.MarshalIndent(item, "", "  ")
			if err != nil {
				return nil, err
			}
			sklog.Errorf("Could not find transition time for pod:\n%s", string(b))
			lastTransitionTime = now
		}
		duration := int64(now.Sub(lastTransitionTime).Seconds())

		for _, container := range item.Status.ContainerStatuses {
			// Find an appropriate status. The possible values for
			// Phase do not provide as much information as we'd
			// like, so dig into the State object when possible.
			status := item.Status.Phase
			if status == POD_PHASE_RUNNING && container.State.Terminated != nil {
				status = "Terminating"
			}
			if status == POD_PHASE_PENDING && container.State.Waiting != nil && container.State.Waiting.Reason != "" {
				status = container.State.Waiting.Reason
			}

			// Update the metric.
			tags := map[string]string{
				"app":       item.Metadata.Labels.App,
				"container": container.Name,
				"pod":       item.Metadata.Name,
				"repo":      *k8sYamlRepo,
				"status":    status,
			}
			m := metrics2.GetInt64Metric(POD_STATUS_METRIC, tags)
			newMetrics[m] = struct{}{}
			delete(oldMetrics, m)
			m.Update(duration)
			sklog.Debugf("  %s:\t%s for %ds", item.Metadata.Name, status, duration)
		}
	}
	for m := range oldMetrics {
		if err := m.Delete(); err != nil {
			sklog.Errorf("Failed to delete metric: %s", err)
			// Add the metric to newMetrics so that we'll
			// have the chance to delete it again on the
			// next cycle.
			newMetrics[m] = struct{}{}
		}
	}
	return newMetrics, nil
}

func main() {
	common.InitWithMust("k8s_checker", common.PrometheusOpt(promPort))
	defer sklog.Flush()
	ctx := context.Background()

	if *serviceAccountKey != "" {
		activationCmd := fmt.Sprintf("gcloud auth activate-service-account --key-file %s", *serviceAccountKey)
		if _, err := exec.RunSimple(ctx, activationCmd); err != nil {
			sklog.Fatal(err)
		}

		// Use the gitcookie created by gitauth package.
		ts, err := auth.NewDefaultTokenSource(false, auth.SCOPE_USERINFO_EMAIL, auth.SCOPE_GERRIT)
		if err != nil {
			sklog.Fatal(err)
		}
		gitcookiesPath := filepath.Join(*workdir, ".gitcookies")
		if _, err := gitauth.New(ts, gitcookiesPath, true, ""); err != nil {
			sklog.Fatalf("Failed to create git cookie updater: %s", err)
		}
	}

	livenessDirtyConfigs := metrics2.NewLiveness(LIVENESS_DIRTY_CONFIGS_METRIC)
	oldMetricsDirtyConfigs := map[metrics2.Int64Metric]struct{}{}
	go util.RepeatCtx(*dirtyConfigChecksPeriod, ctx, func() {
		newMetrics, err := checkForDirtyConfigs(ctx, oldMetricsDirtyConfigs)
		if err != nil {
			sklog.Errorf("Error when checking for dirty configs: %s", err)
		} else {
			livenessDirtyConfigs.Reset()
			oldMetricsDirtyConfigs = newMetrics
		}
	})

	livenessPodStatus := metrics2.NewLiveness(LIVENESS_POD_STATUS_METRIC)
	oldMetricsPodStatus := map[metrics2.Int64Metric]struct{}{}
	go util.RepeatCtx(*podStatusMetricsPeriod, ctx, func() {
		newMetrics, err := updatePodStatusMetrics(ctx, oldMetricsPodStatus)
		if err != nil {
			sklog.Errorf("Error when checking pod statuses: %s", err)
		} else {
			livenessPodStatus.Reset()
			oldMetricsPodStatus = newMetrics
		}
	})

	select {}
}
<|endoftext|>


<commit_before>/*
Copyright 2014 Workiva, LLC

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 palm

import (
	"log"
	"runtime"
	"sync"
	"sync/atomic"

	"github.com/Workiva/go-datastructures/common"
	"github.com/Workiva/go-datastructures/queue"
)

type operation int

const (
	get operation = iota
	add
	remove
)

const multiThreadAt = 1000 // number of keys before we multithread lookups

type recursiveBuild struct {
	keys   common.Comparators
	nodes  []*node
	parent *node
}

type ptree struct {
	root                    *node
	ary, number, bufferSize uint64
	actions                 *queue.RingBuffer
	cache                   []interface{}
	buffer0                 [8]uint64
	disposed                uint64
	buffer1                 [8]uint64
	running                 uint64
}

func (ptree *ptree) checkAndRun(action action) {
	if ptree.actions.Len() > 0 {
		if action != nil {
			ptree.actions.Put(action)
		}
		if atomic.CompareAndSwapUint64(&ptree.running, 0, 1) {
			var a interface{}
			var err error
			for ptree.actions.Len() > 0 {
				a, err = ptree.actions.Get()
				if err != nil {
					return
				}
				ptree.cache = append(ptree.cache, a)
				if uint64(len(ptree.cache)) >= ptree.bufferSize {
					break
				}
			}

			go ptree.operationRunner(ptree.cache, true)
		}
	} else if action != nil {
		if atomic.CompareAndSwapUint64(&ptree.running, 0, 1) {
			switch action.operation() {
			case get:
				ptree.read(action)
				action.complete()
				ptree.reset()
			case add:
				if len(action.keys()) > multiThreadAt {
					ptree.operationRunner(interfaces{action}, true)
				} else {
					ptree.operationRunner(interfaces{action}, false)
				}
			}
		} else {
			ptree.actions.Put(action)
			ptree.checkAndRun(nil)
		}
	}
}

func (ptree *ptree) init(bufferSize, ary uint64) {
	ptree.bufferSize = bufferSize
	ptree.ary = ary
	ptree.cache = make([]interface{}, 0, bufferSize)
	ptree.root = newNode(true, newKeys(ary), newNodes(ary))
	ptree.actions = queue.NewRingBuffer(ptree.bufferSize)
}

func (ptree *ptree) operationRunner(xns interfaces, threaded bool) {
	var writeOperations map[*node]common.Comparators
	var toComplete actions

	if threaded {
		writeOperations, toComplete = ptree.fetchKeys(xns)
	} else {
		writeOperations, toComplete = ptree.singleThreadedFetchKeys(xns)
	}

	ptree.runAdds(writeOperations)
	for _, a := range toComplete {
		a.complete()
	}

	ptree.reset()
}

func (ptree *ptree) read(action action) {
	for i, k := range action.keys() {
		n := getParent(ptree.root, k)
		if n == nil {
			action.keys()[i] = nil
		} else {
			key, _ := n.keys.withPosition(k)
			if key == nil {
				action.keys()[i] = nil
			} else {
				action.keys()[i] = key
			}
		}
	}
}

func (ptree *ptree) singleThreadedFetchKeys(xns interfaces) (map[*node]common.Comparators, actions) {
	for _, ifc := range xns {
		action := ifc.(action)
		for i, key := range action.keys() {
			n := getParent(ptree.root, key)
			switch action.operation() {
			case add:
				action.addNode(int64(i), n)
			case get:
				if n == nil {
					action.keys()[i] = nil
				} else {
					k, _ := n.keys.withPosition(key)
					if k == nil {
						action.keys()[i] = nil
					} else {
						action.keys()[i] = k
					}
				}
			}
		}
	}

	writeOperations := make(map[*node]common.Comparators, len(xns)/2)
	toComplete := make(actions, 0, len(xns)/2)
	for _, ifc := range xns {
		action := ifc.(action)
		switch action.operation() {
		case add:
			for i, n := range action.nodes() {
				writeOperations[n] = append(writeOperations[n], action.keys()[i])
			}
			toComplete = append(toComplete, action)
		case get:
			action.complete()
		}
	}

	return writeOperations, toComplete
}

func (ptree *ptree) reset() {
	for i := range ptree.cache {
		ptree.cache[i] = nil
	}

	ptree.cache = ptree.cache[:0]
	atomic.StoreUint64(&ptree.running, 0)
	ptree.checkAndRun(nil)
}

func (ptree *ptree) fetchKeys(xns []interface{}) (map[*node]common.Comparators, actions) {
	var forCache struct {
		i      int64
		buffer [8]uint64 // different cache lines
		js     []int64
	}

	for j := 0; j < len(xns); j++ {
		forCache.js = append(forCache.js, -1)
	}
	numCPU := runtime.NumCPU()
	if numCPU > 1 {
		numCPU--
	}
	var wg sync.WaitGroup
	wg.Add(numCPU)

	for k := 0; k < numCPU; k++ {
		go func() {
			for {
				index := atomic.LoadInt64(&forCache.i)
				if index >= int64(len(xns)) {
					break
				}
				action := xns[index].(action)

				j := atomic.AddInt64(&forCache.js[index], 1)
				if j > int64(len(action.keys())) { // someone else is updating i
					continue
				} else if j == int64(len(action.keys())) {
					atomic.StoreInt64(&forCache.i, index+1)
					continue
				}

				n := getParent(ptree.root, action.keys()[j])
				switch action.operation() {
				case add:
					action.addNode(j, n)
				case get:
					if n == nil {
						action.keys()[j] = nil
					} else {
						k, _ := n.keys.withPosition(action.keys()[j])
						if k == nil {
							action.keys()[j] = nil
						} else {
							action.keys()[j] = k
						}
					}
				}
			}
			wg.Done()
		}()
	}
	wg.Wait()

	writeOperations := make(map[*node]common.Comparators, len(xns)/2)
	toComplete := make(actions, 0, len(xns)/2)
	for _, ifc := range xns {
		action := ifc.(action)
		switch action.operation() {
		case add:
			for i, n := range action.nodes() {
				writeOperations[n] = append(writeOperations[n], action.keys()[i])
			}
			toComplete = append(toComplete, action)
		case get:
			action.complete()
		}
	}

	return writeOperations, toComplete
}

func (ptree *ptree) recursiveSplit(n, parent, left *node, nodes *[]*node, keys *common.Comparators) {
	if !n.needsSplit(ptree.ary) {
		return
	}

	length := n.keys.len()
	splitAt := ptree.ary - 1

	for i := splitAt; i < length; i += splitAt {
		offset := length - i
		k, left, right := n.split(offset)
		left.right = right
		*keys = append(*keys, k)
		*nodes = append(*nodes, left, right)
		left.parent = parent
		right.parent = parent
	}
}

func (ptree *ptree) recursiveAdd(layer map[*node][]*recursiveBuild, setRoot bool) {
	if len(layer) == 0 {
		return
	}

	if setRoot && len(layer) > 1 {
		panic(`SHOULD ONLY HAVE ONE ROOT`)
	}

	ifs := make(interfaces, 0, len(layer))
	for _, rbs := range layer {
		ifs = append(ifs, rbs)
	}

	var write sync.Mutex
	layer = make(map[*node][]*recursiveBuild, len(layer))
	dummyRoot := &node{
		keys:  newKeys(ptree.ary),
		nodes: newNodes(ptree.ary),
	}
	executeInterfacesInParallel(ifs, func(ifc interface{}) {
		rbs := ifc.([]*recursiveBuild)

		if len(rbs) == 0 {
			return
		}

		n := rbs[0].parent
		if setRoot {
			ptree.root = n
		}

		parent := n.parent
		if parent == nil {
			parent = dummyRoot
			setRoot = true
		}

		for _, rb := range rbs {
			for i, k := range rb.keys {
				if n.keys.len() == 0 {
					n.keys.insert(k)
					n.nodes.push(rb.nodes[i*2])
					n.nodes.push(rb.nodes[i*2+1])
					continue
				}

				n.keys.insert(k)
				index := n.search(k)
				n.nodes.replaceAt(index, rb.nodes[i*2])
				n.nodes.insertAt(index+1, rb.nodes[i*2+1])
			}
		}

		if n.needsSplit(ptree.ary) {
			keys := make(common.Comparators, 0, n.keys.len())
			nodes := make([]*node, 0, n.nodes.len())
			ptree.recursiveSplit(n, parent, nil, &nodes, &keys)
			write.Lock()
			layer[parent] = append(
				layer[parent], &recursiveBuild{keys: keys, nodes: nodes, parent: parent},
			)
			write.Unlock()
		}
	})

	ptree.recursiveAdd(layer, setRoot)
}

func (ptree *ptree) runAdds(addOperations map[*node]common.Comparators) {
	if len(addOperations) == 0 {
		return
	}

	var needRoot bool
	ifs := make(interfaces, 0, len(addOperations))
	for n := range addOperations {
		if n.parent == nil {
			needRoot = true
		}
		ifs = append(ifs, n)
	}

	var dummyRoot *node
	if needRoot {
		dummyRoot = &node{
			keys:  newKeys(ptree.ary),
			nodes: newNodes(ptree.ary),
		}
	}

	var write sync.Mutex
	nextLayer := make(map[*node][]*recursiveBuild)
	executeInterfacesInParallel(ifs, func(ifc interface{}) {
		n := ifc.(*node)
		keys := addOperations[n]

		if len(keys) == 0 {
			return
		}

		parent := n.parent
		if parent == nil {
			parent = dummyRoot
		}

		for _, key := range keys {
			oldKey := n.keys.insert(key)
			if oldKey == nil {
				atomic.AddUint64(&ptree.number, 1)
			}
		}

		if n.needsSplit(ptree.ary) {
			keys := make(common.Comparators, 0, n.keys.len())
			nodes := make([]*node, 0, n.nodes.len())
			ptree.recursiveSplit(n, parent, nil, &nodes, &keys)
			write.Lock()
			nextLayer[parent] = append(
				nextLayer[parent], &recursiveBuild{keys: keys, nodes: nodes, parent: parent},
			)
			write.Unlock()
		}
	})

	ptree.recursiveAdd(nextLayer, needRoot)
}

// Insert will add the provided keys to the tree.
func (ptree *ptree) Insert(keys ...common.Comparator) {
	ia := newInsertAction(keys)
	ptree.checkAndRun(ia)
	ia.completer.Wait()
}

// Get will retrieve a list of keys from the provided keys.
func (ptree *ptree) Get(keys ...common.Comparator) common.Comparators {
	ga := newGetAction(keys)
	ptree.checkAndRun(ga)
	ga.completer.Wait()
	return ga.result
}

// Len returns the number of items in the tree.
func (ptree *ptree) Len() uint64 {
	return atomic.LoadUint64(&ptree.number)
}

// Dispose will clean up any resources used by this tree.  This
// must be called to prevent a memory leak.
func (ptree *ptree) Dispose() {
	ptree.actions.Dispose()
	atomic.StoreUint64(&ptree.disposed, 1)
}

func (ptree *ptree) print(output *log.Logger) {
	println(`PRINTING TREE`)
	if ptree.root == nil {
		return
	}

	ptree.root.print(output)
}

func newTree(bufferSize, ary uint64) *ptree {
	ptree := &ptree{}
	ptree.init(bufferSize, ary)
	return ptree
}

// New will allocate, initialize, and return a new B-Tree based
// on PALM principles.  This type of tree is suited for in-memory
// indices in a multi-threaded environment.
func New(bufferSize, ary uint64) BTree {
	return newTree(bufferSize, ary)
}
<commit_msg>Fixing up split signature.<commit_after>/*
Copyright 2014 Workiva, LLC

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 palm

import (
	"log"
	"runtime"
	"sync"
	"sync/atomic"

	"github.com/Workiva/go-datastructures/common"
	"github.com/Workiva/go-datastructures/queue"
)

type operation int

const (
	get operation = iota
	add
	remove
)

const multiThreadAt = 1000 // number of keys before we multithread lookups

type recursiveBuild struct {
	keys   common.Comparators
	nodes  []*node
	parent *node
}

type ptree struct {
	root                    *node
	ary, number, bufferSize uint64
	actions                 *queue.RingBuffer
	cache                   []interface{}
	buffer0                 [8]uint64
	disposed                uint64
	buffer1                 [8]uint64
	running                 uint64
}

func (ptree *ptree) checkAndRun(action action) {
	if ptree.actions.Len() > 0 {
		if action != nil {
			ptree.actions.Put(action)
		}
		if atomic.CompareAndSwapUint64(&ptree.running, 0, 1) {
			var a interface{}
			var err error
			for ptree.actions.Len() > 0 {
				a, err = ptree.actions.Get()
				if err != nil {
					return
				}
				ptree.cache = append(ptree.cache, a)
				if uint64(len(ptree.cache)) >= ptree.bufferSize {
					break
				}
			}

			go ptree.operationRunner(ptree.cache, true)
		}
	} else if action != nil {
		if atomic.CompareAndSwapUint64(&ptree.running, 0, 1) {
			switch action.operation() {
			case get:
				ptree.read(action)
				action.complete()
				ptree.reset()
			case add:
				if len(action.keys()) > multiThreadAt {
					ptree.operationRunner(interfaces{action}, true)
				} else {
					ptree.operationRunner(interfaces{action}, false)
				}
			}
		} else {
			ptree.actions.Put(action)
			ptree.checkAndRun(nil)
		}
	}
}

func (ptree *ptree) init(bufferSize, ary uint64) {
	ptree.bufferSize = bufferSize
	ptree.ary = ary
	ptree.cache = make([]interface{}, 0, bufferSize)
	ptree.root = newNode(true, newKeys(ary), newNodes(ary))
	ptree.actions = queue.NewRingBuffer(ptree.bufferSize)
}

func (ptree *ptree) operationRunner(xns interfaces, threaded bool) {
	var writeOperations map[*node]common.Comparators
	var toComplete actions

	if threaded {
		writeOperations, toComplete = ptree.fetchKeys(xns)
	} else {
		writeOperations, toComplete = ptree.singleThreadedFetchKeys(xns)
	}

	ptree.runAdds(writeOperations)
	for _, a := range toComplete {
		a.complete()
	}

	ptree.reset()
}

func (ptree *ptree) read(action action) {
	for i, k := range action.keys() {
		n := getParent(ptree.root, k)
		if n == nil {
			action.keys()[i] = nil
		} else {
			key, _ := n.keys.withPosition(k)
			if key == nil {
				action.keys()[i] = nil
			} else {
				action.keys()[i] = key
			}
		}
	}
}

func (ptree *ptree) singleThreadedFetchKeys(xns interfaces) (map[*node]common.Comparators, actions) {
	for _, ifc := range xns {
		action := ifc.(action)
		for i, key := range action.keys() {
			n := getParent(ptree.root, key)
			switch action.operation() {
			case add:
				action.addNode(int64(i), n)
			case get:
				if n == nil {
					action.keys()[i] = nil
				} else {
					k, _ := n.keys.withPosition(key)
					if k == nil {
						action.keys()[i] = nil
					} else {
						action.keys()[i] = k
					}
				}
			}
		}
	}

	writeOperations := make(map[*node]common.Comparators, len(xns)/2)
	toComplete := make(actions, 0, len(xns)/2)
	for _, ifc := range xns {
		action := ifc.(action)
		switch action.operation() {
		case add:
			for i, n := range action.nodes() {
				writeOperations[n] = append(writeOperations[n], action.keys()[i])
			}
			toComplete = append(toComplete, action)
		case get:
			action.complete()
		}
	}

	return writeOperations, toComplete
}

func (ptree *ptree) reset() {
	for i := range ptree.cache {
		ptree.cache[i] = nil
	}

	ptree.cache = ptree.cache[:0]
	atomic.StoreUint64(&ptree.running, 0)
	ptree.checkAndRun(nil)
}

func (ptree *ptree) fetchKeys(xns []interface{}) (map[*node]common.Comparators, actions) {
	var forCache struct {
		i      int64
		buffer [8]uint64 // different cache lines
		js     []int64
	}

	for j := 0; j < len(xns); j++ {
		forCache.js = append(forCache.js, -1)
	}
	numCPU := runtime.NumCPU()
	if numCPU > 1 {
		numCPU--
	}
	var wg sync.WaitGroup
	wg.Add(numCPU)

	for k := 0; k < numCPU; k++ {
		go func() {
			for {
				index := atomic.LoadInt64(&forCache.i)
				if index >= int64(len(xns)) {
					break
				}
				action := xns[index].(action)

				j := atomic.AddInt64(&forCache.js[index], 1)
				if j > int64(len(action.keys())) { // someone else is updating i
					continue
				} else if j == int64(len(action.keys())) {
					atomic.StoreInt64(&forCache.i, index+1)
					continue
				}

				n := getParent(ptree.root, action.keys()[j])
				switch action.operation() {
				case add:
					action.addNode(j, n)
				case get:
					if n == nil {
						action.keys()[j] = nil
					} else {
						k, _ := n.keys.withPosition(action.keys()[j])
						if k == nil {
							action.keys()[j] = nil
						} else {
							action.keys()[j] = k
						}
					}
				}
			}
			wg.Done()
		}()
	}
	wg.Wait()

	writeOperations := make(map[*node]common.Comparators, len(xns)/2)
	toComplete := make(actions, 0, len(xns)/2)
	for _, ifc := range xns {
		action := ifc.(action)
		switch action.operation() {
		case add:
			for i, n := range action.nodes() {
				writeOperations[n] = append(writeOperations[n], action.keys()[i])
			}
			toComplete = append(toComplete, action)
		case get:
			action.complete()
		}
	}

	return writeOperations, toComplete
}

func (ptree *ptree) splitNode(n, parent *node, nodes *[]*node, keys *common.Comparators) {
	if !n.needsSplit(ptree.ary) {
		return
	}

	length := n.keys.len()
	splitAt := ptree.ary - 1

	for i := splitAt; i < length; i += splitAt {
		offset := length - i
		k, left, right := n.split(offset)
		left.right = right
		*keys = append(*keys, k)
		*nodes = append(*nodes, left, right)
		left.parent = parent
		right.parent = parent
	}
}

func (ptree *ptree) recursiveAdd(layer map[*node][]*recursiveBuild, setRoot bool) {
	if len(layer) == 0 {
		return
	}

	if setRoot && len(layer) > 1 {
		panic(`SHOULD ONLY HAVE ONE ROOT`)
	}

	ifs := make(interfaces, 0, len(layer))
	for _, rbs := range layer {
		ifs = append(ifs, rbs)
	}

	var write sync.Mutex
	layer = make(map[*node][]*recursiveBuild, len(layer))
	dummyRoot := &node{
		keys:  newKeys(ptree.ary),
		nodes: newNodes(ptree.ary),
	}
	executeInterfacesInParallel(ifs, func(ifc interface{}) {
		rbs := ifc.([]*recursiveBuild)

		if len(rbs) == 0 {
			return
		}

		n := rbs[0].parent
		if setRoot {
			ptree.root = n
		}

		parent := n.parent
		if parent == nil {
			parent = dummyRoot
			setRoot = true
		}

		for _, rb := range rbs {
			for i, k := range rb.keys {
				if n.keys.len() == 0 {
					n.keys.insert(k)
					n.nodes.push(rb.nodes[i*2])
					n.nodes.push(rb.nodes[i*2+1])
					continue
				}

				n.keys.insert(k)
				index := n.search(k)
				n.nodes.replaceAt(index, rb.nodes[i*2])
				n.nodes.insertAt(index+1, rb.nodes[i*2+1])
			}
		}

		if n.needsSplit(ptree.ary) {
			keys := make(common.Comparators, 0, n.keys.len())
			nodes := make([]*node, 0, n.nodes.len())
			ptree.splitNode(n, parent, &nodes, &keys)
			write.Lock()
			layer[parent] = append(
				layer[parent], &recursiveBuild{keys: keys, nodes: nodes, parent: parent},
			)
			write.Unlock()
		}
	})

	ptree.recursiveAdd(layer, setRoot)
}

func (ptree *ptree) runAdds(addOperations map[*node]common.Comparators) {
	if len(addOperations) == 0 {
		return
	}

	var needRoot bool
	ifs := make(interfaces, 0, len(addOperations))
	for n := range addOperations {
		if n.parent == nil {
			needRoot = true
		}
		ifs = append(ifs, n)
	}

	var dummyRoot *node
	if needRoot {
		dummyRoot = &node{
			keys:  newKeys(ptree.ary),
			nodes: newNodes(ptree.ary),
		}
	}

	var write sync.Mutex
	nextLayer := make(map[*node][]*recursiveBuild)
	executeInterfacesInParallel(ifs, func(ifc interface{}) {
		n := ifc.(*node)
		keys := addOperations[n]

		if len(keys) == 0 {
			return
		}

		parent := n.parent
		if parent == nil {
			parent = dummyRoot
		}

		for _, key := range keys {
			oldKey := n.keys.insert(key)
			if oldKey == nil {
				atomic.AddUint64(&ptree.number, 1)
			}
		}

		if n.needsSplit(ptree.ary) {
			keys := make(common.Comparators, 0, n.keys.len())
			nodes := make([]*node, 0, n.nodes.len())
			ptree.splitNode(n, parent, &nodes, &keys)
			write.Lock()
			nextLayer[parent] = append(
				nextLayer[parent], &recursiveBuild{keys: keys, nodes: nodes, parent: parent},
			)
			write.Unlock()
		}
	})

	ptree.recursiveAdd(nextLayer, needRoot)
}

// Insert will add the provided keys to the tree.
func (ptree *ptree) Insert(keys ...common.Comparator) {
	ia := newInsertAction(keys)
	ptree.checkAndRun(ia)
	ia.completer.Wait()
}

// Get will retrieve a list of keys from the provided keys.
func (ptree *ptree) Get(keys ...common.Comparator) common.Comparators {
	ga := newGetAction(keys)
	ptree.checkAndRun(ga)
	ga.completer.Wait()
	return ga.result
}

// Len returns the number of items in the tree.
func (ptree *ptree) Len() uint64 {
	return atomic.LoadUint64(&ptree.number)
}

// Dispose will clean up any resources used by this tree.  This
// must be called to prevent a memory leak.
func (ptree *ptree) Dispose() {
	ptree.actions.Dispose()
	atomic.StoreUint64(&ptree.disposed, 1)
}

func (ptree *ptree) print(output *log.Logger) {
	println(`PRINTING TREE`)
	if ptree.root == nil {
		return
	}

	ptree.root.print(output)
}

func newTree(bufferSize, ary uint64) *ptree {
	ptree := &ptree{}
	ptree.init(bufferSize, ary)
	return ptree
}

// New will allocate, initialize, and return a new B-Tree based
// on PALM principles.  This type of tree is suited for in-memory
// indices in a multi-threaded environment.
func New(bufferSize, ary uint64) BTree {
	return newTree(bufferSize, ary)
}
<|endoftext|>


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

package casbin

import (
	"reflect"
)

// GetAllSubjects gets the list of subjects that show up in the current policy.
func (e *Enforcer) GetAllSubjects() []string {
	return e.model.GetValuesForFieldInPolicy("p", "p", 0)
}

// GetAllObjects gets the list of objects that show up in the current policy.
func (e *Enforcer) GetAllObjects() []string {
	return e.model.GetValuesForFieldInPolicy("p", "p", 1)
}

// GetAllActions gets the list of actions that show up in the current policy.
func (e *Enforcer) GetAllActions() []string {
	return e.model.GetValuesForFieldInPolicy("p", "p", 2)
}

// GetAllRoles gets the list of roles that show up in the current policy.
func (e *Enforcer) GetAllRoles() []string {
	return e.model.GetValuesForFieldInPolicy("g", "g", 1)
}

// GetPolicy gets all the authorization rules in the policy.
func (e *Enforcer) GetPolicy() [][]string {
	return e.model.GetPolicy("p", "p")
}

// GetFilteredPolicy gets all the authorization rules in the policy, field filters can be specified.
func (e *Enforcer) GetFilteredPolicy(fieldIndex int, fieldValues ...string) [][]string {
	return e.model.GetFilteredPolicy("p", "p", fieldIndex, fieldValues...)
}

// GetGroupingPolicy gets all the role inheritance rules in the policy.
func (e *Enforcer) GetGroupingPolicy() [][]string {
	return e.model.GetPolicy("g", "g")
}

// HasPolicy determines whether an authorization rule exists.
func (e *Enforcer) HasPolicy(params ...interface{}) bool {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		return e.model.HasPolicy("p", "p", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		return e.model.HasPolicy("p", "p", policy)
	}
}

// AddPolicy adds an authorization rule to the current policy.
// If you try to add an existing policy, the call fails and returns false.
func (e *Enforcer) AddPolicy(params ...interface{}) bool {
	res := false
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		res = e.model.AddPolicy("p", "p", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		res = e.model.AddPolicy("p", "p", policy)
	}

	return res
}

// RemovePolicy removes an authorization rule from the current policy.
func (e *Enforcer) RemovePolicy(params ...interface{}) {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		e.model.RemovePolicy("p", "p", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		e.model.RemovePolicy("p", "p", policy)
	}
}

// RemoveFilteredPolicy removes an authorization rule from the current policy, field filters can be specified.
func (e *Enforcer) RemoveFilteredPolicy(fieldIndex int, fieldValues ...string) {
	e.model.RemoveFilteredPolicy("p", "p", fieldIndex, fieldValues...)
}

// HasGroupingPolicy determines whether a role inheritance rule exists.
func (e *Enforcer) HasGroupingPolicy(params ...interface{}) bool {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		return e.model.HasPolicy("g", "g", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		return e.model.HasPolicy("g", "g", policy)
	}
}

// AddGroupingPolicy adds a role inheritance rule to the current policy.
// If you try to add an existing policy, the call fails and returns false.
func (e *Enforcer) AddGroupingPolicy(params ...interface{}) bool {
	res := false
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		res = e.model.AddPolicy("g", "g", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		res = e.model.AddPolicy("g", "g", policy)
	}

	e.model.BuildRoleLinks()
	return res
}

// RemoveGroupingPolicy removes a role inheritance rule from the current policy.
func (e *Enforcer) RemoveGroupingPolicy(params ...interface{}) {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		e.model.RemovePolicy("g", "g", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		e.model.RemovePolicy("g", "g", policy)
	}

	e.model.BuildRoleLinks()
}

// RemoveFilteredGroupingPolicy removes a role inheritance rule from the current policy, field filters can be specified.
func (e *Enforcer) RemoveFilteredGroupingPolicy(fieldIndex int, fieldValues ...string) {
	e.model.RemoveFilteredPolicy("g", "g", fieldIndex, fieldValues...)
	e.model.BuildRoleLinks()
}

// AddSubjectAttributeFunction adds the function that gets attributes for a subject in ABAC.
func (e *Enforcer) AddSubjectAttributeFunction(function func(args ...interface{}) (interface{}, error)) {
	e.fm.AddFunction("subAttr", function)
}

// AddObjectAttributeFunction adds the function that gets attributes for a object in ABAC.
func (e *Enforcer) AddObjectAttributeFunction(function func(args ...interface{}) (interface{}, error)) {
	e.fm.AddFunction("objAttr", function)
}

// AddActionAttributeFunction adds the function that gets attributes for a object in ABAC.
func (e *Enforcer) AddActionAttributeFunction(function func(args ...interface{}) (interface{}, error)) {
	e.fm.AddFunction("actAttr", function)
}

// AddFunction adds a customized function.
func (e *Enforcer) AddFunction(name string, function func(args ...interface{}) (interface{}, error)) {
	e.fm.AddFunction(name, function)
}
<commit_msg>Remove AddSubjectAttributeFunction() and other two legacy ABAC functions in mgmt API. Please use the new ABAC feature.<commit_after>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package casbin

import (
	"reflect"
)

// GetAllSubjects gets the list of subjects that show up in the current policy.
func (e *Enforcer) GetAllSubjects() []string {
	return e.model.GetValuesForFieldInPolicy("p", "p", 0)
}

// GetAllObjects gets the list of objects that show up in the current policy.
func (e *Enforcer) GetAllObjects() []string {
	return e.model.GetValuesForFieldInPolicy("p", "p", 1)
}

// GetAllActions gets the list of actions that show up in the current policy.
func (e *Enforcer) GetAllActions() []string {
	return e.model.GetValuesForFieldInPolicy("p", "p", 2)
}

// GetAllRoles gets the list of roles that show up in the current policy.
func (e *Enforcer) GetAllRoles() []string {
	return e.model.GetValuesForFieldInPolicy("g", "g", 1)
}

// GetPolicy gets all the authorization rules in the policy.
func (e *Enforcer) GetPolicy() [][]string {
	return e.model.GetPolicy("p", "p")
}

// GetFilteredPolicy gets all the authorization rules in the policy, field filters can be specified.
func (e *Enforcer) GetFilteredPolicy(fieldIndex int, fieldValues ...string) [][]string {
	return e.model.GetFilteredPolicy("p", "p", fieldIndex, fieldValues...)
}

// GetGroupingPolicy gets all the role inheritance rules in the policy.
func (e *Enforcer) GetGroupingPolicy() [][]string {
	return e.model.GetPolicy("g", "g")
}

// HasPolicy determines whether an authorization rule exists.
func (e *Enforcer) HasPolicy(params ...interface{}) bool {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		return e.model.HasPolicy("p", "p", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		return e.model.HasPolicy("p", "p", policy)
	}
}

// AddPolicy adds an authorization rule to the current policy.
// If you try to add an existing policy, the call fails and returns false.
func (e *Enforcer) AddPolicy(params ...interface{}) bool {
	res := false
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		res = e.model.AddPolicy("p", "p", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		res = e.model.AddPolicy("p", "p", policy)
	}

	return res
}

// RemovePolicy removes an authorization rule from the current policy.
func (e *Enforcer) RemovePolicy(params ...interface{}) {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		e.model.RemovePolicy("p", "p", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		e.model.RemovePolicy("p", "p", policy)
	}
}

// RemoveFilteredPolicy removes an authorization rule from the current policy, field filters can be specified.
func (e *Enforcer) RemoveFilteredPolicy(fieldIndex int, fieldValues ...string) {
	e.model.RemoveFilteredPolicy("p", "p", fieldIndex, fieldValues...)
}

// HasGroupingPolicy determines whether a role inheritance rule exists.
func (e *Enforcer) HasGroupingPolicy(params ...interface{}) bool {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		return e.model.HasPolicy("g", "g", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		return e.model.HasPolicy("g", "g", policy)
	}
}

// AddGroupingPolicy adds a role inheritance rule to the current policy.
// If you try to add an existing policy, the call fails and returns false.
func (e *Enforcer) AddGroupingPolicy(params ...interface{}) bool {
	res := false
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		res = e.model.AddPolicy("g", "g", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		res = e.model.AddPolicy("g", "g", policy)
	}

	e.model.BuildRoleLinks()
	return res
}

// RemoveGroupingPolicy removes a role inheritance rule from the current policy.
func (e *Enforcer) RemoveGroupingPolicy(params ...interface{}) {
	if len(params) == 1 && reflect.TypeOf(params[0]).Kind() == reflect.Slice {
		e.model.RemovePolicy("g", "g", params[0].([]string))
	} else {
		policy := make([]string, 0)
		for _, param := range params {
			policy = append(policy, param.(string))
		}

		e.model.RemovePolicy("g", "g", policy)
	}

	e.model.BuildRoleLinks()
}

// RemoveFilteredGroupingPolicy removes a role inheritance rule from the current policy, field filters can be specified.
func (e *Enforcer) RemoveFilteredGroupingPolicy(fieldIndex int, fieldValues ...string) {
	e.model.RemoveFilteredPolicy("g", "g", fieldIndex, fieldValues...)
	e.model.BuildRoleLinks()
}

// AddFunction adds a customized function.
func (e *Enforcer) AddFunction(name string, function func(args ...interface{}) (interface{}, error)) {
	e.fm.AddFunction(name, function)
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/codegangsta/cli"
	"log"
)

func ListAction(c *cli.Context, root string) {
	log.Println("LIST")
}

func CleanAction(c *cli.Context, root string) {
	log.Println("CLEAN")
}
<commit_msg>Remove a useless file<commit_after><|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"os/exec"
	"runtime"
	"sort"
)

var cmdActive = &Command{
	Usage: "active -a [account]",
	Short: "Show or set the active force.com account",
	Long: `
Set the active force.com account

Examples:

  force active
  force active -a user@example.org
`,
}
var (
	tojson  bool
	account string
)

func init() {
	cmdActive.Flag.BoolVar(&tojson, "j", false, "output to json")
	cmdActive.Flag.BoolVar(&tojson, "json", false, "output to json")
	cmdActive.Flag.StringVar(&account, "a", "", "output to json")
	cmdActive.Flag.StringVar(&account, "account", "", "output to json")
	cmdActive.Run = runActive
}

func runActive(cmd *Command, args []string) {
	if account == "" {
		account, _ := Config.Load("current", "account")
		data, _ := Config.Load("accounts", account)
		var creds ForceCredentials
		json.Unmarshal([]byte(data), &creds)
		if tojson {
			fmt.Printf(fmt.Sprintf("{ \"login\": \"%s\", \"instanceUrl\": \"%s\", \"namespace\":\"%s\" }", account, creds.InstanceUrl, creds.Namespace))
		} else {
			fmt.Println(fmt.Sprintf("%s - %s - ns:%s", account, creds.InstanceUrl, creds.Namespace))
		}
	} else {
		//account := args[0]
		accounts, _ := Config.List("accounts")
		i := sort.SearchStrings(accounts, account)
		if i < len(accounts) && accounts[i] == account {
			if runtime.GOOS == "windows" {
				cmd := exec.Command("title", account)
				cmd.Run()
			} else {
				title := fmt.Sprintf("\033];%s\007", account)
				fmt.Printf(title)
			}
			fmt.Printf("%s now active", account)
			Config.Save("current", "account", account)
		} else {
			ErrorAndExit(fmt.Sprintf("no such account %s\n", account))
		}
	}
}
<commit_msg>Update "force active -a [username]" CLI output<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"os/exec"
	"runtime"
	"sort"
)

var cmdActive = &Command{
	Usage: "active -a [account]",
	Short: "Show or set the active force.com account",
	Long: `
Set the active force.com account

Examples:

  force active
  force active -a user@example.org
`,
}
var (
	tojson  bool
	account string
)

func init() {
	cmdActive.Flag.BoolVar(&tojson, "j", false, "output to json")
	cmdActive.Flag.BoolVar(&tojson, "json", false, "output to json")
	cmdActive.Flag.StringVar(&account, "a", "", "output to json")
	cmdActive.Flag.StringVar(&account, "account", "", "output to json")
	cmdActive.Run = runActive
}

func runActive(cmd *Command, args []string) {
	if account == "" {
		account, _ := Config.Load("current", "account")
		data, _ := Config.Load("accounts", account)
		var creds ForceCredentials
		json.Unmarshal([]byte(data), &creds)
		if tojson {
			fmt.Printf(fmt.Sprintf("{ \"login\": \"%s\", \"instanceUrl\": \"%s\", \"namespace\":\"%s\" }", account, creds.InstanceUrl, creds.Namespace))
		} else {
			fmt.Println(fmt.Sprintf("%s - %s - ns:%s", account, creds.InstanceUrl, creds.Namespace))
		}
	} else {
		//account := args[0]
		accounts, _ := Config.List("accounts")
		i := sort.SearchStrings(accounts, account)
		if i < len(accounts) && accounts[i] == account {
			if runtime.GOOS == "windows" {
				cmd := exec.Command("title", account)
				cmd.Run()
			} else {
				title := fmt.Sprintf("\033];%s\007", account)
				fmt.Printf(title)
			}
			fmt.Printf("%s now active\n", account)
			Config.Save("current", "account", account)
		} else {
			ErrorAndExit(fmt.Sprintf("no such account %s\n", account))
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2015 The etcd 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 main

import (
	"fmt"
	"time"
)

type tester struct {
	failures         []failure
	cluster          *cluster
	limit            int
	consistencyCheck bool

	status          Status
	currentRevision int64
}

// compactQPS is rough number of compact requests per second.
// Previous tests showed etcd can compact about 60,000 entries per second.
const compactQPS = 50000

func (tt *tester) runLoop() {
	tt.status.Since = time.Now()
	tt.status.RoundLimit = tt.limit
	tt.status.cluster = tt.cluster
	for _, f := range tt.failures {
		tt.status.Failures = append(tt.status.Failures, f.Desc())
	}

	var (
		round          int
		prevCompactRev int64
	)
	for {
		tt.status.setRound(round)
		tt.status.setCase(-1) // -1 so that logPrefix doesn't print out 'case'
		roundTotalCounter.Inc()

		var failed bool
		for j, f := range tt.failures {
			caseTotalCounter.WithLabelValues(f.Desc()).Inc()
			tt.status.setCase(j)

			if err := tt.cluster.WaitHealth(); err != nil {
				plog.Printf("%s wait full health error: %v", tt.logPrefix(), err)
				if err := tt.cleanup(); err != nil {
					return
				}
				failed = true
				break
			}

			plog.Printf("%s injecting failure %q", tt.logPrefix(), f.Desc())
			if err := f.Inject(tt.cluster, round); err != nil {
				plog.Printf("%s injection error: %v", tt.logPrefix(), err)
				if err := tt.cleanup(); err != nil {
					return
				}
				failed = true
				break
			}
			plog.Printf("%s injected failure", tt.logPrefix())

			plog.Printf("%s recovering failure %q", tt.logPrefix(), f.Desc())
			if err := f.Recover(tt.cluster, round); err != nil {
				plog.Printf("%s recovery error: %v", tt.logPrefix(), err)
				if err := tt.cleanup(); err != nil {
					return
				}
				failed = true
				break
			}
			plog.Printf("%s recovered failure", tt.logPrefix())

			if tt.cluster.v2Only {
				plog.Printf("%s succeed!", tt.logPrefix())
				continue
			}

			if !tt.consistencyCheck {
				if err := tt.updateRevision(); err != nil {
					plog.Warningf("%s functional-tester returning with tt.updateRevision error (%v)", tt.logPrefix(), err)
					return
				}
				continue
			}

			var err error
			failed, err = tt.checkConsistency()
			if err != nil {
				plog.Warningf("%s functional-tester returning with tt.checkConsistency error (%v)", tt.logPrefix(), err)
				return
			}
			if failed {
				break
			}
			plog.Printf("%s succeed!", tt.logPrefix())
		}

		// -1 so that logPrefix doesn't print out 'case'
		tt.status.setCase(-1)

		if failed {
			continue
		}

		revToCompact := max(0, tt.currentRevision-10000)
		compactN := revToCompact - prevCompactRev
		timeout := 10 * time.Second
		if prevCompactRev != 0 && compactN > 0 {
			timeout += time.Duration(compactN/compactQPS) * time.Second
		}
		prevCompactRev = revToCompact

		plog.Printf("%s compacting %d entries (timeout %v)", tt.logPrefix(), compactN, timeout)
		if err := tt.compact(revToCompact, timeout); err != nil {
			plog.Warningf("%s functional-tester returning with error (%v)", tt.logPrefix(), err)
			return
		}
		if round > 0 && round%500 == 0 { // every 500 rounds
			if err := tt.defrag(); err != nil {
				plog.Warningf("%s functional-tester returning with error (%v)", tt.logPrefix(), err)
				return
			}
		}

		round++
		if round == tt.limit {
			plog.Printf("%s functional-tester is finished", tt.logPrefix())
			break
		}
	}
}

func (tt *tester) updateRevision() error {
	revs, _, err := tt.cluster.getRevisionHash()
	for _, rev := range revs {
		tt.currentRevision = rev
		break // just need get one of the current revisions
	}
	return err
}

func (tt *tester) checkConsistency() (failed bool, err error) {
	tt.cancelStressers()
	defer tt.startStressers()

	plog.Printf("%s updating current revisions...", tt.logPrefix())
	var (
		revs   map[string]int64
		hashes map[string]int64
		rerr   error
		ok     bool
	)
	for i := 0; i < 7; i++ {
		time.Sleep(time.Second)

		revs, hashes, rerr = tt.cluster.getRevisionHash()
		if rerr != nil {
			plog.Printf("%s #%d failed to get current revisions (%v)", tt.logPrefix(), i, rerr)
			continue
		}
		if tt.currentRevision, ok = getSameValue(revs); ok {
			break
		}

		plog.Printf("%s #%d inconsistent current revisions %+v", tt.logPrefix(), i, revs)
	}
	plog.Printf("%s updated current revisions with %d", tt.logPrefix(), tt.currentRevision)

	if !ok || rerr != nil {
		plog.Printf("%s checking current revisions failed [revisions: %v]", tt.logPrefix(), revs)
		failed = true
		err = tt.cleanup()
		return
	}
	plog.Printf("%s all members are consistent with current revisions [revisions: %v]", tt.logPrefix(), revs)

	plog.Printf("%s checking current storage hashes...", tt.logPrefix())
	if _, ok = getSameValue(hashes); !ok {
		plog.Printf("%s checking current storage hashes failed [hashes: %v]", tt.logPrefix(), hashes)
		failed = true
		err = tt.cleanup()
		return
	}
	plog.Printf("%s all members are consistent with storage hashes", tt.logPrefix())
	return
}

func (tt *tester) compact(rev int64, timeout time.Duration) error {
	plog.Printf("%s compacting storage (current revision %d, compact revision %d)", tt.logPrefix(), tt.currentRevision, rev)
	if err := tt.cluster.compactKV(rev, timeout); err != nil {
		if cerr := tt.cleanup(); cerr != nil {
			return fmt.Errorf("%s, %s", err, cerr)
		}
		return err
	}
	plog.Printf("%s compacted storage (compact revision %d)", tt.logPrefix(), rev)

	plog.Printf("%s checking compaction (compact revision %d)", tt.logPrefix(), rev)
	if err := tt.cluster.checkCompact(rev); err != nil {
		plog.Printf("%s checkCompact error (%v)", tt.logPrefix(), err)
		if cerr := tt.cleanup(); cerr != nil {
			return fmt.Errorf("%s, %s", err, cerr)
		}
		return err
	}

	plog.Printf("%s confirmed compaction (compact revision %d)", tt.logPrefix(), rev)
	return nil
}

func (tt *tester) defrag() error {
	plog.Printf("%s defragmenting...", tt.logPrefix())
	if err := tt.cluster.defrag(); err != nil {
		plog.Printf("%s defrag error (%v)", tt.logPrefix(), err)
		if cerr := tt.cleanup(); cerr != nil {
			return fmt.Errorf("%s, %s", err, cerr)
		}
		return err
	}

	plog.Printf("%s defragmented...", tt.logPrefix())
	return nil
}

func (tt *tester) logPrefix() string {
	var (
		rd     = tt.status.getRound()
		cs     = tt.status.getCase()
		prefix = fmt.Sprintf("[round#%d case#%d]", rd, cs)
	)
	if cs == -1 {
		prefix = fmt.Sprintf("[round#%d]", rd)
	}
	return prefix
}

func (tt *tester) cleanup() error {
	roundFailedTotalCounter.Inc()
	desc := "compact/defrag"
	if tt.status.Case != -1 {
		desc = tt.failures[tt.status.Case].Desc()
	}
	caseFailedTotalCounter.WithLabelValues(desc).Inc()

	plog.Printf("%s cleaning up...", tt.logPrefix())
	if err := tt.cluster.Cleanup(); err != nil {
		plog.Printf("%s cleanup error: %v", tt.logPrefix(), err)
		return err
	}
	return tt.cluster.Bootstrap()
}

func (tt *tester) cancelStressers() {
	plog.Printf("%s canceling the stressers...", tt.logPrefix())
	for _, s := range tt.cluster.Stressers {
		s.Cancel()
	}
	plog.Printf("%s canceled stressers", tt.logPrefix())
}

func (tt *tester) startStressers() {
	plog.Printf("%s starting the stressers...", tt.logPrefix())
	for _, s := range tt.cluster.Stressers {
		go s.Stress()
	}
	plog.Printf("%s started stressers", tt.logPrefix())
}
<commit_msg>etcd-tester: do not exit for compact timeout<commit_after>// Copyright 2015 The etcd 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 main

import (
	"fmt"
	"time"
)

type tester struct {
	failures         []failure
	cluster          *cluster
	limit            int
	consistencyCheck bool

	status          Status
	currentRevision int64
}

// compactQPS is rough number of compact requests per second.
// Previous tests showed etcd can compact about 60,000 entries per second.
const compactQPS = 50000

func (tt *tester) runLoop() {
	tt.status.Since = time.Now()
	tt.status.RoundLimit = tt.limit
	tt.status.cluster = tt.cluster
	for _, f := range tt.failures {
		tt.status.Failures = append(tt.status.Failures, f.Desc())
	}

	var (
		round          int
		prevCompactRev int64
	)
	for {
		tt.status.setRound(round)
		tt.status.setCase(-1) // -1 so that logPrefix doesn't print out 'case'
		roundTotalCounter.Inc()

		var failed bool
		for j, f := range tt.failures {
			caseTotalCounter.WithLabelValues(f.Desc()).Inc()
			tt.status.setCase(j)

			if err := tt.cluster.WaitHealth(); err != nil {
				plog.Printf("%s wait full health error: %v", tt.logPrefix(), err)
				if err := tt.cleanup(); err != nil {
					return
				}
				failed = true
				break
			}

			plog.Printf("%s injecting failure %q", tt.logPrefix(), f.Desc())
			if err := f.Inject(tt.cluster, round); err != nil {
				plog.Printf("%s injection error: %v", tt.logPrefix(), err)
				if err := tt.cleanup(); err != nil {
					return
				}
				failed = true
				break
			}
			plog.Printf("%s injected failure", tt.logPrefix())

			plog.Printf("%s recovering failure %q", tt.logPrefix(), f.Desc())
			if err := f.Recover(tt.cluster, round); err != nil {
				plog.Printf("%s recovery error: %v", tt.logPrefix(), err)
				if err := tt.cleanup(); err != nil {
					return
				}
				failed = true
				break
			}
			plog.Printf("%s recovered failure", tt.logPrefix())

			if tt.cluster.v2Only {
				plog.Printf("%s succeed!", tt.logPrefix())
				continue
			}

			if !tt.consistencyCheck {
				if err := tt.updateRevision(); err != nil {
					plog.Warningf("%s functional-tester returning with tt.updateRevision error (%v)", tt.logPrefix(), err)
					return
				}
				continue
			}

			var err error
			failed, err = tt.checkConsistency()
			if err != nil {
				plog.Warningf("%s functional-tester returning with tt.checkConsistency error (%v)", tt.logPrefix(), err)
				return
			}
			if failed {
				break
			}
			plog.Printf("%s succeed!", tt.logPrefix())
		}

		// -1 so that logPrefix doesn't print out 'case'
		tt.status.setCase(-1)

		if failed {
			continue
		}

		revToCompact := max(0, tt.currentRevision-10000)
		compactN := revToCompact - prevCompactRev
		timeout := 10 * time.Second
		if prevCompactRev != 0 && compactN > 0 {
			timeout += time.Duration(compactN/compactQPS) * time.Second
		}
		prevCompactRev = revToCompact

		plog.Printf("%s compacting %d entries (timeout %v)", tt.logPrefix(), compactN, timeout)
		if err := tt.compact(revToCompact, timeout); err != nil {
			plog.Warningf("%s functional-tester compact got error (%v)", tt.logPrefix(), err)
			if err := tt.cleanup(); err != nil {
				return
			}
		}
		if round > 0 && round%500 == 0 { // every 500 rounds
			if err := tt.defrag(); err != nil {
				plog.Warningf("%s functional-tester returning with error (%v)", tt.logPrefix(), err)
				return
			}
		}

		round++
		if round == tt.limit {
			plog.Printf("%s functional-tester is finished", tt.logPrefix())
			break
		}
	}
}

func (tt *tester) updateRevision() error {
	revs, _, err := tt.cluster.getRevisionHash()
	for _, rev := range revs {
		tt.currentRevision = rev
		break // just need get one of the current revisions
	}
	return err
}

func (tt *tester) checkConsistency() (failed bool, err error) {
	tt.cancelStressers()
	defer tt.startStressers()

	plog.Printf("%s updating current revisions...", tt.logPrefix())
	var (
		revs   map[string]int64
		hashes map[string]int64
		rerr   error
		ok     bool
	)
	for i := 0; i < 7; i++ {
		time.Sleep(time.Second)

		revs, hashes, rerr = tt.cluster.getRevisionHash()
		if rerr != nil {
			plog.Printf("%s #%d failed to get current revisions (%v)", tt.logPrefix(), i, rerr)
			continue
		}
		if tt.currentRevision, ok = getSameValue(revs); ok {
			break
		}

		plog.Printf("%s #%d inconsistent current revisions %+v", tt.logPrefix(), i, revs)
	}
	plog.Printf("%s updated current revisions with %d", tt.logPrefix(), tt.currentRevision)

	if !ok || rerr != nil {
		plog.Printf("%s checking current revisions failed [revisions: %v]", tt.logPrefix(), revs)
		failed = true
		err = tt.cleanup()
		return
	}
	plog.Printf("%s all members are consistent with current revisions [revisions: %v]", tt.logPrefix(), revs)

	plog.Printf("%s checking current storage hashes...", tt.logPrefix())
	if _, ok = getSameValue(hashes); !ok {
		plog.Printf("%s checking current storage hashes failed [hashes: %v]", tt.logPrefix(), hashes)
		failed = true
		err = tt.cleanup()
		return
	}
	plog.Printf("%s all members are consistent with storage hashes", tt.logPrefix())
	return
}

func (tt *tester) compact(rev int64, timeout time.Duration) error {
	plog.Printf("%s compacting storage (current revision %d, compact revision %d)", tt.logPrefix(), tt.currentRevision, rev)
	if err := tt.cluster.compactKV(rev, timeout); err != nil {
		if cerr := tt.cleanup(); cerr != nil {
			return fmt.Errorf("%s, %s", err, cerr)
		}
		return err
	}
	plog.Printf("%s compacted storage (compact revision %d)", tt.logPrefix(), rev)

	plog.Printf("%s checking compaction (compact revision %d)", tt.logPrefix(), rev)
	if err := tt.cluster.checkCompact(rev); err != nil {
		plog.Printf("%s checkCompact error (%v)", tt.logPrefix(), err)
		if cerr := tt.cleanup(); cerr != nil {
			return fmt.Errorf("%s, %s", err, cerr)
		}
		return err
	}

	plog.Printf("%s confirmed compaction (compact revision %d)", tt.logPrefix(), rev)
	return nil
}

func (tt *tester) defrag() error {
	plog.Printf("%s defragmenting...", tt.logPrefix())
	if err := tt.cluster.defrag(); err != nil {
		plog.Printf("%s defrag error (%v)", tt.logPrefix(), err)
		if cerr := tt.cleanup(); cerr != nil {
			return fmt.Errorf("%s, %s", err, cerr)
		}
		return err
	}

	plog.Printf("%s defragmented...", tt.logPrefix())
	return nil
}

func (tt *tester) logPrefix() string {
	var (
		rd     = tt.status.getRound()
		cs     = tt.status.getCase()
		prefix = fmt.Sprintf("[round#%d case#%d]", rd, cs)
	)
	if cs == -1 {
		prefix = fmt.Sprintf("[round#%d]", rd)
	}
	return prefix
}

func (tt *tester) cleanup() error {
	roundFailedTotalCounter.Inc()
	desc := "compact/defrag"
	if tt.status.Case != -1 {
		desc = tt.failures[tt.status.Case].Desc()
	}
	caseFailedTotalCounter.WithLabelValues(desc).Inc()

	plog.Printf("%s cleaning up...", tt.logPrefix())
	if err := tt.cluster.Cleanup(); err != nil {
		plog.Printf("%s cleanup error: %v", tt.logPrefix(), err)
		return err
	}

	if err := tt.cluster.Bootstrap(); err != nil {
		plog.Printf("%s cleanup Bootstrap error: %v", tt.logPrefix(), err)
		return err
	}

	return nil
}

func (tt *tester) cancelStressers() {
	plog.Printf("%s canceling the stressers...", tt.logPrefix())
	for _, s := range tt.cluster.Stressers {
		s.Cancel()
	}
	plog.Printf("%s canceled stressers", tt.logPrefix())
}

func (tt *tester) startStressers() {
	plog.Printf("%s starting the stressers...", tt.logPrefix())
	for _, s := range tt.cluster.Stressers {
		go s.Stress()
	}
	plog.Printf("%s started stressers", tt.logPrefix())
}
<|endoftext|>


<commit_before>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package swarm

import (
	"fmt"
	"github.com/fsouza/go-dockerclient/testing"
	"github.com/tsuru/tsuru/provision"
	"gopkg.in/check.v1"
	"net"
)

func (s *S) TestAddNode(c *check.C) {
	srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	metadata := map[string]string{"m1": "v1", "m2": "v2", "pool": "p1"}
	opts := provision.AddNodeOptions{
		Address:  srv.URL(),
		Metadata: metadata,
	}
	err = s.p.AddNode(opts)
	c.Assert(err, check.IsNil)
	node, err := s.p.GetNode(srv.URL())
	c.Assert(err, check.IsNil)
	c.Assert(node.Address(), check.Equals, srv.URL())
	c.Assert(node.Metadata(), check.DeepEquals, metadata)
	c.Assert(node.Pool(), check.Equals, "p1")
	c.Assert(node.Status(), check.Equals, "ready")
	coll, err := nodeAddrCollection()
	c.Assert(err, check.IsNil)
	var all []NodeAddr
	err = coll.Find(nil).All(&all)
	c.Assert(err, check.IsNil)
	_, port, _ := net.SplitHostPort(srv.SwarmAddress())
	c.Assert(all, check.DeepEquals, []NodeAddr{{DockerAddress: srv.URL(), SwarmAddress: "127.0.0.1:" + port}})
}

func (s *S) TestAddNodeMultiple(c *check.C) {
	for i := 0; i < 5; i++ {
		srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
		c.Assert(err, check.IsNil)
		metadata := map[string]string{"count": fmt.Sprintf("%d", i), "pool": "p1"}
		opts := provision.AddNodeOptions{
			Address:  srv.URL(),
			Metadata: metadata,
		}
		err = s.p.AddNode(opts)
		c.Assert(err, check.IsNil)
	}
	nodes, err := s.p.ListNodes(nil)
	c.Assert(err, check.IsNil)
	c.Assert(nodes, check.HasLen, 5)
	for i, n := range nodes {
		c.Assert(n.Metadata(), check.DeepEquals, map[string]string{
			"count": fmt.Sprintf("%d", i),
			"pool":  "p1",
		})
	}
}
<commit_msg>provision/swarm: fix import formatting<commit_after>// Copyright 2016 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package swarm

import (
	"fmt"
	"net"

	"github.com/fsouza/go-dockerclient/testing"
	"github.com/tsuru/tsuru/provision"
	"gopkg.in/check.v1"
)

func (s *S) TestAddNode(c *check.C) {
	srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
	c.Assert(err, check.IsNil)
	metadata := map[string]string{"m1": "v1", "m2": "v2", "pool": "p1"}
	opts := provision.AddNodeOptions{
		Address:  srv.URL(),
		Metadata: metadata,
	}
	err = s.p.AddNode(opts)
	c.Assert(err, check.IsNil)
	node, err := s.p.GetNode(srv.URL())
	c.Assert(err, check.IsNil)
	c.Assert(node.Address(), check.Equals, srv.URL())
	c.Assert(node.Metadata(), check.DeepEquals, metadata)
	c.Assert(node.Pool(), check.Equals, "p1")
	c.Assert(node.Status(), check.Equals, "ready")
	coll, err := nodeAddrCollection()
	c.Assert(err, check.IsNil)
	var all []NodeAddr
	err = coll.Find(nil).All(&all)
	c.Assert(err, check.IsNil)
	_, port, _ := net.SplitHostPort(srv.SwarmAddress())
	c.Assert(all, check.DeepEquals, []NodeAddr{{DockerAddress: srv.URL(), SwarmAddress: "127.0.0.1:" + port}})
}

func (s *S) TestAddNodeMultiple(c *check.C) {
	for i := 0; i < 5; i++ {
		srv, err := testing.NewServer("127.0.0.1:0", nil, nil)
		c.Assert(err, check.IsNil)
		metadata := map[string]string{"count": fmt.Sprintf("%d", i), "pool": "p1"}
		opts := provision.AddNodeOptions{
			Address:  srv.URL(),
			Metadata: metadata,
		}
		err = s.p.AddNode(opts)
		c.Assert(err, check.IsNil)
	}
	nodes, err := s.p.ListNodes(nil)
	c.Assert(err, check.IsNil)
	c.Assert(nodes, check.HasLen, 5)
	for i, n := range nodes {
		c.Assert(n.Metadata(), check.DeepEquals, map[string]string{
			"count": fmt.Sprintf("%d", i),
			"pool":  "p1",
		})
	}
}
<|endoftext|>


<commit_before>// Copyright ©2013 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Based on the CholeskyDecomposition class from Jama 1.0.3.

package mat64

import (
	"math"

	"github.com/gonum/blas"
	"github.com/gonum/blas/blas64"
	"github.com/gonum/lapack/lapack64"
	"github.com/gonum/matrix"
)

const badTriangle = "mat64: invalid triangle"

// Cholesky is a type for creating and using the Cholesky factorization of a
// symmetric positive definite matrix.
type Cholesky struct {
	chol *TriDense
	cond float64
}

// updateCond updates the condition number of the Cholesky decomposition. If
// norm > 0, then that norm is used as the norm of the original matrix A, otherwise
// the norm is estimated from the decomposition.
func (c *Cholesky) updateCond(norm float64) {
	n := c.chol.mat.N
	work := make([]float64, 3*n)
	if norm < 0 {
		// This is an approximation. By the definition of a norm, ||AB|| <= ||A|| ||B||.
		// Here, A = U^T * U.
		// The condition number is ||A|| || A^-1||, so this will underestimate
		// the condition number somewhat.
		// The norm of the original factorized matrix cannot be stored because of
		// update possibilities.
		unorm := lapack64.Lantr(matrix.CondNorm, c.chol.mat, work)
		lnorm := lapack64.Lantr(matrix.CondNormTrans, c.chol.mat, work)
		norm = unorm * lnorm
	}
	sym := c.chol.asSymBlas()
	iwork := make([]int, n)
	v := lapack64.Pocon(sym, norm, work, iwork)
	c.cond = 1 / v
}

// Factorize calculates the Cholesky decomposition of the matrix A and returns
// whether the matrix is positive definite.
func (c *Cholesky) Factorize(a Symmetric) (ok bool) {
	n := a.Symmetric()
	if c.chol == nil {
		c.chol = NewTriDense(n, true, nil)
	} else {
		c.chol = NewTriDense(n, true, use(c.chol.mat.Data, n*n))
	}
	copySymIntoTriangle(c.chol, a)

	sym := c.chol.asSymBlas()
	work := make([]float64, c.chol.mat.N)
	norm := lapack64.Lansy(matrix.CondNorm, sym, work)
	_, ok = lapack64.Potrf(sym)
	if ok {
		c.updateCond(norm)
	} else {
		c.cond = math.Inf(1)
	}
	return ok
}

// Det returns the determinant of the matrix that has been factorized.
func (c *Cholesky) Det() float64 {
	return math.Exp(c.LogDet())
}

// Size returns the dimension of the factorized matrix.
func (c *Cholesky) Size() int {
	return c.chol.mat.N
}

// LogDet returns the log of the determinant of the matrix that has been factorized.
func (c *Cholesky) LogDet() float64 {
	var det float64
	for i := 0; i < c.chol.mat.N; i++ {
		det += 2 * math.Log(c.chol.mat.Data[i*c.chol.mat.Stride+i])
	}
	return det
}

// SolveCholesky finds the matrix m that solves A * m = b where A is represented
// by the cholesky decomposition, placing the result in the receiver.
func (m *Dense) SolveCholesky(chol *Cholesky, b Matrix) error {
	n := chol.chol.mat.N
	bm, bn := b.Dims()
	if n != bm {
		panic(matrix.ErrShape)
	}

	m.reuseAs(bm, bn)
	if b != m {
		m.Copy(b)
	}
	blas64.Trsm(blas.Left, blas.Trans, 1, chol.chol.mat, m.mat)
	blas64.Trsm(blas.Left, blas.NoTrans, 1, chol.chol.mat, m.mat)
	if chol.cond > matrix.ConditionTolerance {
		return matrix.Condition(chol.cond)
	}
	return nil
}

// SolveCholeskyVec finds the vector v that solves A * v = b where A is represented
// by the Cholesky decomposition, placing the result in the receiver.
func (v *Vector) SolveCholeskyVec(chol *Cholesky, b *Vector) error {
	n := chol.chol.mat.N
	vn := b.Len()
	if vn != n {
		panic(matrix.ErrShape)
	}
	if v != b {
		v.checkOverlap(b.mat)
	}
	v.reuseAs(n)
	if v != b {
		v.CopyVec(b)
	}
	blas64.Trsv(blas.Trans, chol.chol.mat, v.mat)
	blas64.Trsv(blas.NoTrans, chol.chol.mat, v.mat)
	if chol.cond > matrix.ConditionTolerance {
		return matrix.Condition(chol.cond)
	}
	return nil

}

// UFromCholesky extracts the n×n upper triangular matrix U from a Choleksy
// decomposition
//  A = U^T * U.
func (t *TriDense) UFromCholesky(chol *Cholesky) {
	n := chol.chol.mat.N
	t.reuseAs(n, true)
	t.Copy(chol.chol)
}

// LFromCholesky extracts the n×n lower triangular matrix U from a Choleksy
// decomposition
//  A = L * L^T.
func (t *TriDense) LFromCholesky(chol *Cholesky) {
	n := chol.chol.mat.N
	t.reuseAs(n, false)
	t.Copy(chol.chol.TTri())
}

// FromCholesky reconstructs the original positive definite matrix given its
// Cholesky decomposition.
func (s *SymDense) FromCholesky(chol *Cholesky) {
	n := chol.chol.mat.N
	s.reuseAs(n)
	s.SymOuterK(1, chol.chol.T())
}

// InverseCholesky computes the inverse of the matrix represented by its Cholesky
// factorization and stores the result into the receiver. If the factorized
// matrix is ill-conditioned, a Condition error will be returned.
// Note that matrix inversion is numerically unstable, and should generally be
// avoided where possible, for example by using the Solve routines.
func (s *SymDense) InverseCholesky(chol *Cholesky) error {
	// TODO(btracey): Replace this code with a direct call to Dpotri when it
	// is available.
	s.reuseAs(chol.chol.mat.N)
	// If:
	//  chol(A) = U^T * U
	// Then:
	//  chol(A^-1) = S * S^T
	// where S = U^-1
	var t TriDense
	err := t.InverseTri(chol.chol)
	s.SymOuterK(1, &t)
	return err
}
<commit_msg>mat64: use matrix.Upper and matrix.Lower constants in Cholesky<commit_after>// Copyright ©2013 The gonum Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Based on the CholeskyDecomposition class from Jama 1.0.3.

package mat64

import (
	"math"

	"github.com/gonum/blas"
	"github.com/gonum/blas/blas64"
	"github.com/gonum/lapack/lapack64"
	"github.com/gonum/matrix"
)

const badTriangle = "mat64: invalid triangle"

// Cholesky is a type for creating and using the Cholesky factorization of a
// symmetric positive definite matrix.
type Cholesky struct {
	chol *TriDense
	cond float64
}

// updateCond updates the condition number of the Cholesky decomposition. If
// norm > 0, then that norm is used as the norm of the original matrix A, otherwise
// the norm is estimated from the decomposition.
func (c *Cholesky) updateCond(norm float64) {
	n := c.chol.mat.N
	work := make([]float64, 3*n)
	if norm < 0 {
		// This is an approximation. By the definition of a norm, ||AB|| <= ||A|| ||B||.
		// Here, A = U^T * U.
		// The condition number is ||A|| || A^-1||, so this will underestimate
		// the condition number somewhat.
		// The norm of the original factorized matrix cannot be stored because of
		// update possibilities.
		unorm := lapack64.Lantr(matrix.CondNorm, c.chol.mat, work)
		lnorm := lapack64.Lantr(matrix.CondNormTrans, c.chol.mat, work)
		norm = unorm * lnorm
	}
	sym := c.chol.asSymBlas()
	iwork := make([]int, n)
	v := lapack64.Pocon(sym, norm, work, iwork)
	c.cond = 1 / v
}

// Factorize calculates the Cholesky decomposition of the matrix A and returns
// whether the matrix is positive definite.
func (c *Cholesky) Factorize(a Symmetric) (ok bool) {
	n := a.Symmetric()
	if c.chol == nil {
		c.chol = NewTriDense(n, matrix.Upper, nil)
	} else {
		c.chol = NewTriDense(n, matrix.Upper, use(c.chol.mat.Data, n*n))
	}
	copySymIntoTriangle(c.chol, a)

	sym := c.chol.asSymBlas()
	work := make([]float64, c.chol.mat.N)
	norm := lapack64.Lansy(matrix.CondNorm, sym, work)
	_, ok = lapack64.Potrf(sym)
	if ok {
		c.updateCond(norm)
	} else {
		c.cond = math.Inf(1)
	}
	return ok
}

// Det returns the determinant of the matrix that has been factorized.
func (c *Cholesky) Det() float64 {
	return math.Exp(c.LogDet())
}

// Size returns the dimension of the factorized matrix.
func (c *Cholesky) Size() int {
	return c.chol.mat.N
}

// LogDet returns the log of the determinant of the matrix that has been factorized.
func (c *Cholesky) LogDet() float64 {
	var det float64
	for i := 0; i < c.chol.mat.N; i++ {
		det += 2 * math.Log(c.chol.mat.Data[i*c.chol.mat.Stride+i])
	}
	return det
}

// SolveCholesky finds the matrix m that solves A * m = b where A is represented
// by the cholesky decomposition, placing the result in the receiver.
func (m *Dense) SolveCholesky(chol *Cholesky, b Matrix) error {
	n := chol.chol.mat.N
	bm, bn := b.Dims()
	if n != bm {
		panic(matrix.ErrShape)
	}

	m.reuseAs(bm, bn)
	if b != m {
		m.Copy(b)
	}
	blas64.Trsm(blas.Left, blas.Trans, 1, chol.chol.mat, m.mat)
	blas64.Trsm(blas.Left, blas.NoTrans, 1, chol.chol.mat, m.mat)
	if chol.cond > matrix.ConditionTolerance {
		return matrix.Condition(chol.cond)
	}
	return nil
}

// SolveCholeskyVec finds the vector v that solves A * v = b where A is represented
// by the Cholesky decomposition, placing the result in the receiver.
func (v *Vector) SolveCholeskyVec(chol *Cholesky, b *Vector) error {
	n := chol.chol.mat.N
	vn := b.Len()
	if vn != n {
		panic(matrix.ErrShape)
	}
	if v != b {
		v.checkOverlap(b.mat)
	}
	v.reuseAs(n)
	if v != b {
		v.CopyVec(b)
	}
	blas64.Trsv(blas.Trans, chol.chol.mat, v.mat)
	blas64.Trsv(blas.NoTrans, chol.chol.mat, v.mat)
	if chol.cond > matrix.ConditionTolerance {
		return matrix.Condition(chol.cond)
	}
	return nil

}

// UFromCholesky extracts the n×n upper triangular matrix U from a Choleksy
// decomposition
//  A = U^T * U.
func (t *TriDense) UFromCholesky(chol *Cholesky) {
	n := chol.chol.mat.N
	t.reuseAs(n, matrix.Upper)
	t.Copy(chol.chol)
}

// LFromCholesky extracts the n×n lower triangular matrix U from a Choleksy
// decomposition
//  A = L * L^T.
func (t *TriDense) LFromCholesky(chol *Cholesky) {
	n := chol.chol.mat.N
	t.reuseAs(n, matrix.Lower)
	t.Copy(chol.chol.TTri())
}

// FromCholesky reconstructs the original positive definite matrix given its
// Cholesky decomposition.
func (s *SymDense) FromCholesky(chol *Cholesky) {
	n := chol.chol.mat.N
	s.reuseAs(n)
	s.SymOuterK(1, chol.chol.T())
}

// InverseCholesky computes the inverse of the matrix represented by its Cholesky
// factorization and stores the result into the receiver. If the factorized
// matrix is ill-conditioned, a Condition error will be returned.
// Note that matrix inversion is numerically unstable, and should generally be
// avoided where possible, for example by using the Solve routines.
func (s *SymDense) InverseCholesky(chol *Cholesky) error {
	// TODO(btracey): Replace this code with a direct call to Dpotri when it
	// is available.
	s.reuseAs(chol.chol.mat.N)
	// If:
	//  chol(A) = U^T * U
	// Then:
	//  chol(A^-1) = S * S^T
	// where S = U^-1
	var t TriDense
	err := t.InverseTri(chol.chol)
	s.SymOuterK(1, &t)
	return err
}
<|endoftext|>


<commit_before>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package message

import (
	"encoding/xml"
	"io"

	"github.com/issue9/wechat/common/result"
)

// 消息类型
const (
	MsgTypeText                    = "text"
	MsgTypeImage                   = "image"
	MsgTypeVoice                   = "voice"
	MsgTypeShortVideo              = "shortvideo"
	MsgTypeLocation                = "location"
	MsgTypeLink                    = "link"
	MsgTypeEvent                   = "event"
	MsgTypeTransferCustomerService = "transfer_customer_service" // 只能用于回复消息中
)

type Messager interface {
	To() string
	From() string
	Type() string
	Created() int64
}

// MsgText 文本消息
type MsgText struct {
	ToUserName   string `xml:"ToUserName,cdata"`   // 开发者微信号
	FromUserName string `xml:"FromUserName,cdata"` // 发送方帐号（一个 OpenID）
	CreateTime   int64  `xml:"CreateTime"`         // 消息创建时间 （整型）
	Content      string `xml:"Content,cdata"`      // 文本消息内容
	MsgType      string `xml:"MsgType,cdata"`      // 消息类型
	MsgId        int64  `xml:"MsgId"`              // 消息 id，64 位整型
}

// msgType 这不是一个真实存在的消息类型，
// 只是用于解析 xml 中的 MsgType 字段的具体值用的。
type msgType struct {
	MsgType string `xml:"MsgType"`
}

func (m *MsgText) To() string {
	return m.ToUserName
}

func (m *MsgText) From() string {
	return m.FromUserName
}

func (m *MsgText) Type() string {
	// 不采用 m.MsgType，而是直接返回常量
	return MsgTypeText
}

func (m *MsgText) Created() int64 {
	return m.CreateTime
}

// 从指定的数据中分析其消息的类型
func getMsgType(data []byte) (string, error) {
	obj := &msgType{}
	if err := xml.Unmarshal(data, obj); err != nil {
		return "", result.New(600)
	}

	return obj.MsgType, nil
}

func getMsgObj(r io.Reader) (Messager, error) {
	data := make([]byte, 0, 1000)
	_, err := io.ReadFull(r, data)
	if err != nil {
		return nil, err
	}

	typ, err := getMsgType(data)
	if err != nil {
		return nil, err
	}

	var obj Messager
	switch typ {
	case MsgTypeText:
		obj = &MsgText{}
		err = xml.Unmarshal(data, obj)
	case MsgTypeLink:
		// TODO
	}

	if err != nil {
		return nil, err
	}
	return obj, nil
}
<commit_msg>添加一个 MsgImage 类型<commit_after>// Copyright 2016 by caixw, All rights reserved.
// Use of this source code is governed by a MIT
// license that can be found in the LICENSE file.

package message

import (
	"encoding/xml"
	"io"

	"github.com/issue9/wechat/common/result"
)

// 消息类型
const (
	MsgTypeText                    = "text"
	MsgTypeImage                   = "image"
	MsgTypeVoice                   = "voice"
	MsgTypeShortVideo              = "shortvideo"
	MsgTypeLocation                = "location"
	MsgTypeLink                    = "link"
	MsgTypeEvent                   = "event"
	MsgTypeTransferCustomerService = "transfer_customer_service" // 只能用于回复消息中
)

type Messager interface {
	To() string
	From() string
	Type() string
	Created() int64
}

// MsgText 文本消息
type MsgText struct {
	ToUserName   string `xml:"ToUserName,cdata"`   // 开发者微信号
	FromUserName string `xml:"FromUserName,cdata"` // 发送方帐号（一个 OpenID）
	CreateTime   int64  `xml:"CreateTime"`         // 消息创建时间 （整型）
	MsgType      string `xml:"MsgType,cdata"`      // 消息类型
	Content      string `xml:"Content,cdata"`      // 文本消息内容
	MsgID        int64  `xml:"MsgId"`              // 消息 id，64 位整型
}

type MsgImage struct {
	ToUserName   string `xml:"ToUserName,cdata"`   // 开发者微信号
	FromUserName string `xml:"FromUserName,cdata"` // 发送方帐号（一个 OpenID）
	CreateTime   int64  `xml:"CreateTime"`         // 消息创建时间 （整型）
	MsgType      string `xml:"MsgType,cdata"`      // 消息类型
	MsgID        int64  `xml:"MsgId"`              // 消息 id，64 位整型
	PicUrl       string `xml:"PicUrl,cdata"`
	MediaID      string `xml:"MediaId,cdata"`
}

// msgType 这不是一个真实存在的消息类型，
// 只是用于解析 xml 中的 MsgType 字段的具体值用的。
type msgType struct {
	MsgType string `xml:"MsgType"`
}

func (m *MsgText) To() string {
	return m.ToUserName
}

func (m *MsgText) From() string {
	return m.FromUserName
}

func (m *MsgText) Type() string {
	// 不采用 m.MsgType，而是直接返回常量
	return MsgTypeText
}

func (m *MsgText) Created() int64 {
	return m.CreateTime
}

func (m *MsgImage) To() string {
	return m.ToUserName
}

func (m *MsgImage) From() string {
	return m.FromUserName
}

func (m *MsgImage) Type() string {
	// 不采用 m.MsgType，而是直接返回常量
	return MsgTypeImage
}

func (m *MsgImage) Created() int64 {
	return m.CreateTime
}

// 从指定的数据中分析其消息的类型
func getMsgType(data []byte) (string, error) {
	obj := &msgType{}
	if err := xml.Unmarshal(data, obj); err != nil {
		return "", result.New(600)
	}

	return obj.MsgType, nil
}

func getMsgObj(r io.Reader) (Messager, error) {
	data := make([]byte, 0, 1000)
	_, err := io.ReadFull(r, data)
	if err != nil {
		return nil, err
	}

	typ, err := getMsgType(data)
	if err != nil {
		return nil, err
	}

	var obj Messager
	switch typ {
	case MsgTypeText:
		obj = &MsgText{}
		err = xml.Unmarshal(data, obj)
	case MsgTypeImage:
		obj = &MsgImage{}
		err = xml.Unmarshal(data, obj)
	}

	if err != nil {
		return nil, err
	}
	return obj, nil
}
<|endoftext|>