<commit_before>/*
Copyright 2015 Juniper Networks, Inc. All rights reserved.

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

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

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

package opencontrail

import (
	flag "github.com/spf13/pflag"
)

const (
	DefaultDomain = "default-domain"
)

type Config struct {
	ApiAddress string
	ApiPort    int

	DefaultProject string
	PublicNetwork  string
	ServiceNetwork  string

	PublicSubnet  string
	PrivateSubnet string
	ServiceSubnet string

	NetworkTag       string
	NetworkAccessTag string
}

func NewConfig() *Config {
	config := &Config{
		ApiAddress:       "localhost",
		ApiPort:          8082,
		DefaultProject:   "default-domain:default-project",
		PublicNetwork:    "default-domain:default-project:Public",
		ServiceNetwork:   "default-domain:default-project:Service",
		PrivateSubnet:    "10.10.10.0/24",
		ServiceSubnet:    "10.10.20.0/24",
		PublicSubnet:     "10.0.2.0/24",
		NetworkTag:       "name",
		NetworkAccessTag: "uses",
	}
	return config
}

func (c *Config) Parse(args []string) {
	fs := flag.NewFlagSet("opencontrail", flag.ExitOnError)
	fs.StringVar(&c.ApiAddress, "contrail_api", c.ApiAddress,
		"Hostname or address for the OpenContrail API server.")
	fs.IntVar(&c.ApiPort, "contrail_port", 8082,
		"OpenContrail API port.")
	fs.StringVar(&c.PublicNetwork, "public_name", c.PublicNetwork,
		"External network name.")
	fs.StringVar(&c.PublicSubnet, "public_net", "",
		"External network subnet prefix used when provisioning the cluster.")
	fs.StringVar(&c.PrivateSubnet, "private_net", c.PrivateSubnet,
		"Address range to use for private IP addresses.")
	fs.StringVar(&c.ServiceSubnet, "portal_net", c.ServiceSubnet,
		"Address range to use for services.")
	fs.StringVar(&c.NetworkTag, "network_label", c.NetworkTag,
		"Label used to specify the network used by the resource (pod or service).")
	fs.StringVar(&c.NetworkAccessTag, "access_label", c.NetworkAccessTag,
		"Label used to determine what services this resource (pod/rc) accesses.")
	fs.Parse(args)
}
<commit_msg>update service network subnet<commit_after>/*
Copyright 2015 Juniper Networks, Inc. All rights reserved.

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

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

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

package opencontrail

import (
	flag "github.com/spf13/pflag"
)

const (
	DefaultDomain = "default-domain"
)

type Config struct {
	ApiAddress string
	ApiPort    int

	DefaultProject string
	PublicNetwork  string
	ServiceNetwork  string

	PublicSubnet  string
	PrivateSubnet string
	ServiceSubnet string

	NetworkTag       string
	NetworkAccessTag string
}

func NewConfig() *Config {
	config := &Config{
		ApiAddress:       "localhost",
		ApiPort:          8082,
		DefaultProject:   "default-domain:default-project",
		PublicNetwork:    "default-domain:default-project:Public",
		ServiceNetwork:   "default-domain:default-project:Service",
		PrivateSubnet:    "10.10.10.0/24",
		ServiceSubnet:    "10.247.0.0/24",
		PublicSubnet:     "10.0.2.0/24",
		NetworkTag:       "name",
		NetworkAccessTag: "uses",
	}
	return config
}

func (c *Config) Parse(args []string) {
	fs := flag.NewFlagSet("opencontrail", flag.ExitOnError)
	fs.StringVar(&c.ApiAddress, "contrail_api", c.ApiAddress,
		"Hostname or address for the OpenContrail API server.")
	fs.IntVar(&c.ApiPort, "contrail_port", 8082,
		"OpenContrail API port.")
	fs.StringVar(&c.PublicNetwork, "public_name", c.PublicNetwork,
		"External network name.")
	fs.StringVar(&c.PublicSubnet, "public_net", "",
		"External network subnet prefix used when provisioning the cluster.")
	fs.StringVar(&c.PrivateSubnet, "private_net", c.PrivateSubnet,
		"Address range to use for private IP addresses.")
	fs.StringVar(&c.ServiceSubnet, "portal_net", c.ServiceSubnet,
		"Address range to use for services.")
	fs.StringVar(&c.NetworkTag, "network_label", c.NetworkTag,
		"Label used to specify the network used by the resource (pod or service).")
	fs.StringVar(&c.NetworkAccessTag, "access_label", c.NetworkAccessTag,
		"Label used to determine what services this resource (pod/rc) accesses.")
	fs.Parse(args)
}
<|endoftext|>


<commit_before>// Package observability provides tools for working with open census.
package observability

import (
	"context"
	"strconv"

	"github.com/google/exposure-notifications-server/pkg/logging"
	"github.com/google/exposure-notifications-verification-server/pkg/buildinfo"
	"go.opencensus.io/tag"
)

const (
	MetricRoot = "en-verification-server"
)

var (
	BuildIDTagKey  = tag.MustNewKey("build_id")
	BuildTagTagKey = tag.MustNewKey("build_tag")
	RealmTagKey    = tag.MustNewKey("realm")
)

// CommonTagKeys returns the slice of common tag keys that should used in all
// views.
func CommonTagKeys() []tag.Key {
	return []tag.Key{
		BuildIDTagKey,
		BuildTagTagKey,
		RealmTagKey,
	}
}

// WithRealmID creates a new context with the realm id attached to the
// observability context.
func WithRealmID(octx context.Context, realmID uint) context.Context {
	realmIDStr := strconv.FormatUint(uint64(realmID), 10)
	ctx, err := tag.New(octx, tag.Upsert(RealmTagKey, realmIDStr))
	if err != nil {
		logger := logging.FromContext(octx).Named("observability.WithRealmID")
		logger.Errorw("failed to upsert realm on observability context",
			"error", err,
			"realm_id", realmIDStr)
		return octx
	}
	return ctx
}

// WithBuildInfo creates a new context with the build info attached to the
// observability context.
func WithBuildInfo(octx context.Context) context.Context {
	ctx, err := tag.New(octx,
		tag.Upsert(BuildIDTagKey, buildinfo.BuildID),
		tag.Upsert(BuildTagTagKey, buildinfo.BuildTag))
	if err != nil {
		logger := logging.FromContext(octx).Named("observability.WithBuildInfo")
		logger.Errorw("failed to upsert buildinfo on observability context", "error", err)
		return octx
	}
	return ctx
}
<commit_msg>Include knative info in metrics if present (#736)<commit_after>// Package observability provides tools for working with open census.
package observability

import (
	"context"
	"os"
	"strconv"

	"github.com/google/exposure-notifications-server/pkg/logging"
	"github.com/google/exposure-notifications-verification-server/pkg/buildinfo"
	"go.opencensus.io/tag"
)

const (
	MetricRoot = "en-verification-server"
)

var (
	BuildIDTagKey  = tag.MustNewKey("build_id")
	BuildTagTagKey = tag.MustNewKey("build_tag")

	KnativeServiceTagKey       = tag.MustNewKey("k_service")
	KnativeRevisionTagKey      = tag.MustNewKey("k_revision")
	KnativeConfigurationTagKey = tag.MustNewKey("k_configuration")

	RealmTagKey = tag.MustNewKey("realm")

	knativeService       = os.Getenv("K_SERVICE")
	knativeRevision      = os.Getenv("K_REVISION")
	knativeConfiguration = os.Getenv("K_CONFIGURATION")
)

// CommonTagKeys returns the slice of common tag keys that should used in all
// views.
func CommonTagKeys() []tag.Key {
	return []tag.Key{
		BuildIDTagKey,
		BuildTagTagKey,
		RealmTagKey,
	}
}

// WithRealmID creates a new context with the realm id attached to the
// observability context.
func WithRealmID(octx context.Context, realmID uint) context.Context {
	realmIDStr := strconv.FormatUint(uint64(realmID), 10)
	ctx, err := tag.New(octx, tag.Upsert(RealmTagKey, realmIDStr))
	if err != nil {
		logger := logging.FromContext(octx).Named("observability.WithRealmID")
		logger.Errorw("failed to upsert realm on observability context",
			"error", err,
			"realm_id", realmIDStr)
		return octx
	}
	return ctx
}

// WithBuildInfo creates a new context with the build and revision info attached
// to the observability context.
func WithBuildInfo(octx context.Context) context.Context {
	tags := make([]tag.Mutator, 0, 5)
	tags = append(tags, tag.Upsert(BuildIDTagKey, buildinfo.BuildID))
	tags = append(tags, tag.Upsert(BuildTagTagKey, buildinfo.BuildTag))

	if knativeService != "" {
		tags = append(tags, tag.Upsert(KnativeServiceTagKey, knativeService))
	}

	if knativeRevision != "" {
		tags = append(tags, tag.Upsert(KnativeRevisionTagKey, knativeRevision))
	}

	if knativeConfiguration != "" {
		tags = append(tags, tag.Upsert(KnativeConfigurationTagKey, knativeConfiguration))
	}

	ctx, err := tag.New(octx, tags...)
	if err != nil {
		logger := logging.FromContext(octx).Named("observability.WithBuildInfo")
		logger.Errorw("failed to upsert buildinfo on observability context", "error", err)
		return octx
	}

	return ctx
}
<|endoftext|>


<commit_before>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package registrydisk

import (
	"errors"
	"fmt"
	"os"

	"github.com/jeevatkm/go-model"

	kubev1 "k8s.io/api/core/v1"

	"kubevirt.io/kubevirt/pkg/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/precond"
)

const registryDiskV1Alpha = "RegistryDisk:v1alpha"
const defaultIqn = "iqn.2017-01.io.kubevirt:wrapper/1"
const defaultPort = 3261
const defaultPortStr = "3261"
const filePrefix = "disk-image"
const defaultHost = "127.0.0.1"

var registryDiskOwner = "qemu"

var mountBaseDir = "/var/run/libvirt/kubevirt-disk-dir"

func generateVMBaseDir(vm *v1.VirtualMachine) string {
	domain := precond.MustNotBeEmpty(vm.GetObjectMeta().GetName())
	namespace := precond.MustNotBeEmpty(vm.GetObjectMeta().GetNamespace())
	return fmt.Sprintf("%s/%s/%s", mountBaseDir, namespace, domain)
}
func generateVolumeMountDir(vm *v1.VirtualMachine, diskCount int) string {
	baseDir := generateVMBaseDir(vm)
	return fmt.Sprintf("%s/disk%d", baseDir, diskCount)
}

func k8sSecretName(vm *v1.VirtualMachine) string {
	return fmt.Sprintf("registrydisk-iscsi-%s-%s", vm.GetObjectMeta().GetNamespace(), vm.GetObjectMeta().GetName())
}

func SetLocalDirectory(dir string) error {
	mountBaseDir = dir
	return nil
}

// The unit test suite uses this function
func SetLocalDataOwner(user string) {
	registryDiskOwner = user
}

func getFilePath(basePath string) (string, string, error) {
	rawPath := basePath + "/" + filePrefix + ".raw"
	qcow2Path := basePath + "/" + filePrefix + ".qcow2"

	exists, err := diskutils.FileExists(rawPath)
	if err != nil {
		return "", "", err
	} else if exists {
		return rawPath, "raw", nil
	}

	exists, err = diskutils.FileExists(qcow2Path)
	if err != nil {
		return "", "", err
	} else if exists {
		return qcow2Path, "qcow2", nil
	}

	return "", "", errors.New(fmt.Sprintf("no supported file disk found in directory %s", basePath))
}

func CleanupEphemeralDisks(vm *v1.VirtualMachine) error {
	volumeMountDir := generateVMBaseDir(vm)
	err := os.RemoveAll(volumeMountDir)
	if err != nil && os.IsNotExist(err) {
		return nil
	}
	return err
}

// The virt-handler converts registry disks to their corresponding iscsi network
// disks when the VM spec is being defined as a domain with libvirt.
// The ports and host of the iscsi disks are already provided here by the controller.
func MapRegistryDisks(vm *v1.VirtualMachine) (*v1.VirtualMachine, error) {
	vmCopy := &v1.VirtualMachine{}
	model.Copy(vmCopy, vm)

	for diskCount, disk := range vmCopy.Spec.Domain.Devices.Disks {
		if disk.Type == registryDiskV1Alpha {
			volumeMountDir := generateVolumeMountDir(vm, diskCount)

			diskPath, diskType, err := getFilePath(volumeMountDir)
			if err != nil {
				return vm, err
			}

			// Rename file to release management of it from container process.
			oldDiskPath := diskPath
			diskPath = oldDiskPath + ".virt"
			err = os.Rename(oldDiskPath, diskPath)
			if err != nil {
				return vm, err
			}

			err = diskutils.SetFileOwnership(registryDiskOwner, diskPath)
			if err != nil {
				return vm, err
			}

			newDisk := v1.Disk{}
			newDisk.Type = "file"
			newDisk.Device = "disk"
			newDisk.Driver = &v1.DiskDriver{
				Type: diskType,
				Name: "qemu",
			}
			newDisk.Source.File = diskPath
			newDisk.Target = disk.Target
			vmCopy.Spec.Domain.Devices.Disks[diskCount] = newDisk
		}
	}

	return vmCopy, nil
}

// The controller uses this function to generate the container
// specs for hosting the container registry disks.
func GenerateContainers(vm *v1.VirtualMachine) ([]kubev1.Container, []kubev1.Volume, error) {
	var containers []kubev1.Container
	var volumes []kubev1.Volume

	initialDelaySeconds := 2
	timeoutSeconds := 5
	periodSeconds := 5
	successThreshold := 2
	failureThreshold := 5

	// Make VM Image Wrapper Containers
	for diskCount, disk := range vm.Spec.Domain.Devices.Disks {
		if disk.Type == registryDiskV1Alpha {

			volumeMountDir := generateVolumeMountDir(vm, diskCount)
			volumeName := fmt.Sprintf("disk%d-volume", diskCount)
			diskContainerName := fmt.Sprintf("disk%d", diskCount)
			// container image is disk.Source.Name
			diskContainerImage := disk.Source.Name

			volumes = append(volumes, kubev1.Volume{
				Name: volumeName,
				VolumeSource: kubev1.VolumeSource{
					HostPath: &kubev1.HostPathVolumeSource{
						Path: volumeMountDir,
					},
				},
			})
			containers = append(containers, kubev1.Container{
				Name:            diskContainerName,
				Image:           diskContainerImage,
				ImagePullPolicy: kubev1.PullIfNotPresent,
				Command:         []string{"/entry-point.sh"},
				Env: []kubev1.EnvVar{
					kubev1.EnvVar{
						Name:  "COPY_PATH",
						Value: volumeMountDir + "/" + filePrefix,
					},
				},
				VolumeMounts: []kubev1.VolumeMount{
					{
						Name:      volumeName,
						MountPath: volumeMountDir,
					},
				},
				// The readiness probes ensure the disk coversion and copy finished
				// before the container is marked as "Ready: True"
				ReadinessProbe: &kubev1.Probe{
					Handler: kubev1.Handler{
						Exec: &kubev1.ExecAction{
							Command: []string{
								"cat",
								"/tmp/healthy",
							},
						},
					},
					InitialDelaySeconds: int32(initialDelaySeconds),
					PeriodSeconds:       int32(periodSeconds),
					TimeoutSeconds:      int32(timeoutSeconds),
					SuccessThreshold:    int32(successThreshold),
					FailureThreshold:    int32(failureThreshold),
				},
			})
		}
	}
	return containers, volumes, nil
}
<commit_msg>Remove dead RegistryDisk code from before rewrite<commit_after>/*
 * This file is part of the KubeVirt project
 *
 * 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.
 *
 * Copyright 2017 Red Hat, Inc.
 *
 */

package registrydisk

import (
	"errors"
	"fmt"
	"os"

	"github.com/jeevatkm/go-model"

	kubev1 "k8s.io/api/core/v1"

	"kubevirt.io/kubevirt/pkg/api/v1"
	diskutils "kubevirt.io/kubevirt/pkg/ephemeral-disk-utils"
	"kubevirt.io/kubevirt/pkg/precond"
)

const registryDiskV1Alpha = "RegistryDisk:v1alpha"
const filePrefix = "disk-image"

var registryDiskOwner = "qemu"

var mountBaseDir = "/var/run/libvirt/kubevirt-disk-dir"

func generateVMBaseDir(vm *v1.VirtualMachine) string {
	domain := precond.MustNotBeEmpty(vm.GetObjectMeta().GetName())
	namespace := precond.MustNotBeEmpty(vm.GetObjectMeta().GetNamespace())
	return fmt.Sprintf("%s/%s/%s", mountBaseDir, namespace, domain)
}
func generateVolumeMountDir(vm *v1.VirtualMachine, diskCount int) string {
	baseDir := generateVMBaseDir(vm)
	return fmt.Sprintf("%s/disk%d", baseDir, diskCount)
}

func SetLocalDirectory(dir string) error {
	mountBaseDir = dir
	return nil
}

// The unit test suite uses this function
func SetLocalDataOwner(user string) {
	registryDiskOwner = user
}

func getFilePath(basePath string) (string, string, error) {
	rawPath := basePath + "/" + filePrefix + ".raw"
	qcow2Path := basePath + "/" + filePrefix + ".qcow2"

	exists, err := diskutils.FileExists(rawPath)
	if err != nil {
		return "", "", err
	} else if exists {
		return rawPath, "raw", nil
	}

	exists, err = diskutils.FileExists(qcow2Path)
	if err != nil {
		return "", "", err
	} else if exists {
		return qcow2Path, "qcow2", nil
	}

	return "", "", errors.New(fmt.Sprintf("no supported file disk found in directory %s", basePath))
}

func CleanupEphemeralDisks(vm *v1.VirtualMachine) error {
	volumeMountDir := generateVMBaseDir(vm)
	err := os.RemoveAll(volumeMountDir)
	if err != nil && os.IsNotExist(err) {
		return nil
	}
	return err
}

// The virt-handler converts registry disks to their corresponding iscsi network
// disks when the VM spec is being defined as a domain with libvirt.
// The ports and host of the iscsi disks are already provided here by the controller.
func MapRegistryDisks(vm *v1.VirtualMachine) (*v1.VirtualMachine, error) {
	vmCopy := &v1.VirtualMachine{}
	model.Copy(vmCopy, vm)

	for diskCount, disk := range vmCopy.Spec.Domain.Devices.Disks {
		if disk.Type == registryDiskV1Alpha {
			volumeMountDir := generateVolumeMountDir(vm, diskCount)

			diskPath, diskType, err := getFilePath(volumeMountDir)
			if err != nil {
				return vm, err
			}

			// Rename file to release management of it from container process.
			oldDiskPath := diskPath
			diskPath = oldDiskPath + ".virt"
			err = os.Rename(oldDiskPath, diskPath)
			if err != nil {
				return vm, err
			}

			err = diskutils.SetFileOwnership(registryDiskOwner, diskPath)
			if err != nil {
				return vm, err
			}

			newDisk := v1.Disk{}
			newDisk.Type = "file"
			newDisk.Device = "disk"
			newDisk.Driver = &v1.DiskDriver{
				Type: diskType,
				Name: "qemu",
			}
			newDisk.Source.File = diskPath
			newDisk.Target = disk.Target
			vmCopy.Spec.Domain.Devices.Disks[diskCount] = newDisk
		}
	}

	return vmCopy, nil
}

// The controller uses this function to generate the container
// specs for hosting the container registry disks.
func GenerateContainers(vm *v1.VirtualMachine) ([]kubev1.Container, []kubev1.Volume, error) {
	var containers []kubev1.Container
	var volumes []kubev1.Volume

	initialDelaySeconds := 2
	timeoutSeconds := 5
	periodSeconds := 5
	successThreshold := 2
	failureThreshold := 5

	// Make VM Image Wrapper Containers
	for diskCount, disk := range vm.Spec.Domain.Devices.Disks {
		if disk.Type == registryDiskV1Alpha {

			volumeMountDir := generateVolumeMountDir(vm, diskCount)
			volumeName := fmt.Sprintf("disk%d-volume", diskCount)
			diskContainerName := fmt.Sprintf("disk%d", diskCount)
			// container image is disk.Source.Name
			diskContainerImage := disk.Source.Name

			volumes = append(volumes, kubev1.Volume{
				Name: volumeName,
				VolumeSource: kubev1.VolumeSource{
					HostPath: &kubev1.HostPathVolumeSource{
						Path: volumeMountDir,
					},
				},
			})
			containers = append(containers, kubev1.Container{
				Name:            diskContainerName,
				Image:           diskContainerImage,
				ImagePullPolicy: kubev1.PullIfNotPresent,
				Command:         []string{"/entry-point.sh"},
				Env: []kubev1.EnvVar{
					kubev1.EnvVar{
						Name:  "COPY_PATH",
						Value: volumeMountDir + "/" + filePrefix,
					},
				},
				VolumeMounts: []kubev1.VolumeMount{
					{
						Name:      volumeName,
						MountPath: volumeMountDir,
					},
				},
				// The readiness probes ensure the disk coversion and copy finished
				// before the container is marked as "Ready: True"
				ReadinessProbe: &kubev1.Probe{
					Handler: kubev1.Handler{
						Exec: &kubev1.ExecAction{
							Command: []string{
								"cat",
								"/tmp/healthy",
							},
						},
					},
					InitialDelaySeconds: int32(initialDelaySeconds),
					PeriodSeconds:       int32(periodSeconds),
					TimeoutSeconds:      int32(timeoutSeconds),
					SuccessThreshold:    int32(successThreshold),
					FailureThreshold:    int32(failureThreshold),
				},
			})
		}
	}
	return containers, volumes, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Skaffold 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 label

import (
	"encoding/json"
	"time"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/deploy"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/kubernetes"
	"github.com/sirupsen/logrus"

	clientgo "k8s.io/client-go/kubernetes"

	appsv1 "k8s.io/api/apps/v1"
	appsv1beta1 "k8s.io/api/apps/v1beta1"
	appsv1beta2 "k8s.io/api/apps/v1beta2"
	corev1 "k8s.io/api/core/v1"
	extensionsv1beta1 "k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"k8s.io/apimachinery/pkg/types"
	patch "k8s.io/apimachinery/pkg/util/strategicpatch"
)

// retry 3 times to give the object time to propagate to the API server
const tries int = 3
const sleeptime time.Duration = 300 * time.Millisecond

//nolint
func LabelDeployResults(labels map[string]string, results []deploy.Artifact) {
	// use the kubectl client to update all k8s objects with a skaffold watermark
	client, err := kubernetes.Client()
	if err != nil {
		logrus.Warnf("error retrieving kubernetes client: %s", err.Error())
		return
	}
	for _, res := range results {
		err = nil
		for i := 0; i < tries; i++ {
			if err = updateRuntimeObject(client, labels, res); err == nil {
				break
			}
			time.Sleep(sleeptime)
		}
		if err != nil {
			logrus.Warnf("error adding label to runtime object: %s", err.Error())
		}
	}
}

func addSkaffoldLabels(labels map[string]string, m *metav1.ObjectMeta) {
	if m.Labels == nil {
		m.Labels = map[string]string{}
	}
	for k, v := range labels {
		m.Labels[k] = v
	}
}

func retrieveNamespace(ns string, m metav1.ObjectMeta) string {
	if ns != "" {
		return ns
	}
	if m.Namespace != "" {
		return m.Namespace
	}
	return "default"
}

// TODO(nkubala): change this to use the client-go dynamic client or something equally clean
func updateRuntimeObject(client clientgo.Interface, labels map[string]string, res deploy.Artifact) error {
	for k, v := range constants.Labels.DefaultLabels {
		labels[k] = v
	}
	var err error
	obj := *res.Obj
	originalJSON, _ := json.Marshal(obj)
	modifiedObj := obj.DeepCopyObject()
	switch obj.(type) {
	case *corev1.Pod:
		apiObject := modifiedObj.(*corev1.Pod)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.CoreV1().Pods(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1.Deployment:
		apiObject := modifiedObj.(*appsv1.Deployment)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1().Deployments(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1beta1.Deployment:
		apiObject := modifiedObj.(*appsv1beta1.Deployment)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1beta1().Deployments(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1beta2.Deployment:
		apiObject := modifiedObj.(*appsv1beta2.Deployment)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1beta2().Deployments(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *extensionsv1beta1.Deployment:
		apiObject := modifiedObj.(*extensionsv1beta1.Deployment)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.ExtensionsV1beta1().Deployments(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *corev1.Service:
		apiObject := modifiedObj.(*corev1.Service)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.CoreV1().Services(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1.StatefulSet:
		apiObject := modifiedObj.(*appsv1.StatefulSet)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1().StatefulSets(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1beta1.StatefulSet:
		apiObject := modifiedObj.(*appsv1beta1.StatefulSet)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1beta1().StatefulSets(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1beta2.StatefulSet:
		apiObject := modifiedObj.(*appsv1beta2.StatefulSet)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1beta2().StatefulSets(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *extensionsv1beta1.DaemonSet:
		apiObject := modifiedObj.(*extensionsv1beta1.DaemonSet)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.ExtensionsV1beta1().DaemonSets(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1.ReplicaSet:
		apiObject := modifiedObj.(*appsv1.ReplicaSet)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1().ReplicaSets(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	case *appsv1beta2.ReplicaSet:
		apiObject := modifiedObj.(*appsv1beta2.ReplicaSet)
		addSkaffoldLabels(labels, &apiObject.ObjectMeta)
		modifiedJSON, _ := json.Marshal(modifiedObj)
		p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, apiObject)
		_, err = client.AppsV1beta2().ReplicaSets(retrieveNamespace(res.Namespace, apiObject.ObjectMeta)).Patch(apiObject.Name, types.StrategicMergePatchType, p)
	default:
		logrus.Infof("unknown runtime.Object, skipping label")
		return nil
	}
	return err
}
<commit_msg>refactor switch statement to remove boilerplate<commit_after>/*
Copyright 2018 The Skaffold 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 label

import (
	"encoding/json"
	"time"

	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/constants"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/deploy"
	"github.com/GoogleContainerTools/skaffold/pkg/skaffold/kubernetes"
	"github.com/sirupsen/logrus"

	clientgo "k8s.io/client-go/kubernetes"

	appsv1 "k8s.io/api/apps/v1"
	appsv1beta1 "k8s.io/api/apps/v1beta1"
	appsv1beta2 "k8s.io/api/apps/v1beta2"
	corev1 "k8s.io/api/core/v1"
	extensionsv1beta1 "k8s.io/api/extensions/v1beta1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	patch "k8s.io/apimachinery/pkg/util/strategicpatch"
)

type objectType int

// List of API Objects supported by the Skaffold Labeler
const (
	_ = iota
	corev1Pod
	appsv1Deployment
	appsv1Beta1Deployment
	appsv1Beta2Deployment
	extensionsv1Beta1Deployment
	corev1Service
	appv1StatefulSet
	appsv1Beta1StatefulSet
	appsv1Beta2StatefulSet
	extensionsv1Beta1DaemonSet
	appsv1ReplicaSet
	appsv1Beta2ReplicaSet
)

// converter is responsible for determining whether an object can be converted to a given type
type converter func(runtime.Object) bool

// patcher is responsible for applying a given patch to the provided object
type patcher func(clientgo.Interface, string, string, []byte) error

// objectMeta is responsible for returning a generic runtime.Object's metadata
type objectMeta func(runtime.Object) *metav1.ObjectMeta

var converters = map[objectType]converter{
	corev1Pod: func(r runtime.Object) bool {
		_, ok := r.(*corev1.Pod)
		return ok
	},
	appsv1Deployment: func(r runtime.Object) bool {
		_, ok := r.(*appsv1.Deployment)
		return ok
	},
	appsv1Beta1Deployment: func(r runtime.Object) bool {
		_, ok := r.(*appsv1beta1.Deployment)
		return ok
	},
	appsv1Beta2Deployment: func(r runtime.Object) bool {
		_, ok := r.(*appsv1beta2.Deployment)
		return ok
	},
	extensionsv1Beta1Deployment: func(r runtime.Object) bool {
		_, ok := r.(*extensionsv1beta1.Deployment)
		return ok
	},
	corev1Service: func(r runtime.Object) bool {
		_, ok := r.(*corev1.Service)
		return ok
	},
	appv1StatefulSet: func(r runtime.Object) bool {
		_, ok := r.(*appsv1.StatefulSet)
		return ok
	},
	appsv1Beta1StatefulSet: func(r runtime.Object) bool {
		_, ok := r.(*appsv1beta1.StatefulSet)
		return ok
	},
	appsv1Beta2StatefulSet: func(r runtime.Object) bool {
		_, ok := r.(*appsv1beta2.StatefulSet)
		return ok
	},
	extensionsv1Beta1DaemonSet: func(r runtime.Object) bool {
		_, ok := r.(*extensionsv1beta1.DaemonSet)
		return ok
	},
	appsv1ReplicaSet: func(r runtime.Object) bool {
		_, ok := r.(*appsv1.ReplicaSet)
		return ok
	},
	appsv1Beta2ReplicaSet: func(r runtime.Object) bool {
		_, ok := r.(*appsv1beta2.ReplicaSet)
		return ok
	},
}

var patchers = map[objectType]patcher{
	corev1Pod: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.CoreV1().Pods(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1Beta1Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1beta1().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1Beta2Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1beta2().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	extensionsv1Beta1Deployment: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.ExtensionsV1beta1().Deployments(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	corev1Service: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.CoreV1().Services(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appv1StatefulSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1().StatefulSets(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1Beta1StatefulSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1beta1().StatefulSets(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1Beta2StatefulSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1beta2().StatefulSets(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	extensionsv1Beta1DaemonSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.ExtensionsV1beta1().DaemonSets(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1ReplicaSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1().ReplicaSets(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
	appsv1Beta2ReplicaSet: func(client clientgo.Interface, ns string, name string, p []byte) error {
		_, err := client.AppsV1beta2().ReplicaSets(ns).Patch(name, types.StrategicMergePatchType, p)
		return err
	},
}

var objectMetas = map[objectType]objectMeta{
	corev1Pod: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*corev1.Pod).ObjectMeta)
	},
	appsv1Deployment: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1.Deployment).ObjectMeta)
	},
	appsv1Beta1Deployment: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1beta1.Deployment).ObjectMeta)
	},
	appsv1Beta2Deployment: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1beta2.Deployment).ObjectMeta)
	},
	extensionsv1Beta1Deployment: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*extensionsv1beta1.Deployment).ObjectMeta)
	},
	corev1Service: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*corev1.Service).ObjectMeta)
	},
	appv1StatefulSet: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1.StatefulSet).ObjectMeta)
	},
	appsv1Beta1StatefulSet: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1beta1.StatefulSet).ObjectMeta)
	},
	appsv1Beta2StatefulSet: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1beta1.StatefulSet).ObjectMeta)
	},
	extensionsv1Beta1DaemonSet: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*extensionsv1beta1.DaemonSet).ObjectMeta)
	},
	appsv1ReplicaSet: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1.ReplicaSet).ObjectMeta)
	},
	appsv1Beta2ReplicaSet: func(r runtime.Object) *metav1.ObjectMeta {
		return &(r.(*appsv1beta2.ReplicaSet).ObjectMeta)
	},
}

// retry 3 times to give the object time to propagate to the API server
const tries int = 3
const sleeptime time.Duration = 300 * time.Millisecond

//nolint
func LabelDeployResults(labels map[string]string, results []deploy.Artifact) {
	// use the kubectl client to update all k8s objects with a skaffold watermark
	client, err := kubernetes.Client()
	if err != nil {
		logrus.Warnf("error retrieving kubernetes client: %s", err.Error())
		return
	}
	for _, res := range results {
		err = nil
		for i := 0; i < tries; i++ {
			if err = updateRuntimeObject(client, labels, res); err == nil {
				break
			}
			time.Sleep(sleeptime)
		}
		if err != nil {
			logrus.Warnf("error adding label to runtime object: %s", err.Error())
		}
	}
}

func addSkaffoldLabels(labels map[string]string, m *metav1.ObjectMeta) {
	if m.Labels == nil {
		m.Labels = map[string]string{}
	}
	for k, v := range labels {
		m.Labels[k] = v
	}
}

func retrieveNamespace(ns string, m metav1.ObjectMeta) string {
	if ns != "" {
		return ns
	}
	if m.Namespace != "" {
		return m.Namespace
	}
	return "default"
}

// TODO(nkubala): change this to use the client-go dynamic client or something equally clean
func updateRuntimeObject(client clientgo.Interface, labels map[string]string, res deploy.Artifact) error {
	for k, v := range constants.Labels.DefaultLabels {
		labels[k] = v
	}
	var err error
	applied := false
	obj := *res.Obj
	originalJSON, _ := json.Marshal(obj)
	modifiedObj := obj.DeepCopyObject()
	for typeStr, c := range converters {
		if applied = c(modifiedObj); applied {
			metadata := objectMetas[typeStr](modifiedObj)
			addSkaffoldLabels(labels, metadata)
			modifiedJSON, _ := json.Marshal(modifiedObj)
			p, _ := patch.CreateTwoWayMergePatch(originalJSON, modifiedJSON, modifiedObj)
			err = patchers[typeStr](client, retrieveNamespace(res.Namespace, *metadata), metadata.GetName(), p)
			break // we should only ever apply one patch, so stop here
		}
	}
	if !applied {
		logrus.Infof("unknown runtime.Object, skipping label")
	}
	return err
}
<|endoftext|>


<commit_before>package donut

import (
	"encoding/json"
	"io/ioutil"
	"os"
	"path"
)

func newDonutObjectWriter(objectDir string) (Writer, error) {
	dataFile, err := os.OpenFile(path.Join(objectDir, "data"), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if err != nil {
		return nil, err
	}
	return donutObjectWriter{
		root:          objectDir,
		file:          dataFile,
		metadata:      make(map[string]string),
		donutMetadata: make(map[string]string),
	}, nil
}

type donutObjectWriter struct {
	root          string
	file          *os.File
	metadata      map[string]string
	donutMetadata map[string]string
	err           error
}

func (d donutObjectWriter) Write(data []byte) (int, error) {
	return d.file.Write(data)
}

func (d donutObjectWriter) Close() error {
	if d.err != nil {
		return d.err
	}
	metadata, _ := json.Marshal(d.metadata)
	ioutil.WriteFile(path.Join(d.root, "metadata.json"), metadata, 0600)
	donutMetadata, _ := json.Marshal(d.donutMetadata)
	ioutil.WriteFile(path.Join(d.root, "donutMetadata.json"), donutMetadata, 0600)

	return d.file.Close()
}

func (d donutObjectWriter) CloseWithError(err error) error {
	if d.err != nil {
		d.err = err
	}
	return d.Close()
}

func (d donutObjectWriter) SetMetadata(metadata map[string]string) error {
	for k := range d.metadata {
		delete(d.metadata, k)
	}
	for k, v := range metadata {
		d.metadata[k] = v
	}
	return nil
}

func (d donutObjectWriter) GetMetadata() (map[string]string, error) {
	metadata := make(map[string]string)
	for k, v := range d.metadata {
		metadata[k] = v
	}
	return metadata, nil
}

func (d donutObjectWriter) SetDonutMetadata(metadata map[string]string) error {
	for k := range d.donutMetadata {
		delete(d.donutMetadata, k)
	}
	for k, v := range metadata {
		d.donutMetadata[k] = v
	}
	return nil
}

func (d donutObjectWriter) GetDonutMetadata() (map[string]string, error) {
	donutMetadata := make(map[string]string)
	for k, v := range d.donutMetadata {
		donutMetadata[k] = v
	}
	return donutMetadata, nil
}
<commit_msg>Adding iodine to object_writer<commit_after>package donut

import (
	"encoding/json"
	"github.com/minio-io/iodine"
	"io/ioutil"
	"os"
	"path"
)

func newDonutObjectWriter(objectDir string) (Writer, error) {
	dataFile, err := os.OpenFile(path.Join(objectDir, "data"), os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0600)
	if err != nil {
		return nil, iodine.Error(err, map[string]string{"objectDir": objectDir})
	}
	return donutObjectWriter{
		root:          objectDir,
		file:          dataFile,
		metadata:      make(map[string]string),
		donutMetadata: make(map[string]string),
	}, nil
}

type donutObjectWriter struct {
	root          string
	file          *os.File
	metadata      map[string]string
	donutMetadata map[string]string
	err           error
}

func (d donutObjectWriter) Write(data []byte) (int, error) {
	written, err := d.file.Write(data)
	iodine.Error(err, nil)
	return written, err
}

func (d donutObjectWriter) Close() error {
	if d.err != nil {
		return iodine.Error(d.err, nil)
	}
	metadata, _ := json.Marshal(d.metadata)
	ioutil.WriteFile(path.Join(d.root, "metadata.json"), metadata, 0600)
	donutMetadata, _ := json.Marshal(d.donutMetadata)
	ioutil.WriteFile(path.Join(d.root, "donutMetadata.json"), donutMetadata, 0600)

	return iodine.Error(d.file.Close(), nil)
}

func (d donutObjectWriter) CloseWithError(err error) error {
	if d.err != nil {
		d.err = err
	}
	// TODO semantics between this and d.Close are weird, work out something better
	return iodine.Error(d.Close(), nil)
}

func (d donutObjectWriter) SetMetadata(metadata map[string]string) error {
	for k := range d.metadata {
		delete(d.metadata, k)
	}
	for k, v := range metadata {
		d.metadata[k] = v
	}
	return nil
}

func (d donutObjectWriter) GetMetadata() (map[string]string, error) {
	metadata := make(map[string]string)
	for k, v := range d.metadata {
		metadata[k] = v
	}
	return metadata, nil
}

func (d donutObjectWriter) SetDonutMetadata(metadata map[string]string) error {
	for k := range d.donutMetadata {
		delete(d.donutMetadata, k)
	}
	for k, v := range metadata {
		d.donutMetadata[k] = v
	}
	return nil
}

func (d donutObjectWriter) GetDonutMetadata() (map[string]string, error) {
	donutMetadata := make(map[string]string)
	for k, v := range d.donutMetadata {
		donutMetadata[k] = v
	}
	return donutMetadata, nil
}
<|endoftext|>


<commit_before>// Copyright 2016 The Gosl 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 gm

import (
	"github.com/cpmech/gosl/chk"
	"github.com/cpmech/gosl/fun"
	"github.com/cpmech/gosl/la"
	"github.com/cpmech/gosl/plt"
)

// Transfinite maps a reference square [-1,1]×[-1,1] into a curve-bounded quadrilateral
//
//                                   Γ1(η) _,B
//       c———————b                C--.___-'  |
//       |       |                 \         | Γ0(ξ)
//   η   |       |      y     Γ2(ξ) \        |
//   ↑   |       |      ↑           /      __A
//   |   d———————a      |          / ____,'
//   |                  |         D-'   Γ3(η)
//   +————→ξ            +————→x
//
type Transfinite struct {
	Ndim int         // space dimension
	Γ    []fun.Vs    // the boundary functions
	Γd   []fun.Vs    // derivatives of boundary functions
	C    []la.Vector // "corner" points
	X    []la.Vector // points at arbitrary positions along edges/faces
	Xd   []la.Vector // derivatives at arbitrary positions along edges/faces
}

// NewTransfinite allocates a new structure
//  Γ  -- boundary functions x(s) = Γ(s)
//  Γd -- derivative functions dxds(s) = Γ'(s)
func NewTransfinite(ndim int, Γ, Γd []fun.Vs) (o *Transfinite) {
	o = new(Transfinite)
	o.Ndim = ndim
	o.Γ = Γ
	o.Γd = Γd
	if o.Ndim == 2 {
		if len(Γ) != 4 || len(Γd) != 4 {
			chk.Panic("in 2D, four boundary functions Γ are required\n")
		}
		o.C = make([]la.Vector, 4)
		o.X = make([]la.Vector, 4)
		o.Xd = make([]la.Vector, 4)
		for i := 0; i < len(o.C); i++ {
			o.C[i] = la.NewVector(o.Ndim)
			o.X[i] = la.NewVector(o.Ndim)
			o.Xd[i] = la.NewVector(o.Ndim)
		}
		o.Γ[0](o.C[0], -1)
		o.Γ[0](o.C[1], +1)
		o.Γ[2](o.C[2], +1)
		o.Γ[2](o.C[3], -1)
	} else if o.Ndim == 3 {
		if len(Γ) != 6 {
			chk.Panic("in 3D, six boundary functions Γ are required\n")
		}
	} else {
		chk.Panic("space dimension (ndim) must be 2 or 3\n")
	}
	return
}

// Point computes "real" position x(ξ,η)
//  Input:
//    r -- the "reference" coordinates {ξ,η}
//  Output:
//    x -- the "real" coordinates {x,y}
func (o *Transfinite) Point(x, r la.Vector) {
	if o.Ndim == 2 {
		ξ, η := r[0], r[1]
		A, B, C, D := o.X[0], o.X[1], o.X[2], o.X[3]
		m, n, p, q := o.C[0], o.C[1], o.C[2], o.C[3]
		o.Γ[0](A, ξ)
		o.Γ[1](B, η)
		o.Γ[2](C, ξ)
		o.Γ[3](D, η)
		for i := 0; i < o.Ndim; i++ {
			x[i] = 0.5*((1-η)*A[i]+(1+ξ)*B[i]+(1+η)*C[i]+(1-ξ)*D[i]) -
				0.25*((1-ξ)*((1-η)*m[i]+(1+η)*q[i])+(1+ξ)*((1-η)*n[i]+(1+η)*p[i]))
		}
		return
	}
	chk.Panic("Point function is not ready for 3D yet\n")
}

// Derivs calculates derivatives (=metric terms) @ r={ξ,η}
//  Input:
//    r -- the "reference" coordinates {ξ,η}
//  Output:
//    dxdr -- the derivatives [dx/dr]ij = dxi/drj
//    x    -- the "real" coordinates {x,y}
func (o *Transfinite) Derivs(dxdr *la.Matrix, x, r la.Vector) {
	if o.Ndim == 2 {
		ξ, η := r[0], r[1]
		A, B, C, D := o.X[0], o.X[1], o.X[2], o.X[3]
		a, b, c, d := o.Xd[0], o.Xd[1], o.Xd[2], o.Xd[3]
		m, n, p, q := o.C[0], o.C[1], o.C[2], o.C[3]
		o.Γ[0](A, ξ)
		o.Γ[1](B, η)
		o.Γ[2](C, ξ)
		o.Γ[3](D, η)
		o.Γd[0](a, ξ)
		o.Γd[1](b, η)
		o.Γd[2](c, ξ)
		o.Γd[3](d, η)
		var dxidξ, dxidη float64
		for i := 0; i < o.Ndim; i++ {

			x[i] = 0.5*((1-η)*A[i]+(1+ξ)*B[i]+(1+η)*C[i]+(1-ξ)*D[i]) -
				0.25*((1-ξ)*((1-η)*m[i]+(1+η)*q[i])+(1+ξ)*((1-η)*n[i]+(1+η)*p[i]))

			dxidξ = 0.5*((1-η)*a[i]+B[i]+(1+η)*c[i]-D[i]) -
				0.25*((1-η)*(n[i]-m[i])+(1+η)*(p[i]-q[i]))

			dxidη = 0.5*(-A[i]+(1+ξ)*b[i]+C[i]+(1-ξ)*d[i]) -
				0.25*((1-ξ)*(q[i]-m[i])+(1+ξ)*(p[i]-n[i]))

			dxdr.Set(i, 0, dxidξ)
			dxdr.Set(i, 1, dxidη)
		}
		return
	}
	chk.Panic("Derivs function is not ready for 3D yet\n")
}

// Draw draws figure formed by Γ
func (o *Transfinite) Draw(npts []int, args, argsBry *plt.A) {

	// auxiliary
	if len(npts) != o.Ndim {
		npts = make([]int, o.Ndim)
	}
	for i := 0; i < o.Ndim; i++ {
		if npts[i] < 3 {
			npts[i] = 3
		}
	}
	if args == nil {
		args = &plt.A{C: plt.C(0, 0), NoClip: true}
	}
	if argsBry == nil {
		argsBry = &plt.A{C: plt.C(0, 0), Lw: 2, NoClip: true}
	}
	x := la.NewVector(o.Ndim)
	r := la.NewVector(o.Ndim)

	// draw 0-lines
	x0 := make([]float64, npts[0])
	y0 := make([]float64, npts[0])
	for j := 0; j < npts[1]; j++ {
		r[1] = -1 + 2*float64(j)/float64(npts[1]-1)
		for i := 0; i < npts[0]; i++ {
			r[0] = -1 + 2*float64(i)/float64(npts[0]-1)
			o.Point(x, r)
			x0[i] = x[0]
			y0[i] = x[1]
		}
		plt.Plot(x0, y0, args)
	}

	// draw 1-lines
	x1 := make([]float64, npts[1])
	y1 := make([]float64, npts[1])
	for i := 0; i < npts[0]; i++ {
		r[0] = -1 + 2*float64(i)/float64(npts[0]-1)
		for j := 0; j < npts[1]; j++ {
			r[1] = -1 + 2*float64(j)/float64(npts[1]-1)
			o.Point(x, r)
			x1[j] = x[0]
			y1[j] = x[1]
		}
		plt.Plot(x1, y1, args)
	}

	// draw Γ0(ξ)
	for i := 0; i < npts[0]; i++ {
		ξ := -1 + 2*float64(i)/float64(npts[0]-1)
		o.Γ[0](x, ξ)
		x0[i] = x[0]
		y0[i] = x[1]
	}
	plt.Plot(x0, y0, argsBry)

	// draw Γ1(η)
	for j := 0; j < npts[1]; j++ {
		η := -1 + 2*float64(j)/float64(npts[1]-1)
		o.Γ[1](x, η)
		x1[j] = x[0]
		y1[j] = x[1]
	}
	plt.Plot(x1, y1, argsBry)

	// draw Γ2(ξ)
	for i := 0; i < npts[0]; i++ {
		ξ := -1 + 2*float64(i)/float64(npts[0]-1)
		o.Γ[2](x, ξ)
		x0[i] = x[0]
		y0[i] = x[1]
	}
	plt.Plot(x0, y0, argsBry)

	// draw Γ3(η)
	for j := 0; j < npts[1]; j++ {
		η := -1 + 2*float64(j)/float64(npts[1]-1)
		o.Γ[3](x, η)
		x1[j] = x[0]
		y1[j] = x[1]
	}
	plt.Plot(x1, y1, argsBry)
}
<commit_msg>Rename greek variables to r,s,t<commit_after>// Copyright 2016 The Gosl 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 gm

import (
	"github.com/cpmech/gosl/chk"
	"github.com/cpmech/gosl/fun"
	"github.com/cpmech/gosl/la"
	"github.com/cpmech/gosl/plt"
)

// Transfinite maps a reference square [-1,1]×[-1,1] into a curve-bounded quadrilateral
//
//                                   B1(s) _,B
//       c———————b                C--.___-'  |
//       |       |                 \         | B0(r)
//   s   |       |      y     B2(r) \        |
//   ↑   |       |      ↑           /      __A
//   |   d———————a      |          / ____,'
//   |                  |         D-'   B3(s)
//   +————→r            +————→x
//
type Transfinite struct {
	Ndim int         // space dimension
	B    []fun.Vs    // the boundary functions
	Bd   []fun.Vs    // derivatives of boundary functions
	C    []la.Vector // "corner" points
	X    []la.Vector // points at arbitrary positions along edges/faces
	Xd   []la.Vector // derivatives at arbitrary positions along edges/faces
}

// NewTransfinite allocates a new structure
//  B  -- boundary functions x(s) = B(s)
//  Bd -- derivative functions dxds(s) = B'(s)
func NewTransfinite(ndim int, B, Bd []fun.Vs) (o *Transfinite) {
	o = new(Transfinite)
	o.Ndim = ndim
	o.B = B
	o.Bd = Bd
	if o.Ndim == 2 {
		if len(B) != 4 || len(Bd) != 4 {
			chk.Panic("in 2D, four boundary functions B are required\n")
		}
		o.C = make([]la.Vector, 4)
		o.X = make([]la.Vector, 4)
		o.Xd = make([]la.Vector, 4)
		for i := 0; i < len(o.C); i++ {
			o.C[i] = la.NewVector(o.Ndim)
			o.X[i] = la.NewVector(o.Ndim)
			o.Xd[i] = la.NewVector(o.Ndim)
		}
		o.B[0](o.C[0], -1)
		o.B[0](o.C[1], +1)
		o.B[2](o.C[2], +1)
		o.B[2](o.C[3], -1)
	} else if o.Ndim == 3 {
		if len(B) != 6 {
			chk.Panic("in 3D, six boundary functions B are required\n")
		}
	} else {
		chk.Panic("space dimension (ndim) must be 2 or 3\n")
	}
	return
}

// Point computes "real" position x(r,s,t)
//  Input:
//    u -- the "reference" coordinates {r,s,t} ϵ [-1,+1]×[-1,+1]×[-1,+1]
//  Output:
//    x -- the "real" coordinates {x,y,z}
func (o *Transfinite) Point(x, u la.Vector) {
	if o.Ndim == 2 {
		r, s := u[0], u[1]
		A, B, C, D := o.X[0], o.X[1], o.X[2], o.X[3]
		m, n, p, q := o.C[0], o.C[1], o.C[2], o.C[3]
		o.B[0](A, r)
		o.B[1](B, s)
		o.B[2](C, r)
		o.B[3](D, s)
		for i := 0; i < o.Ndim; i++ {
			x[i] = 0.5*((1-s)*A[i]+(1+r)*B[i]+(1+s)*C[i]+(1-r)*D[i]) -
				0.25*((1-r)*((1-s)*m[i]+(1+s)*q[i])+(1+r)*((1-s)*n[i]+(1+s)*p[i]))
		}
		return
	}
	chk.Panic("Point function is not ready for 3D yet\n")
}

// Derivs calculates derivatives (=metric terms) @ u={r,s,t}
//  Input:
//    u -- the "reference" coordinates {r,s,t} ϵ [-1,+1]×[-1,+1]×[-1,+1]
//  Output:
//    dxdu -- the derivatives [dx/du]ij = dxi/duj
//    x    -- the "real" coordinates {x,y,z}
func (o *Transfinite) Derivs(dxdu *la.Matrix, x, u la.Vector) {
	if o.Ndim == 2 {
		r, s := u[0], u[1]
		A, B, C, D := o.X[0], o.X[1], o.X[2], o.X[3]
		a, b, c, d := o.Xd[0], o.Xd[1], o.Xd[2], o.Xd[3]
		m, n, p, q := o.C[0], o.C[1], o.C[2], o.C[3]
		o.B[0](A, r)
		o.B[1](B, s)
		o.B[2](C, r)
		o.B[3](D, s)
		o.Bd[0](a, r)
		o.Bd[1](b, s)
		o.Bd[2](c, r)
		o.Bd[3](d, s)
		var dxidr, dxids float64
		for i := 0; i < o.Ndim; i++ {

			x[i] = 0.5*((1-s)*A[i]+(1+r)*B[i]+(1+s)*C[i]+(1-r)*D[i]) -
				0.25*((1-r)*((1-s)*m[i]+(1+s)*q[i])+(1+r)*((1-s)*n[i]+(1+s)*p[i]))

			dxidr = 0.5*((1-s)*a[i]+B[i]+(1+s)*c[i]-D[i]) -
				0.25*((1-s)*(n[i]-m[i])+(1+s)*(p[i]-q[i]))

			dxids = 0.5*(-A[i]+(1+r)*b[i]+C[i]+(1-r)*d[i]) -
				0.25*((1-r)*(q[i]-m[i])+(1+r)*(p[i]-n[i]))

			dxdu.Set(i, 0, dxidr)
			dxdu.Set(i, 1, dxids)
		}
		return
	}
	chk.Panic("Derivs function is not ready for 3D yet\n")
}

// Draw draws figure formed by B
func (o *Transfinite) Draw(npts []int, args, argsBry *plt.A) {

	// auxiliary
	if len(npts) != o.Ndim {
		npts = make([]int, o.Ndim)
	}
	for i := 0; i < o.Ndim; i++ {
		if npts[i] < 3 {
			npts[i] = 3
		}
	}
	if args == nil {
		args = &plt.A{C: plt.C(0, 0), NoClip: true}
	}
	if argsBry == nil {
		argsBry = &plt.A{C: plt.C(0, 0), Lw: 2, NoClip: true}
	}
	x := la.NewVector(o.Ndim)
	u := la.NewVector(o.Ndim)

	// draw 0-lines
	x0 := make([]float64, npts[0])
	y0 := make([]float64, npts[0])
	for j := 0; j < npts[1]; j++ {
		u[1] = -1 + 2*float64(j)/float64(npts[1]-1)
		for i := 0; i < npts[0]; i++ {
			u[0] = -1 + 2*float64(i)/float64(npts[0]-1)
			o.Point(x, u)
			x0[i] = x[0]
			y0[i] = x[1]
		}
		plt.Plot(x0, y0, args)
	}

	// draw 1-lines
	x1 := make([]float64, npts[1])
	y1 := make([]float64, npts[1])
	for i := 0; i < npts[0]; i++ {
		u[0] = -1 + 2*float64(i)/float64(npts[0]-1)
		for j := 0; j < npts[1]; j++ {
			u[1] = -1 + 2*float64(j)/float64(npts[1]-1)
			o.Point(x, u)
			x1[j] = x[0]
			y1[j] = x[1]
		}
		plt.Plot(x1, y1, args)
	}

	// draw B0(r)
	for i := 0; i < npts[0]; i++ {
		r := -1 + 2*float64(i)/float64(npts[0]-1)
		o.B[0](x, r)
		x0[i] = x[0]
		y0[i] = x[1]
	}
	plt.Plot(x0, y0, argsBry)

	// draw B1(s)
	for j := 0; j < npts[1]; j++ {
		s := -1 + 2*float64(j)/float64(npts[1]-1)
		o.B[1](x, s)
		x1[j] = x[0]
		y1[j] = x[1]
	}
	plt.Plot(x1, y1, argsBry)

	// draw B2(r)
	for i := 0; i < npts[0]; i++ {
		r := -1 + 2*float64(i)/float64(npts[0]-1)
		o.B[2](x, r)
		x0[i] = x[0]
		y0[i] = x[1]
	}
	plt.Plot(x0, y0, argsBry)

	// draw B3(s)
	for j := 0; j < npts[1]; j++ {
		s := -1 + 2*float64(j)/float64(npts[1]-1)
		o.B[3](x, s)
		x1[j] = x[0]
		y1[j] = x[1]
	}
	plt.Plot(x1, y1, argsBry)
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package gnmi

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"path"

	pb "github.com/openconfig/gnmi/proto/gnmi"
	"google.golang.org/grpc/codes"
)

// Get sents a GetRequest to the given client.
func Get(ctx context.Context, client pb.GNMIClient, paths [][]string) error {
	req, err := NewGetRequest(paths)
	if err != nil {
		return err
	}
	resp, err := client.Get(ctx, req)
	if err != nil {
		return err
	}
	for _, notif := range resp.Notification {
		for _, update := range notif.Update {
			fmt.Printf("%s:\n", StrPath(update.Path))
			fmt.Println(strUpdateVal(update))
		}
	}
	return nil
}

// Capabilities retuns the capabilities of the client.
func Capabilities(ctx context.Context, client pb.GNMIClient) error {
	resp, err := client.Capabilities(ctx, &pb.CapabilityRequest{})
	if err != nil {
		return err
	}
	fmt.Printf("Version: %s\n", resp.GNMIVersion)
	for _, mod := range resp.SupportedModels {
		fmt.Printf("SupportedModel: %s\n", mod)
	}
	for _, enc := range resp.SupportedEncodings {
		fmt.Printf("SupportedEncoding: %s\n", enc)
	}
	return nil
}

// val may be a path to a file or it may be json. First see if it is a
// file, if so return its contents, otherwise return val
func extractJSON(val string) []byte {
	jsonBytes, err := ioutil.ReadFile(val)
	if err != nil {
		jsonBytes = []byte(val)
	}
	return jsonBytes
}

// strUpdateVal will return a string representing the value within the supplied update
func strUpdateVal(u *pb.Update) string {
	if u.Value != nil {
		return string(u.Value.Value) // Backwards compatibility with pre-v0.4 gnmi
	}
	return strVal(u.Val)
}

// strVal will return a string representing the supplied value
func strVal(val *pb.TypedValue) string {
	switch v := val.GetValue().(type) {
	case *pb.TypedValue_StringVal:
		return v.StringVal
	case *pb.TypedValue_JsonIetfVal:
		return string(v.JsonIetfVal)
	case *pb.TypedValue_JsonVal:
		return string(v.JsonVal)
	case *pb.TypedValue_IntVal:
		return fmt.Sprintf("%v", v.IntVal)
	case *pb.TypedValue_UintVal:
		return fmt.Sprintf("%v", v.UintVal)
	case *pb.TypedValue_BoolVal:
		return fmt.Sprintf("%v", v.BoolVal)
	case *pb.TypedValue_BytesVal:
		return string(v.BytesVal)
	case *pb.TypedValue_DecimalVal:
		return strDecimal64(v.DecimalVal)
	case *pb.TypedValue_FloatVal:
		return fmt.Sprintf("%v", v.FloatVal)
	case *pb.TypedValue_LeaflistVal:
		return strLeaflist(v.LeaflistVal)
	case *pb.TypedValue_AsciiVal:
		return v.AsciiVal
	case *pb.TypedValue_AnyVal:
		return v.AnyVal.String()
	default:
		panic(v)
	}
}

func strDecimal64(d *pb.Decimal64) string {
	var i, frac uint64
	if d.Precision > 0 {
		div := uint64(10)
		it := d.Precision - 1
		for it > 0 {
			div *= 10
			it--
		}
		i = d.Digits / div
		frac = d.Digits % div
	} else {
		i = d.Digits
	}
	return fmt.Sprintf("%d.%d", i, frac)
}

// strLeafList builds a human-readable form of a leaf-list. e.g. [1,2,3] or [a,b,c]
func strLeaflist(v *pb.ScalarArray) string {
	s := make([]string, 0, len(v.Element))
	sz := 2 // []

	// convert arbitrary TypedValues to string form
	for _, elm := range v.Element {
		str := strVal(elm)
		s = append(s, str)
		sz += len(str) + 1 // %v + ,
	}

	b := make([]byte, sz)
	buf := bytes.NewBuffer(b)

	buf.WriteRune('[')
	for i := range v.Element {
		buf.WriteString(s[i])
		if i < len(v.Element)-1 {
			buf.WriteRune(',')
		}
	}
	buf.WriteRune(']')
	return buf.String()
}

func update(p *pb.Path, val string) *pb.Update {
	var v *pb.TypedValue
	switch p.Origin {
	case "":
		v = &pb.TypedValue{
			Value: &pb.TypedValue_JsonIetfVal{JsonIetfVal: extractJSON(val)}}
	case "cli":
		v = &pb.TypedValue{
			Value: &pb.TypedValue_AsciiVal{AsciiVal: val}}
	default:
		panic(fmt.Errorf("unexpected origin: %q", p.Origin))
	}

	return &pb.Update{Path: p, Val: v}
}

// Operation describes an gNMI operation.
type Operation struct {
	Type string
	Path []string
	Val  string
}

func newSetRequest(setOps []*Operation) (*pb.SetRequest, error) {
	req := &pb.SetRequest{}
	for _, op := range setOps {
		p, err := ParseGNMIElements(op.Path)
		if err != nil {
			return nil, err
		}

		switch op.Type {
		case "delete":
			req.Delete = append(req.Delete, p)
		case "update":
			req.Update = append(req.Update, update(p, op.Val))
		case "replace":
			req.Replace = append(req.Replace, update(p, op.Val))
		}
	}
	return req, nil
}

// Set sends a SetRequest to the given client.
func Set(ctx context.Context, client pb.GNMIClient, setOps []*Operation) error {
	req, err := newSetRequest(setOps)
	if err != nil {
		return err
	}
	resp, err := client.Set(ctx, req)
	if err != nil {
		return err
	}
	if resp.Message != nil && codes.Code(resp.Message.Code) != codes.OK {
		return errors.New(resp.Message.Message)
	}
	// TODO: Iterate over SetResponse.Response for more detailed error message?

	return nil
}

// Subscribe sends a SubscribeRequest to the given client.
func Subscribe(ctx context.Context, client pb.GNMIClient, paths [][]string,
	respChan chan<- *pb.SubscribeResponse, errChan chan<- error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	stream, err := client.Subscribe(ctx)
	if err != nil {
		errChan <- err
		return
	}
	req, err := NewSubscribeRequest(paths)
	if err != nil {
		errChan <- err
		return
	}
	if err := stream.Send(req); err != nil {
		errChan <- err
		return
	}

	for {
		resp, err := stream.Recv()
		if err != nil {
			if err == io.EOF {
				return
			}
			errChan <- err
			return
		}
		respChan <- resp
	}
}

// LogSubscribeResponse logs update responses to stderr.
func LogSubscribeResponse(response *pb.SubscribeResponse) error {
	switch resp := response.Response.(type) {
	case *pb.SubscribeResponse_Error:
		return errors.New(resp.Error.Message)
	case *pb.SubscribeResponse_SyncResponse:
		if !resp.SyncResponse {
			return errors.New("initial sync failed")
		}
	case *pb.SubscribeResponse_Update:
		prefix := StrPath(resp.Update.Prefix)
		for _, update := range resp.Update.Update {
			fmt.Printf("%s = %s\n", path.Join(prefix, StrPath(update.Path)),
				strUpdateVal(update))
		}
	}
	return nil
}
<commit_msg>gnmi: Pretty-print JSON returned from Get<commit_after>// Copyright (c) 2017 Arista Networks, Inc.
// Use of this source code is governed by the Apache License 2.0
// that can be found in the COPYING file.

package gnmi

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"path"

	pb "github.com/openconfig/gnmi/proto/gnmi"
	"google.golang.org/grpc/codes"
)

// Get sents a GetRequest to the given client.
func Get(ctx context.Context, client pb.GNMIClient, paths [][]string) error {
	req, err := NewGetRequest(paths)
	if err != nil {
		return err
	}
	resp, err := client.Get(ctx, req)
	if err != nil {
		return err
	}
	for _, notif := range resp.Notification {
		for _, update := range notif.Update {
			fmt.Printf("%s:\n", StrPath(update.Path))
			fmt.Println(strUpdateVal(update))
		}
	}
	return nil
}

// Capabilities retuns the capabilities of the client.
func Capabilities(ctx context.Context, client pb.GNMIClient) error {
	resp, err := client.Capabilities(ctx, &pb.CapabilityRequest{})
	if err != nil {
		return err
	}
	fmt.Printf("Version: %s\n", resp.GNMIVersion)
	for _, mod := range resp.SupportedModels {
		fmt.Printf("SupportedModel: %s\n", mod)
	}
	for _, enc := range resp.SupportedEncodings {
		fmt.Printf("SupportedEncoding: %s\n", enc)
	}
	return nil
}

// val may be a path to a file or it may be json. First see if it is a
// file, if so return its contents, otherwise return val
func extractJSON(val string) []byte {
	jsonBytes, err := ioutil.ReadFile(val)
	if err != nil {
		jsonBytes = []byte(val)
	}
	return jsonBytes
}

// strUpdateVal will return a string representing the value within the supplied update
func strUpdateVal(u *pb.Update) string {
	if u.Value != nil {
		return string(u.Value.Value) // Backwards compatibility with pre-v0.4 gnmi
	}
	return strVal(u.Val)
}

// strVal will return a string representing the supplied value
func strVal(val *pb.TypedValue) string {
	switch v := val.GetValue().(type) {
	case *pb.TypedValue_StringVal:
		return v.StringVal
	case *pb.TypedValue_JsonIetfVal:
		return strJSON(v.JsonIetfVal)
	case *pb.TypedValue_JsonVal:
		return strJSON(v.JsonVal)
	case *pb.TypedValue_IntVal:
		return fmt.Sprintf("%v", v.IntVal)
	case *pb.TypedValue_UintVal:
		return fmt.Sprintf("%v", v.UintVal)
	case *pb.TypedValue_BoolVal:
		return fmt.Sprintf("%v", v.BoolVal)
	case *pb.TypedValue_BytesVal:
		return string(v.BytesVal)
	case *pb.TypedValue_DecimalVal:
		return strDecimal64(v.DecimalVal)
	case *pb.TypedValue_FloatVal:
		return fmt.Sprintf("%v", v.FloatVal)
	case *pb.TypedValue_LeaflistVal:
		return strLeaflist(v.LeaflistVal)
	case *pb.TypedValue_AsciiVal:
		return v.AsciiVal
	case *pb.TypedValue_AnyVal:
		return v.AnyVal.String()
	default:
		panic(v)
	}
}

func strJSON(inJSON []byte) string {
	var out bytes.Buffer
	err := json.Indent(&out, inJSON, "", "  ")
	if err != nil {
		return fmt.Sprintf("(error unmarshalling json: %s)\n", err) + string(inJSON)
	}
	return out.String()
}

func strDecimal64(d *pb.Decimal64) string {
	var i, frac uint64
	if d.Precision > 0 {
		div := uint64(10)
		it := d.Precision - 1
		for it > 0 {
			div *= 10
			it--
		}
		i = d.Digits / div
		frac = d.Digits % div
	} else {
		i = d.Digits
	}
	return fmt.Sprintf("%d.%d", i, frac)
}

// strLeafList builds a human-readable form of a leaf-list. e.g. [1,2,3] or [a,b,c]
func strLeaflist(v *pb.ScalarArray) string {
	s := make([]string, 0, len(v.Element))
	sz := 2 // []

	// convert arbitrary TypedValues to string form
	for _, elm := range v.Element {
		str := strVal(elm)
		s = append(s, str)
		sz += len(str) + 1 // %v + ,
	}

	b := make([]byte, sz)
	buf := bytes.NewBuffer(b)

	buf.WriteRune('[')
	for i := range v.Element {
		buf.WriteString(s[i])
		if i < len(v.Element)-1 {
			buf.WriteRune(',')
		}
	}
	buf.WriteRune(']')
	return buf.String()
}

func update(p *pb.Path, val string) *pb.Update {
	var v *pb.TypedValue
	switch p.Origin {
	case "":
		v = &pb.TypedValue{
			Value: &pb.TypedValue_JsonIetfVal{JsonIetfVal: extractJSON(val)}}
	case "cli":
		v = &pb.TypedValue{
			Value: &pb.TypedValue_AsciiVal{AsciiVal: val}}
	default:
		panic(fmt.Errorf("unexpected origin: %q", p.Origin))
	}

	return &pb.Update{Path: p, Val: v}
}

// Operation describes an gNMI operation.
type Operation struct {
	Type string
	Path []string
	Val  string
}

func newSetRequest(setOps []*Operation) (*pb.SetRequest, error) {
	req := &pb.SetRequest{}
	for _, op := range setOps {
		p, err := ParseGNMIElements(op.Path)
		if err != nil {
			return nil, err
		}

		switch op.Type {
		case "delete":
			req.Delete = append(req.Delete, p)
		case "update":
			req.Update = append(req.Update, update(p, op.Val))
		case "replace":
			req.Replace = append(req.Replace, update(p, op.Val))
		}
	}
	return req, nil
}

// Set sends a SetRequest to the given client.
func Set(ctx context.Context, client pb.GNMIClient, setOps []*Operation) error {
	req, err := newSetRequest(setOps)
	if err != nil {
		return err
	}
	resp, err := client.Set(ctx, req)
	if err != nil {
		return err
	}
	if resp.Message != nil && codes.Code(resp.Message.Code) != codes.OK {
		return errors.New(resp.Message.Message)
	}
	// TODO: Iterate over SetResponse.Response for more detailed error message?

	return nil
}

// Subscribe sends a SubscribeRequest to the given client.
func Subscribe(ctx context.Context, client pb.GNMIClient, paths [][]string,
	respChan chan<- *pb.SubscribeResponse, errChan chan<- error) {
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()
	stream, err := client.Subscribe(ctx)
	if err != nil {
		errChan <- err
		return
	}
	req, err := NewSubscribeRequest(paths)
	if err != nil {
		errChan <- err
		return
	}
	if err := stream.Send(req); err != nil {
		errChan <- err
		return
	}

	for {
		resp, err := stream.Recv()
		if err != nil {
			if err == io.EOF {
				return
			}
			errChan <- err
			return
		}
		respChan <- resp
	}
}

// LogSubscribeResponse logs update responses to stderr.
func LogSubscribeResponse(response *pb.SubscribeResponse) error {
	switch resp := response.Response.(type) {
	case *pb.SubscribeResponse_Error:
		return errors.New(resp.Error.Message)
	case *pb.SubscribeResponse_SyncResponse:
		if !resp.SyncResponse {
			return errors.New("initial sync failed")
		}
	case *pb.SubscribeResponse_Update:
		prefix := StrPath(resp.Update.Prefix)
		for _, update := range resp.Update.Update {
			fmt.Printf("%s = %s\n", path.Join(prefix, StrPath(update.Path)),
				strUpdateVal(update))
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"github.com/keybase/client/go/citogo/types"
)

type opts struct {
	Flakes               int
	Fails                int
	Prefix               string
	S3Bucket             string
	ReportLambdaFunction string
	DirBasename          string
	BuildID              string
	Branch               string
	Preserve             bool
	BuildURL             string
	NoCompile            bool
	TestBinary           string
	Timeout              string
	Pause                time.Duration
}

func logError(f string, args ...interface{}) {
	s := fmt.Sprintf(f, args...)
	if s[len(s)-1] != '\n' {
		s += "\n"
	}
	fmt.Fprintf(os.Stderr, "%s", s)
}

type runner struct {
	opts   opts
	flakes []string
	fails  []string
	tests  []string
}

func convertBreakingChars(s string) string {
	// replace either the unix or the DOS directory marker
	// with an underscore, so as not to break the directory
	// structure of where we are storing the log
	s = strings.ReplaceAll(s, "/", "_")
	s = strings.ReplaceAll(s, "\\", "_")
	s = strings.ReplaceAll(s, "-", "_")
	return s
}

func (r *runner) parseArgs() (err error) {
	flag.IntVar(&r.opts.Flakes, "flakes", 3, "number of allowed flakes")
	flag.IntVar(&r.opts.Fails, "fails", -1, "number of fails allowed before quitting")
	flag.StringVar(&r.opts.Prefix, "prefix", "", "test set prefix")
	flag.StringVar(&r.opts.S3Bucket, "s3bucket", "", "AWS S3 bucket to write failures to")
	flag.StringVar(&r.opts.ReportLambdaFunction, "report-lambda-function", "", "lambda function to report results to")
	flag.StringVar(&r.opts.BuildID, "build-id", "", "build ID of the current build")
	flag.StringVar(&r.opts.Branch, "branch", "", "the branch of the current build")
	flag.BoolVar(&r.opts.Preserve, "preserve", false, "preserve test binary after done")
	flag.StringVar(&r.opts.BuildURL, "build-url", "", "URL for this build (in CI mainly)")
	flag.BoolVar(&r.opts.NoCompile, "no-compile", false, "specify flag if you've pre-compiled the test")
	flag.StringVar(&r.opts.TestBinary, "test-binary", "", "specify the test binary to run")
	flag.StringVar(&r.opts.Timeout, "timeout", "60s", "timeout (in seconds) for any one individual test")
	flag.DurationVar(&r.opts.Pause, "pause", 0, "pause duration between each test (default 0)")
	flag.Parse()
	var d string
	d, err = os.Getwd()
	if err != nil {
		return err
	}
	r.opts.DirBasename = filepath.Base(d)
	return nil
}

func (r *runner) compile() error {
	if r.opts.NoCompile {
		return nil
	}
	fmt.Printf("CMPL: %s\n", r.testerName())
	cmd := exec.Command("go", "test", "-c")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func filter(v []string) []string {
	var ret []string
	for _, s := range v {
		if s != "" {
			ret = append(ret, s)
		}
	}
	return ret
}

func (r *runner) testerName() string {
	if r.opts.TestBinary != "" {
		return r.opts.TestBinary
	}
	return fmt.Sprintf(".%c%s.test", os.PathSeparator, r.opts.DirBasename)
}

func (r *runner) listTests() error {
	cmd := exec.Command(r.testerName(), "-test.list", ".")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return err
	}
	r.tests = filter(strings.Split(out.String(), "\n"))
	return nil
}

func (r *runner) flushTestLogs(test string, log bytes.Buffer) (string, error) {
	logName := fmt.Sprintf("citogo-%s-%s-%s-%s", convertBreakingChars(r.opts.Branch),
		convertBreakingChars(r.opts.BuildID), convertBreakingChars(r.opts.Prefix), test)
	if r.opts.S3Bucket != "" {
		return r.flushLogsToS3(logName, log)
	}
	return r.flushTestLogsToTemp(logName, log)
}

func (r *runner) flushLogsToS3(logName string, log bytes.Buffer) (string, error) {
	return s3put(&log, r.opts.S3Bucket, logName)
}

func (r *runner) flushTestLogsToTemp(logName string, log bytes.Buffer) (string, error) {
	tmpfile, err := ioutil.TempFile("", fmt.Sprintf("%s-", logName))
	if err != nil {
		return "", err
	}
	_, err = tmpfile.Write(log.Bytes())
	if err != nil {
		return "", err
	}
	err = tmpfile.Close()
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("see log: %s", tmpfile.Name()), nil
}

func (r *runner) report(result types.TestResult) {
	if r.opts.ReportLambdaFunction == "" {
		return
	}

	b, err := json.Marshal(result)
	if err != nil {
		logError("error marshalling result: %s", err.Error())
		return
	}

	err = lambdaInvoke(r.opts.ReportLambdaFunction, b)
	if err != nil {
		logError("error reporting flake: %s", err.Error())
	}
}

func (r *runner) newTestResult(outcome types.Outcome, testName string, where string) types.TestResult {
	return types.TestResult{
		Outcome:  outcome,
		TestName: testName,
		Where:    where,
		Branch:   r.opts.Branch,
		BuildID:  r.opts.BuildID,
		Prefix:   r.opts.Prefix,
		BuildURL: r.opts.BuildURL,
	}
}

func (r *runner) runTest(test string) error {
	canRerun := len(r.flakes) < r.opts.Flakes
	outcome, where, err := r.runTestOnce(test, false /* isRerun */, canRerun)
	if err != nil {
		return err
	}
	if outcome == types.OutcomeSuccess {
		return nil
	}
	if len(r.flakes) >= r.opts.Flakes {
		return errTestFailed
	}
	outcome2, _, err2 := r.runTestOnce(test, true /* isRerun */, false /* canRerun */)
	if err2 != nil {
		return err2
	}
	switch outcome2 {
	case types.OutcomeFail:
		return errTestFailed
	case types.OutcomeSuccess:
		r.report(r.newTestResult(types.OutcomeFlake, test, where))
		r.flakes = append(r.flakes, test)
	}
	return nil
}

var errTestFailed = errors.New("test failed")

// runTestOnce only returns an error if there was a problem with the test
// harness code; it does not return an error if the test failed.
func (r *runner) runTestOnce(test string, isRerun bool, canRerun bool) (outcome types.Outcome, where string, err error) {
	defer func() {
		logOutcome := outcome
		if outcome == types.OutcomeFail && canRerun {
			logOutcome = types.OutcomeFlake
		}
		fmt.Printf("%s: %s %s\n", logOutcome.Abbrv(), test, where)
		if logOutcome != types.OutcomeFlake && r.opts.Branch == "master" && err == nil {
			r.report(r.newTestResult(logOutcome, test, where))
		}
	}()

	cmd := exec.Command(r.testerName(), "-test.run", "^"+test+"$", "-test.timeout", r.opts.Timeout)
	if isRerun {
		cmd.Env = append(os.Environ(), "CITOGO_FLAKE_RERUN=1")
	}
	var combined bytes.Buffer
	cmd.Stdout = &combined
	cmd.Stderr = &combined
	testErr := cmd.Run()
	if testErr != nil {
		err = errTestFailed

		var flushErr error
		where, flushErr := r.flushTestLogs(test, combined)
		if flushErr != nil {
			return types.OutcomeFail, "", flushErr
		}
		return types.OutcomeFail, where, nil
	}
	return types.OutcomeSuccess, "", nil
}

func (r *runner) runTestFixError(t string) error {
	err := r.runTest(t)
	if err == nil {
		return nil
	}
	if err != errTestFailed {
		return err
	}
	r.fails = append(r.fails, t)
	if r.opts.Fails < 0 {
		// We have an infinite fail budget, so keep plowing through
		// failed tests. This test run is still going to fail.
		return nil
	}
	if r.opts.Fails >= len(r.fails) {
		// We've failed less than our budget, so we can still keep going.
		// This test run is still going to fail.
		return nil
	}
	// We ate up our fail budget.
	return err
}

func (r *runner) runTests() error {
	for _, f := range r.tests {
		err := r.runTestFixError(f)
		if err != nil {
			return err
		}
		if r.opts.Pause > 0 {
			time.Sleep(r.opts.Pause)
		}
	}
	return nil
}

func (r *runner) cleanup() {
	if r.opts.Preserve || r.opts.NoCompile {
		return
	}
	n := r.testerName()
	err := os.Remove(n)
	fmt.Printf("RMOV: %s\n", n)
	if err != nil {
		logError("could not remove %s: %s", n, err.Error())
	}
}

func (r *runner) debugStartup() {
	dir, _ := os.Getwd()
	fmt.Printf("WDIR: %s\n", dir)
}

func (r *runner) testExists() (bool, error) {
	f := r.testerName()
	info, err := os.Stat(f)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	if info.Mode().IsRegular() {
		return true, nil
	}
	return false, fmt.Errorf("%s: file of wrong type", f)

}

func (r *runner) run() error {
	start := time.Now()
	err := r.parseArgs()
	if err != nil {
		return err
	}

	r.debugStartup()
	err = r.compile()
	if err != nil {
		return err
	}
	exists, err := r.testExists()
	if exists {
		err = r.listTests()
		if err != nil {
			return err
		}
		err = r.runTests()
		r.cleanup()
	}
	end := time.Now()
	diff := end.Sub(start)
	fmt.Printf("DONE: in %s\n", diff)
	if err != nil {
		return err
	}
	if len(r.fails) > 0 {
		// If we have more than 15 tests, repeat at the end which tests failed,
		// so we don't have to scroll all the way up.
		if len(r.tests) > 15 {
			for _, t := range r.fails {
				fmt.Printf("FAIL: %s\n", t)
			}
		}
		return fmt.Errorf("RED!: %d total tests failed", len(r.fails))
	}
	return nil
}

func main2() error {
	runner := runner{}
	return runner.run()
}

func main() {
	err := main2()
	if err != nil {
		logError(err.Error())
		fmt.Printf("EXIT: 2\n")
		os.Exit(2)
	}
	fmt.Printf("EXIT: 0\n")
	os.Exit(0)
}
<commit_msg>support parallelizing citogo (#23819)<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/keybase/client/go/citogo/types"
)

type opts struct {
	Flakes               int
	Fails                int
	Prefix               string
	S3Bucket             string
	ReportLambdaFunction string
	DirBasename          string
	BuildID              string
	Branch               string
	Parallel             int
	Preserve             bool
	BuildURL             string
	NoCompile            bool
	TestBinary           string
	Timeout              string
	Pause                time.Duration
}

func logError(f string, args ...interface{}) {
	s := fmt.Sprintf(f, args...)
	if s[len(s)-1] != '\n' {
		s += "\n"
	}
	fmt.Fprintf(os.Stderr, "%s", s)
}

type runner struct {
	opts   opts
	flakes []string
	fails  []string
	tests  []string
}

func convertBreakingChars(s string) string {
	// replace either the unix or the DOS directory marker
	// with an underscore, so as not to break the directory
	// structure of where we are storing the log
	s = strings.ReplaceAll(s, "/", "_")
	s = strings.ReplaceAll(s, "\\", "_")
	s = strings.ReplaceAll(s, "-", "_")
	return s
}

func (r *runner) parseArgs() (err error) {
	flag.IntVar(&r.opts.Flakes, "flakes", 3, "number of allowed flakes")
	flag.IntVar(&r.opts.Fails, "fails", -1, "number of fails allowed before quitting")
	flag.IntVar(&r.opts.Parallel, "parallel", 1, "number of tests to run in parallel")
	flag.StringVar(&r.opts.Prefix, "prefix", "", "test set prefix")
	flag.StringVar(&r.opts.S3Bucket, "s3bucket", "", "AWS S3 bucket to write failures to")
	flag.StringVar(&r.opts.ReportLambdaFunction, "report-lambda-function", "", "lambda function to report results to")
	flag.StringVar(&r.opts.BuildID, "build-id", "", "build ID of the current build")
	flag.StringVar(&r.opts.Branch, "branch", "", "the branch of the current build")
	flag.BoolVar(&r.opts.Preserve, "preserve", false, "preserve test binary after done")
	flag.StringVar(&r.opts.BuildURL, "build-url", "", "URL for this build (in CI mainly)")
	flag.BoolVar(&r.opts.NoCompile, "no-compile", false, "specify flag if you've pre-compiled the test")
	flag.StringVar(&r.opts.TestBinary, "test-binary", "", "specify the test binary to run")
	flag.StringVar(&r.opts.Timeout, "timeout", "60s", "timeout (in seconds) for any one individual test")
	flag.DurationVar(&r.opts.Pause, "pause", 0, "pause duration between each test (default 0)")
	flag.Parse()
	var d string
	d, err = os.Getwd()
	if err != nil {
		return err
	}
	r.opts.DirBasename = filepath.Base(d)
	return nil
}

func (r *runner) compile() error {
	if r.opts.NoCompile {
		return nil
	}
	fmt.Printf("CMPL: %s\n", r.testerName())
	cmd := exec.Command("go", "test", "-c")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	return cmd.Run()
}

func filter(v []string) []string {
	var ret []string
	for _, s := range v {
		if s != "" {
			ret = append(ret, s)
		}
	}
	return ret
}

func (r *runner) testerName() string {
	if r.opts.TestBinary != "" {
		return r.opts.TestBinary
	}
	return fmt.Sprintf(".%c%s.test", os.PathSeparator, r.opts.DirBasename)
}

func (r *runner) listTests() error {
	cmd := exec.Command(r.testerName(), "-test.list", ".")
	var out bytes.Buffer
	cmd.Stdout = &out
	err := cmd.Run()
	if err != nil {
		return err
	}
	r.tests = filter(strings.Split(out.String(), "\n"))
	return nil
}

func (r *runner) flushTestLogs(test string, log bytes.Buffer) (string, error) {
	logName := fmt.Sprintf("citogo-%s-%s-%s-%s", convertBreakingChars(r.opts.Branch),
		convertBreakingChars(r.opts.BuildID), convertBreakingChars(r.opts.Prefix), test)
	if r.opts.S3Bucket != "" {
		return r.flushLogsToS3(logName, log)
	}
	return r.flushTestLogsToTemp(logName, log)
}

func (r *runner) flushLogsToS3(logName string, log bytes.Buffer) (string, error) {
	return s3put(&log, r.opts.S3Bucket, logName)
}

func (r *runner) flushTestLogsToTemp(logName string, log bytes.Buffer) (string, error) {
	tmpfile, err := ioutil.TempFile("", fmt.Sprintf("%s-", logName))
	if err != nil {
		return "", err
	}
	_, err = tmpfile.Write(log.Bytes())
	if err != nil {
		return "", err
	}
	err = tmpfile.Close()
	if err != nil {
		return "", err
	}
	return fmt.Sprintf("see log: %s", tmpfile.Name()), nil
}

func (r *runner) report(result types.TestResult) {
	if r.opts.ReportLambdaFunction == "" {
		return
	}

	b, err := json.Marshal(result)
	if err != nil {
		logError("error marshalling result: %s", err.Error())
		return
	}

	err = lambdaInvoke(r.opts.ReportLambdaFunction, b)
	if err != nil {
		logError("error reporting flake: %s", err.Error())
	}
}

func (r *runner) newTestResult(outcome types.Outcome, testName string, where string) types.TestResult {
	return types.TestResult{
		Outcome:  outcome,
		TestName: testName,
		Where:    where,
		Branch:   r.opts.Branch,
		BuildID:  r.opts.BuildID,
		Prefix:   r.opts.Prefix,
		BuildURL: r.opts.BuildURL,
	}
}

func (r *runner) runTest(test string) error {
	canRerun := len(r.flakes) < r.opts.Flakes
	outcome, where, err := r.runTestOnce(test, false /* isRerun */, canRerun)
	if err != nil {
		return err
	}
	if outcome == types.OutcomeSuccess {
		return nil
	}
	if len(r.flakes) >= r.opts.Flakes {
		return errTestFailed
	}
	outcome2, _, err2 := r.runTestOnce(test, true /* isRerun */, false /* canRerun */)
	if err2 != nil {
		return err2
	}
	switch outcome2 {
	case types.OutcomeFail:
		return errTestFailed
	case types.OutcomeSuccess:
		r.report(r.newTestResult(types.OutcomeFlake, test, where))
		r.flakes = append(r.flakes, test)
	}
	return nil
}

var errTestFailed = errors.New("test failed")

// runTestOnce only returns an error if there was a problem with the test
// harness code; it does not return an error if the test failed.
func (r *runner) runTestOnce(test string, isRerun bool, canRerun bool) (outcome types.Outcome, where string, err error) {
	defer func() {
		logOutcome := outcome
		if outcome == types.OutcomeFail && canRerun {
			logOutcome = types.OutcomeFlake
		}
		fmt.Printf("%s: %s %s\n", logOutcome.Abbrv(), test, where)
		if logOutcome != types.OutcomeFlake && r.opts.Branch == "master" && err == nil {
			r.report(r.newTestResult(logOutcome, test, where))
		}
	}()

	cmd := exec.Command(r.testerName(), "-test.run", "^"+test+"$", "-test.timeout", r.opts.Timeout)
	if isRerun {
		cmd.Env = append(os.Environ(), "CITOGO_FLAKE_RERUN=1")
	}
	var combined bytes.Buffer
	cmd.Stdout = &combined
	cmd.Stderr = &combined
	testErr := cmd.Run()
	if testErr != nil {
		err = errTestFailed

		var flushErr error
		where, flushErr := r.flushTestLogs(test, combined)
		if flushErr != nil {
			return types.OutcomeFail, "", flushErr
		}
		return types.OutcomeFail, where, nil
	}
	return types.OutcomeSuccess, "", nil
}

func (r *runner) runTestFixError(t string) error {
	err := r.runTest(t)
	if err == nil {
		return nil
	}
	if err != errTestFailed {
		return err
	}
	r.fails = append(r.fails, t)
	if r.opts.Fails < 0 {
		// We have an infinite fail budget, so keep plowing through
		// failed tests. This test run is still going to fail.
		return nil
	}
	if r.opts.Fails >= len(r.fails) {
		// We've failed less than our budget, so we can still keep going.
		// This test run is still going to fail.
		return nil
	}
	// We ate up our fail budget.
	return err
}

func (r *runner) runTests() error {
	var eg errgroup.Group
	q := make(chan string, len(r.tests))
	for i := 0; i < r.opts.Parallel; i++ {
		eg.Go(func() error {
			for f := range q {
				err := r.runTestFixError(f)
				if err != nil {
					return err
				}
				if r.opts.Pause > 0 {
					time.Sleep(r.opts.Pause)
				}
			}
			return nil
		})
	}
	for _, f := range r.tests {
		q <- f
	}
	close(q)
	return eg.Wait()
}

func (r *runner) cleanup() {
	if r.opts.Preserve || r.opts.NoCompile {
		return
	}
	n := r.testerName()
	err := os.Remove(n)
	fmt.Printf("RMOV: %s\n", n)
	if err != nil {
		logError("could not remove %s: %s", n, err.Error())
	}
}

func (r *runner) debugStartup() {
	dir, _ := os.Getwd()
	fmt.Printf("WDIR: %s\n", dir)
}

func (r *runner) testExists() (bool, error) {
	f := r.testerName()
	info, err := os.Stat(f)
	if os.IsNotExist(err) {
		return false, nil
	}
	if err != nil {
		return false, err
	}
	if info.Mode().IsRegular() {
		return true, nil
	}
	return false, fmt.Errorf("%s: file of wrong type", f)

}

func (r *runner) run() error {
	start := time.Now()
	err := r.parseArgs()
	if err != nil {
		return err
	}

	r.debugStartup()
	err = r.compile()
	if err != nil {
		return err
	}
	exists, err := r.testExists()
	if exists {
		err = r.listTests()
		if err != nil {
			return err
		}
		err = r.runTests()
		r.cleanup()
	}
	end := time.Now()
	diff := end.Sub(start)
	fmt.Printf("DONE: in %s\n", diff)
	if err != nil {
		return err
	}
	if len(r.fails) > 0 {
		// If we have more than 15 tests, repeat at the end which tests failed,
		// so we don't have to scroll all the way up.
		if len(r.tests) > 15 {
			for _, t := range r.fails {
				fmt.Printf("FAIL: %s\n", t)
			}
		}
		return fmt.Errorf("RED!: %d total tests failed", len(r.fails))
	}
	return nil
}

func main2() error {
	runner := runner{}
	return runner.run()
}

func main() {
	err := main2()
	if err != nil {
		logError(err.Error())
		fmt.Printf("EXIT: 2\n")
		os.Exit(2)
	}
	fmt.Printf("EXIT: 0\n")
	os.Exit(0)
}
<|endoftext|>


<commit_before>package geom

import (
	"fmt"
	"math"
)

type Margin struct {
	Top    float64
	Bottom float64
	Left   float64
	Right  float64
}

func MakeMargin(dist float64) Margin {
	return Margin{
		Top:    dist,
		Bottom: dist,
		Left:   dist,
		Right:  dist,
	}
}

func (m Margin) TopLeft(p Point) Point {
	return Point{
		X: p.X + m.Left,
		Y: p.Y + m.Top,
	}
}

func (m Margin) BottomRight(p Point) Point {
	return Point{
		X: p.X - m.Right,
		Y: p.Y - m.Bottom,
	}
}

func (m Margin) GrowRect(r Rect) Rect {
	return Rect{
		X: r.X - m.Left,
		Y: r.Y - m.Top,
		W: r.W + m.Width(),
		H: r.H + m.Height(),
	}
}

func (m Margin) GrowSize(s Size) Size {
	return Size{
		W: s.W + m.Width(),
		H: s.H + m.Height(),
	}
}

func (m Margin) ShrinkRect(r Rect) Rect {
	s := Rect{
		X: r.X + m.Left,
		Y: r.Y + m.Top,
		W: r.W - m.Width(),
		H: r.H - m.Height(),
	}

	if s.W < 0 {
		r.X = r.X + (r.W / 2)
		s.W = 0
	}

	if s.H < 0 {
		r.Y = r.Y + (r.H / 2)
		s.H = 0
	}

	return s
}

func (m Margin) ShrinkSize(s Size) Size {
	return Size{
		W: math.Max(0, s.W-m.Width()),
		H: math.Max(0, s.H-m.Height()),
	}
}

func (m Margin) Width() float64 {
	return m.Left + m.Right
}

func (m Margin) Height() float64 {
	return m.Top + m.Bottom
}

func (m Margin) Size() Size {
	return Size{
		W: m.Width(),
		H: m.Height(),
	}
}

func (m Margin) String() string {
	return fmt.Sprintf("margins { top = %g, bottom = %g, left = %g, right = %g }", m.Top, m.Bottom, m.Left, m.Right)
}
<commit_msg>udpate margin documentation<commit_after>package geom

import (
	"fmt"
	"math"
)

// The Margin type represent 2D margins of a rectangle area.
type Margin struct {
	Top    float64
	Bottom float64
	Left   float64
	Right  float64
}

// Makes a new Margin value where every side of a rectangle would be applied the
// margin value passed as argument.
func MakeMargin(m float64) Margin {
	return Margin{
		Top:    m,
		Bottom: m,
		Left:   m,
		Right:  m,
	}
}

// Given a point that would be the top-left corner of a rectangle, applies the
// margin and return the modified coordinates.
func (m Margin) TopLeft(p Point) Point {
	return Point{
		X: p.X + m.Left,
		Y: p.Y + m.Top,
	}
}

// Given a point that would be the bottom-right corner of a rectangle, applies
// the margin and return the modified corrdinates.
func (m Margin) BottomRight(p Point) Point {
	return Point{
		X: p.X - m.Right,
		Y: p.Y - m.Bottom,
	}
}

// Grows the given rectangle by applying the margin and returns the modified
// rectangle.
func (m Margin) GrowRect(r Rect) Rect {
	return Rect{
		X: r.X - m.Left,
		Y: r.Y - m.Top,
		W: r.W + m.Width(),
		H: r.H + m.Height(),
	}
}

// Grows a size value by applying the margin and returns the modified size.
func (m Margin) GrowSize(s Size) Size {
	return Size{
		W: s.W + m.Width(),
		H: s.H + m.Height(),
	}
}

// Shrinks the given rectangle by applying the margin and returns the modified
// rectangle.
func (m Margin) ShrinkRect(r Rect) Rect {
	s := Rect{
		X: r.X + m.Left,
		Y: r.Y + m.Top,
		W: r.W - m.Width(),
		H: r.H - m.Height(),
	}

	if s.W < 0 {
		r.X = r.X + (r.W / 2)
		s.W = 0
	}

	if s.H < 0 {
		r.Y = r.Y + (r.H / 2)
		s.H = 0
	}

	return s
}

// Shrinks the given size by applying the margin and returns the modified size.
func (m Margin) ShrinkSize(s Size) Size {
	return Size{
		W: math.Max(0, s.W-m.Width()),
		H: math.Max(0, s.H-m.Height()),
	}
}

// Returns the sum of the left and right values of the given margin.
func (m Margin) Width() float64 {
	return m.Left + m.Right
}

// Returns the sum of the top and bottom values of the given margin.
func (m Margin) Height() float64 {
	return m.Top + m.Bottom
}

// Returns the combined width and height of the margin as a size value.
func (m Margin) Size() Size {
	return Size{
		W: m.Width(),
		H: m.Height(),
	}
}

// Returns a string representation of the margin value.
func (m Margin) String() string {
	return fmt.Sprintf("margins { top = %g, bottom = %g, left = %g, right = %g }", m.Top, m.Bottom, m.Left, m.Right)
}
<|endoftext|>