<commit_before>// Copyright 2009 The Go 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 os

import (
	"sync"
	"syscall"
)

var getwdCache struct {
	sync.Mutex
	dir string
}

// useSyscallwd determines whether to use the return value of
// syscall.Getwd based on its error.
var useSyscallwd = func(error) bool { return true }

// Getwd returns a rooted path name corresponding to the
// current directory.  If the current directory can be
// reached via multiple paths (due to symbolic links),
// Getwd may return any one of them.
func Getwd() (dir string, err error) {
	// Clumsy but widespread kludge:
	// if $PWD is set and matches ".", use it.
	dot, err := Stat(".")
	if err != nil {
		return "", err
	}
	dir = Getenv("PWD")
	if len(dir) > 0 && dir[0] == '/' {
		d, err := Stat(dir)
		if err == nil && SameFile(dot, d) {
			return dir, nil
		}
	}

	// If the operating system provides a Getwd call, use it.
	// Otherwise, we're trying to find our way back to ".".
	if syscall.ImplementsGetwd {
		s, e := syscall.Getwd()
		if useSyscallwd(e) {
			return s, NewSyscallError("getwd", e)
		}
	}

	// Apply same kludge but to cached dir instead of $PWD.
	getwdCache.Lock()
	dir = getwdCache.dir
	getwdCache.Unlock()
	if len(dir) > 0 {
		d, err := Stat(dir)
		if err == nil && SameFile(dot, d) {
			return dir, nil
		}
	}

	// Root is a special case because it has no parent
	// and ends in a slash.
	root, err := Stat("/")
	if err != nil {
		// Can't stat root - no hope.
		return "", err
	}
	if SameFile(root, dot) {
		return "/", nil
	}

	// General algorithm: find name in parent
	// and then find name of parent.  Each iteration
	// adds /name to the beginning of dir.
	dir = ""
	for parent := ".."; ; parent = "../" + parent {
		if len(parent) >= 1024 { // Sanity check
			return "", syscall.ENAMETOOLONG
		}
		fd, err := Open(parent)
		if err != nil {
			return "", err
		}

		for {
			names, err := fd.Readdirnames(100)
			if err != nil {
				fd.Close()
				return "", err
			}
			for _, name := range names {
				d, _ := Lstat(parent + "/" + name)
				if SameFile(d, dot) {
					dir = "/" + name + dir
					goto Found
				}
			}
		}

	Found:
		pd, err := fd.Stat()
		if err != nil {
			return "", err
		}
		fd.Close()
		if SameFile(pd, root) {
			break
		}
		// Set up for next round.
		dot = pd
	}

	// Save answer as hint to avoid the expensive path next time.
	getwdCache.Lock()
	getwdCache.dir = dir
	getwdCache.Unlock()

	return dir, nil
}
<commit_msg>os: simplify windows Getwd (fixes build)<commit_after>// Copyright 2009 The Go 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 os

import (
	"runtime"
	"sync"
	"syscall"
)

var getwdCache struct {
	sync.Mutex
	dir string
}

// useSyscallwd determines whether to use the return value of
// syscall.Getwd based on its error.
var useSyscallwd = func(error) bool { return true }

// Getwd returns a rooted path name corresponding to the
// current directory.  If the current directory can be
// reached via multiple paths (due to symbolic links),
// Getwd may return any one of them.
func Getwd() (dir string, err error) {
	if runtime.GOOS == "windows" {
		return syscall.Getwd()
	}

	// Clumsy but widespread kludge:
	// if $PWD is set and matches ".", use it.
	dot, err := Stat(".")
	if err != nil {
		return "", err
	}
	dir = Getenv("PWD")
	if len(dir) > 0 && dir[0] == '/' {
		d, err := Stat(dir)
		if err == nil && SameFile(dot, d) {
			return dir, nil
		}
	}

	// If the operating system provides a Getwd call, use it.
	// Otherwise, we're trying to find our way back to ".".
	if syscall.ImplementsGetwd {
		s, e := syscall.Getwd()
		if useSyscallwd(e) {
			return s, NewSyscallError("getwd", e)
		}
	}

	// Apply same kludge but to cached dir instead of $PWD.
	getwdCache.Lock()
	dir = getwdCache.dir
	getwdCache.Unlock()
	if len(dir) > 0 {
		d, err := Stat(dir)
		if err == nil && SameFile(dot, d) {
			return dir, nil
		}
	}

	// Root is a special case because it has no parent
	// and ends in a slash.
	root, err := Stat("/")
	if err != nil {
		// Can't stat root - no hope.
		return "", err
	}
	if SameFile(root, dot) {
		return "/", nil
	}

	// General algorithm: find name in parent
	// and then find name of parent.  Each iteration
	// adds /name to the beginning of dir.
	dir = ""
	for parent := ".."; ; parent = "../" + parent {
		if len(parent) >= 1024 { // Sanity check
			return "", syscall.ENAMETOOLONG
		}
		fd, err := Open(parent)
		if err != nil {
			return "", err
		}

		for {
			names, err := fd.Readdirnames(100)
			if err != nil {
				fd.Close()
				return "", err
			}
			for _, name := range names {
				d, _ := Lstat(parent + "/" + name)
				if SameFile(d, dot) {
					dir = "/" + name + dir
					goto Found
				}
			}
		}

	Found:
		pd, err := fd.Stat()
		if err != nil {
			return "", err
		}
		fd.Close()
		if SameFile(pd, root) {
			break
		}
		// Set up for next round.
		dot = pd
	}

	// Save answer as hint to avoid the expensive path next time.
	getwdCache.Lock()
	getwdCache.dir = dir
	getwdCache.Unlock()

	return dir, nil
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go 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 os

import "syscall"

// An operating-system independent representation of Unix data structures.
// OS-specific routines in this directory convert the OS-local versions to these.

// Getpagesize returns the underlying system's memory page size.
func Getpagesize() int { return syscall.Getpagesize() }

// A FileInfo describes a file and is returned by Stat, Fstat, and Lstat
type FileInfo struct {
	Dev             uint64 // device number of file system holding file.
	Ino             uint64 // inode number.
	Nlink           uint64 // number of hard links.
	Mode            uint32 // permission and mode bits.
	Uid             int    // user id of owner.
	Gid             int    // group id of owner.
	Rdev            uint64 // device type for special file.
	Size            int64  // length in bytes.
	Blksize         int64  // size of blocks, in bytes.
	Blocks          int64  // number of blocks allocated for file.
	Atime_ns        int64  // access time; nanoseconds since epoch.
	Mtime_ns        int64  // modified time; nanoseconds since epoch.
	Ctime_ns        int64  // status change time; nanoseconds since epoch.
	Name            string // name of file as presented to Open.
	FollowedSymlink bool   // followed a symlink to get this information
}

// IsFifo reports whether the FileInfo describes a FIFO file.
func (f *FileInfo) IsFifo() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFIFO }

// IsChar reports whether the FileInfo describes a character special file.
func (f *FileInfo) IsChar() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFCHR }

// IsDirectory reports whether the FileInfo describes a directory.
func (f *FileInfo) IsDirectory() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFDIR }

// IsBlock reports whether the FileInfo describes a block special file.
func (f *FileInfo) IsBlock() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFBLK }

// IsRegular reports whether the FileInfo describes a regular file.
func (f *FileInfo) IsRegular() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFREG }

// IsSymlink reports whether the FileInfo describes a symbolic link.
func (f *FileInfo) IsSymlink() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFLNK }

// IsSocket reports whether the FileInfo describes a socket.
func (f *FileInfo) IsSocket() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFSOCK }

// Permission returns the file permission bits.
func (f *FileInfo) Permission() uint32 { return f.Mode & 0777 }
<commit_msg>os: fix documentation for FileInfo.Name. It's the base name, not the full name. Fixes issue 2047.<commit_after>// Copyright 2009 The Go 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 os

import "syscall"

// An operating-system independent representation of Unix data structures.
// OS-specific routines in this directory convert the OS-local versions to these.

// Getpagesize returns the underlying system's memory page size.
func Getpagesize() int { return syscall.Getpagesize() }

// A FileInfo describes a file and is returned by Stat, Fstat, and Lstat
type FileInfo struct {
	Dev             uint64 // device number of file system holding file.
	Ino             uint64 // inode number.
	Nlink           uint64 // number of hard links.
	Mode            uint32 // permission and mode bits.
	Uid             int    // user id of owner.
	Gid             int    // group id of owner.
	Rdev            uint64 // device type for special file.
	Size            int64  // length in bytes.
	Blksize         int64  // size of blocks, in bytes.
	Blocks          int64  // number of blocks allocated for file.
	Atime_ns        int64  // access time; nanoseconds since epoch.
	Mtime_ns        int64  // modified time; nanoseconds since epoch.
	Ctime_ns        int64  // status change time; nanoseconds since epoch.
	Name            string // base name of the file name provided in Open, Stat, etc.
	FollowedSymlink bool   // followed a symlink to get this information
}

// IsFifo reports whether the FileInfo describes a FIFO file.
func (f *FileInfo) IsFifo() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFIFO }

// IsChar reports whether the FileInfo describes a character special file.
func (f *FileInfo) IsChar() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFCHR }

// IsDirectory reports whether the FileInfo describes a directory.
func (f *FileInfo) IsDirectory() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFDIR }

// IsBlock reports whether the FileInfo describes a block special file.
func (f *FileInfo) IsBlock() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFBLK }

// IsRegular reports whether the FileInfo describes a regular file.
func (f *FileInfo) IsRegular() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFREG }

// IsSymlink reports whether the FileInfo describes a symbolic link.
func (f *FileInfo) IsSymlink() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFLNK }

// IsSocket reports whether the FileInfo describes a socket.
func (f *FileInfo) IsSocket() bool { return (f.Mode & syscall.S_IFMT) == syscall.S_IFSOCK }

// Permission returns the file permission bits.
func (f *FileInfo) Permission() uint32 { return f.Mode & 0777 }
<|endoftext|>


<commit_before>// puzzle at http://adventofcode.com/day/4

package main

import (
	"crypto/md5"
	"fmt"
	"strconv"
)

var key string = "ckczppom"

func main() {

	for i := 0; ; i++ {

		if i%10001 == 10000 {
			fmt.Printf("i = %v...\n", i)
		}

		test := key + strconv.Itoa(i)
		md5String := MD5Hash(test)

		if md5String[:5] == "00000" {
			fmt.Printf("First hash with prefix '00000' is %v, from %v\n", md5String, test)
			break
		}
	}
}

func MD5Hash(s string) string {
	md5Hash := md5.Sum([]byte(s))
	return fmt.Sprintf("%x", md5Hash)
}
<commit_msg>day 4 part 2<commit_after>// puzzle at http://adventofcode.com/day/4

package main

import (
	"crypto/md5"
	"fmt"
	"strconv"
)

var key string = "ckczppom"

func main() {

	for i := 0; ; i++ {

		if i%100001 == 100000 {
			fmt.Printf("i = %v...\n", i)
		}

		test := key + strconv.Itoa(i)
		md5String := MD5Hash(test)

		targetPrefix := "000000"
		if md5String[:len(targetPrefix)] == targetPrefix {
			fmt.Printf("First hash with prefix '000000' is %v, from %v\n", md5String, test)
			break
		}
	}
}

func MD5Hash(s string) string {
	md5Hash := md5.Sum([]byte(s))
	return fmt.Sprintf("%x", md5Hash)
}
<|endoftext|>


<commit_before>package q

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
)

var (
	LogFile = "/var/log/q"
)

func Println(a ...interface{}) {
	fd, err := os.OpenFile(LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	ptr, file, line, ok := runtime.Caller(1)
	if ok {
		file = filepath.Base(file)
		s := []interface{}{
			fmt.Sprintf("%s:%d", file, line), // filename:number
			runtime.FuncForPC(ptr).Name(),    // caller name
		}
		s = append(s, a...)

		_, err = fmt.Fprintln(fd, s...)
	} else {
		_, err = fmt.Fprintln(fd, a...)
	}

	if err != nil {
		panic(err)
	}
}

func Printf(format string, a ...interface{}) {
	fd, err := os.OpenFile(LogFile, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	_, err = fmt.Fprintf(fd, format, a...)
	if err != nil {
		panic(err)
	}
}
<commit_msg>Put log in /tmp to avoid perm issues<commit_after>package q

import (
	"fmt"
	"os"
	"path/filepath"
	"runtime"
)

var (
	LogFile = "q.log"
)

func Println(a ...interface{}) {
	f := filepath.Join("/tmp", LogFile)
	fd, err := os.OpenFile(f, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	ptr, file, line, ok := runtime.Caller(1)
	if ok {
		file = filepath.Base(file)
		s := []interface{}{
			fmt.Sprintf("%s:%d", file, line), // filename:number
			runtime.FuncForPC(ptr).Name(),    // caller name
		}
		s = append(s, a...)

		_, err = fmt.Fprintln(fd, s...)
	} else {
		_, err = fmt.Fprintln(fd, a...)
	}

	if err != nil {
		panic(err)
	}
}

func Printf(format string, a ...interface{}) {
	f := filepath.Join("/tmp", LogFile)
	fd, err := os.OpenFile(f, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	_, err = fmt.Fprintf(fd, format, a...)
	if err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>package benchmark_bbs_test

import (
	"crypto/rand"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/cloudfoundry-incubator/bbs"
	etcddb "github.com/cloudfoundry-incubator/bbs/db/etcd"
	"github.com/cloudfoundry-incubator/bbs/encryption"
	"github.com/cloudfoundry-incubator/bbs/format"
	"github.com/cloudfoundry-incubator/benchmark-bbs/generator"
	"github.com/cloudfoundry-incubator/benchmark-bbs/reporter"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/coreos/go-etcd/etcd"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/zorkian/go-datadog-api"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

var (
	bbsAddress           string
	bbsClientCert        string
	bbsClientKey         string
	etcdFlags            *ETCDFlags
	dataDogAPIKey        string
	dataDogAppKey        string
	awsAccessKeyID       string
	awsSecretAccessKey   string
	awsBucketName        string
	awsRegion            string
	desiredLRPs          int
	encryptionFlags      *encryption.EncryptionFlags
	metricPrefix         string
	numTrials            int
	numPopulateWorkers   int
	expectedLRPCount     int
	expectedLRPTolerance float64
	errorTolerance       float64
	logLevel             string
	logFilename          string

	logger               lager.Logger
	etcdClient           *etcd.Client
	etcdDB               *etcddb.ETCDDB
	bbsClient            bbs.Client
	bbsClientHTTPTimeout time.Duration
	dataDogClient        *datadog.Client
	dataDogReporter      reporter.DataDogReporter
	reporters            []Reporter
)

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.IntVar(&numTrials, "numTrials", 10, "number of benchmark trials to average across")
	flag.IntVar(&numPopulateWorkers, "numPopulateWorkers", 10, "number of workers generating desired LRPs during setup")
	flag.IntVar(&desiredLRPs, "desiredLRPs", 0, "number of single instance DesiredLRPs to create")

	flag.StringVar(&bbsAddress, "bbsAddress", "", "Address of the BBS Server")
	flag.StringVar(&bbsClientCert, "bbsClientCert", "", "BBS client SSL certificate")
	flag.StringVar(&bbsClientKey, "bbsClientKey", "", "BBS client SSL key")
	flag.DurationVar(&bbsClientHTTPTimeout, "bbsClientHTTPTimeout", 0, "BBS client HTTP timeout")

	flag.StringVar(&dataDogAPIKey, "dataDogAPIKey", "", "DataDog API key")
	flag.StringVar(&dataDogAppKey, "dataDogAppKey", "", "DataDog app Key")
	flag.StringVar(&metricPrefix, "metricPrefix", "", "DataDog metric prefix")

	flag.StringVar(&awsAccessKeyID, "awsAccessKeyID", "", "AWS Access Key ID")
	flag.StringVar(&awsSecretAccessKey, "awsSecretAccessKey", "", "AWS Secret Access Key")
	flag.StringVar(&awsBucketName, "awsBucketName", "", "AWS Bucket to store metrics")
	flag.StringVar(&awsRegion, "awsRegion", "us-west-1", "AWS Bucket to store metrics")

	flag.StringVar(&logLevel, "logLevel", string(INFO), "log level: debug, info, error or fatal")
	flag.StringVar(&logFilename, "logFilename", "", "Name of local file to save logs to")
	flag.Float64Var(&errorTolerance, "errorTolerance", 0.05, "error tollerance rate")

	etcdFlags = AddETCDFlags(flag.CommandLine)
	encryptionFlags = encryption.AddEncryptionFlags(flag.CommandLine)

	flag.Parse()

	if bbsAddress == "" {
		log.Fatal("bbsAddress is required")
	}

	BenchmarkConvergenceGathering(numTrials)
	BenchmarkNsyncFetching(numTrials)
	BenchmarkRouteEmitterFetching(numTrials)
}

func TestBenchmarkBbs(t *testing.T) {
	var lagerLogLevel lager.LogLevel
	switch logLevel {
	case DEBUG:
		lagerLogLevel = lager.DEBUG
	case INFO:
		lagerLogLevel = lager.INFO
	case ERROR:
		lagerLogLevel = lager.ERROR
	case FATAL:
		lagerLogLevel = lager.FATAL
	default:
		panic(fmt.Errorf("unknown log level: %s", logLevel))
	}

	var logWriter io.Writer
	if logFilename == "" {
		logWriter = GinkgoWriter
	} else {
		logFile, err := os.Create(logFilename)
		if err != nil {
			panic(fmt.Errorf("Error opening file '%s': %s", logFilename, err.Error()))
		}
		defer logFile.Close()

		logWriter = logFile
	}

	logger = lager.NewLogger("bbs-benchmarks-test")
	logger.RegisterSink(lager.NewWriterSink(logWriter, lagerLogLevel))

	reporters = []Reporter{}

	if dataDogAPIKey != "" && dataDogAppKey != "" {
		dataDogClient = datadog.NewClient(dataDogAPIKey, dataDogAppKey)
		dataDogReporter = reporter.NewDataDogReporter(logger, metricPrefix, dataDogClient)
		reporters = append(reporters, &dataDogReporter)
	}

	if awsAccessKeyID != "" && awsSecretAccessKey != "" && awsBucketName != "" {
		creds := credentials.NewStaticCredentials(awsAccessKeyID, awsSecretAccessKey, "")
		s3Client := s3.New(&aws.Config{
			Region:      &awsRegion,
			Credentials: creds,
		})
		uploader := s3manager.NewUploader(&s3manager.UploadOptions{S3: s3Client})
		reporter := reporter.NewS3Reporter(logger, awsBucketName, uploader)
		reporters = append(reporters, &reporter)
	}

	RegisterFailHandler(Fail)
	RunSpecsWithDefaultAndCustomReporters(t, "Benchmark BBS Suite", reporters)
}

var _ = SynchronizedBeforeSuite(func() []byte {
	etcdOptions, err := etcdFlags.Validate()
	Expect(err).NotTo(HaveOccurred())

	etcdClient = initializeEtcdClient(logger, etcdOptions)
	bbsClient = initializeBBSClient(logger, bbsClientHTTPTimeout)
	etcdDB = initializeETCDDB(logger, etcdClient)

	purge("/v1/desired_lrp")
	purge("/v1/actual")

	_, err = bbsClient.Domains()
	Expect(err).NotTo(HaveOccurred())
	if desiredLRPs > 0 {
		desiredLRPGenerator := generator.NewDesiredLRPGenerator(errorTolerance, metricPrefix, numPopulateWorkers, bbsClient, dataDogClient)
		expectedLRPCount, err = desiredLRPGenerator.Generate(logger, desiredLRPs)
		Expect(err).NotTo(HaveOccurred())
		expectedLRPTolerance = float64(expectedLRPCount) * errorTolerance
	}
	values := fmt.Sprintf("%d %f", expectedLRPCount, expectedLRPTolerance)
	return []byte(values)
}, func(data []byte) {
	values := string(data)
	fmt.Sscanf(values, "%d %f", &expectedLRPCount, &expectedLRPTolerance)
	if etcdClient == nil {
		etcdOptions, err := etcdFlags.Validate()
		Expect(err).NotTo(HaveOccurred())
		etcdClient = initializeEtcdClient(logger, etcdOptions)
		bbsClient = initializeBBSClient(logger, bbsClientHTTPTimeout)
		etcdDB = initializeETCDDB(logger, etcdClient)
	}
})

var _ = SynchronizedAfterSuite(func() {
}, func() {
	purge("/v1/desired_lrp")
	purge("/v1/actual")
})

type ETCDFlags struct {
	etcdCertFile           string
	etcdKeyFile            string
	etcdCaFile             string
	clusterUrls            string
	clientSessionCacheSize int
	maxIdleConnsPerHost    int
}

func AddETCDFlags(flagSet *flag.FlagSet) *ETCDFlags {
	flags := &ETCDFlags{}

	flagSet.StringVar(
		&flags.clusterUrls,
		"etcdCluster",
		"http://127.0.0.1:4001",
		"comma-separated list of etcd URLs (scheme://ip:port)",
	)
	flagSet.StringVar(
		&flags.etcdCertFile,
		"etcdCertFile",
		"",
		"Location of the client certificate for mutual auth",
	)
	flagSet.StringVar(
		&flags.etcdKeyFile,
		"etcdKeyFile",
		"",
		"Location of the client key for mutual auth",
	)
	flagSet.StringVar(
		&flags.etcdCaFile,
		"etcdCaFile",
		"",
		"Location of the CA certificate for mutual auth",
	)

	flagSet.IntVar(
		&flags.clientSessionCacheSize,
		"etcdSessionCacheSize",
		0,
		"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
	)
	flagSet.IntVar(
		&flags.maxIdleConnsPerHost,
		"etcdMaxIdleConnsPerHost",
		0,
		"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
	)
	return flags
}

func (flags *ETCDFlags) Validate() (*etcddb.ETCDOptions, error) {
	scheme := ""
	clusterUrls := strings.Split(flags.clusterUrls, ",")
	for i, uString := range clusterUrls {
		uString = strings.TrimSpace(uString)
		clusterUrls[i] = uString
		u, err := url.Parse(uString)
		if err != nil {
			return nil, fmt.Errorf("Invalid cluster URL: '%s', error: [%s]", uString, err.Error())
		}
		if scheme == "" {
			if u.Scheme != "http" && u.Scheme != "https" {
				return nil, errors.New("Invalid scheme: " + uString)
			}
			scheme = u.Scheme
		} else if scheme != u.Scheme {
			return nil, fmt.Errorf("Multiple url schemes provided: %s", flags.clusterUrls)
		}
	}

	isSSL := false
	if scheme == "https" {
		isSSL = true
		if flags.etcdCertFile == "" {
			return nil, errors.New("Cert file must be provided for https connections")
		}
		if flags.etcdKeyFile == "" {
			return nil, errors.New("Key file must be provided for https connections")
		}
	}

	return &etcddb.ETCDOptions{
		CertFile:    flags.etcdCertFile,
		KeyFile:     flags.etcdKeyFile,
		CAFile:      flags.etcdCaFile,
		ClusterUrls: clusterUrls,
		IsSSL:       isSSL,
		ClientSessionCacheSize: flags.clientSessionCacheSize,
		MaxIdleConnsPerHost:    flags.maxIdleConnsPerHost,
	}, nil
}

func initializeEtcdClient(logger lager.Logger, etcdOptions *etcddb.ETCDOptions) *etcd.Client {
	var etcdClient *etcd.Client
	var tr *http.Transport

	if etcdOptions.IsSSL {
		if etcdOptions.CertFile == "" || etcdOptions.KeyFile == "" {
			logger.Fatal("failed-to-construct-etcd-tls-client", errors.New("Require both cert and key path"))
		}

		var err error
		etcdClient, err = etcd.NewTLSClient(etcdOptions.ClusterUrls, etcdOptions.CertFile, etcdOptions.KeyFile, etcdOptions.CAFile)
		if err != nil {
			logger.Fatal("failed-to-construct-etcd-tls-client", err)
		}

		tlsCert, err := tls.LoadX509KeyPair(etcdOptions.CertFile, etcdOptions.KeyFile)
		if err != nil {
			logger.Fatal("failed-to-construct-etcd-tls-client", err)
		}

		tlsConfig := &tls.Config{
			Certificates:       []tls.Certificate{tlsCert},
			InsecureSkipVerify: true,
			ClientSessionCache: tls.NewLRUClientSessionCache(etcdOptions.ClientSessionCacheSize),
		}
		tr = &http.Transport{
			TLSClientConfig:     tlsConfig,
			Dial:                etcdClient.DefaultDial,
			MaxIdleConnsPerHost: etcdOptions.MaxIdleConnsPerHost,
		}
		etcdClient.SetTransport(tr)
	} else {
		etcdClient = etcd.NewClient(etcdOptions.ClusterUrls)
	}
	etcdClient.SetConsistency(etcd.STRONG_CONSISTENCY)

	return etcdClient
}

func initializeETCDDB(logger lager.Logger, etcdClient *etcd.Client) *etcddb.ETCDDB {
	keyManager, err := encryptionFlags.Validate()
	Expect(err).NotTo(HaveOccurred())
	cryptor := encryption.NewCryptor(keyManager, rand.Reader)

	return etcddb.NewETCD(format.ENCRYPTED_PROTO, 1000, 1000, 1*time.Minute, cryptor, etcddb.NewStoreClient(etcdClient), nil, nil, clock.NewClock(), nil, nil)
}

func initializeBBSClient(logger lager.Logger, bbsClientHTTPTimeout time.Duration) bbs.Client {
	bbsURL, err := url.Parse(bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(bbsAddress)
	}

	cf_http.Initialize(bbsClientHTTPTimeout)
	bbsClient, err := bbs.NewSecureSkipVerifyClient(bbsAddress, bbsClientCert, bbsClientKey, 1, 1)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}

func purge(key string) {
	_, err := etcdClient.Delete(key, true)
	if err != nil {
		matches, matchErr := regexp.Match(".*Key not found.*", []byte(err.Error()))
		if matchErr != nil {
			Fail(matchErr.Error())
		}
		if !matches {
			Fail(err.Error())
		}
	}
}
<commit_msg>Use new encryptionFlags parsing method<commit_after>package benchmark_bbs_test

import (
	"crypto/rand"
	"crypto/tls"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"runtime"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/service/s3"
	"github.com/aws/aws-sdk-go/service/s3/s3manager"
	"github.com/cloudfoundry-incubator/bbs"
	etcddb "github.com/cloudfoundry-incubator/bbs/db/etcd"
	"github.com/cloudfoundry-incubator/bbs/encryption"
	"github.com/cloudfoundry-incubator/bbs/format"
	"github.com/cloudfoundry-incubator/benchmark-bbs/generator"
	"github.com/cloudfoundry-incubator/benchmark-bbs/reporter"
	"github.com/cloudfoundry-incubator/cf_http"
	"github.com/coreos/go-etcd/etcd"
	"github.com/pivotal-golang/clock"
	"github.com/pivotal-golang/lager"
	"github.com/zorkian/go-datadog-api"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"

	"testing"
)

var (
	bbsAddress           string
	bbsClientCert        string
	bbsClientKey         string
	etcdFlags            *ETCDFlags
	dataDogAPIKey        string
	dataDogAppKey        string
	awsAccessKeyID       string
	awsSecretAccessKey   string
	awsBucketName        string
	awsRegion            string
	desiredLRPs          int
	encryptionFlags      *encryption.EncryptionFlags
	metricPrefix         string
	numTrials            int
	numPopulateWorkers   int
	expectedLRPCount     int
	expectedLRPTolerance float64
	errorTolerance       float64
	logLevel             string
	logFilename          string

	logger               lager.Logger
	etcdClient           *etcd.Client
	etcdDB               *etcddb.ETCDDB
	bbsClient            bbs.Client
	bbsClientHTTPTimeout time.Duration
	dataDogClient        *datadog.Client
	dataDogReporter      reporter.DataDogReporter
	reporters            []Reporter
)

const (
	DEBUG = "debug"
	INFO  = "info"
	ERROR = "error"
	FATAL = "fatal"
)

func init() {
	runtime.GOMAXPROCS(runtime.NumCPU())
	flag.IntVar(&numTrials, "numTrials", 10, "number of benchmark trials to average across")
	flag.IntVar(&numPopulateWorkers, "numPopulateWorkers", 10, "number of workers generating desired LRPs during setup")
	flag.IntVar(&desiredLRPs, "desiredLRPs", 0, "number of single instance DesiredLRPs to create")

	flag.StringVar(&bbsAddress, "bbsAddress", "", "Address of the BBS Server")
	flag.StringVar(&bbsClientCert, "bbsClientCert", "", "BBS client SSL certificate")
	flag.StringVar(&bbsClientKey, "bbsClientKey", "", "BBS client SSL key")
	flag.DurationVar(&bbsClientHTTPTimeout, "bbsClientHTTPTimeout", 0, "BBS client HTTP timeout")

	flag.StringVar(&dataDogAPIKey, "dataDogAPIKey", "", "DataDog API key")
	flag.StringVar(&dataDogAppKey, "dataDogAppKey", "", "DataDog app Key")
	flag.StringVar(&metricPrefix, "metricPrefix", "", "DataDog metric prefix")

	flag.StringVar(&awsAccessKeyID, "awsAccessKeyID", "", "AWS Access Key ID")
	flag.StringVar(&awsSecretAccessKey, "awsSecretAccessKey", "", "AWS Secret Access Key")
	flag.StringVar(&awsBucketName, "awsBucketName", "", "AWS Bucket to store metrics")
	flag.StringVar(&awsRegion, "awsRegion", "us-west-1", "AWS Bucket to store metrics")

	flag.StringVar(&logLevel, "logLevel", string(INFO), "log level: debug, info, error or fatal")
	flag.StringVar(&logFilename, "logFilename", "", "Name of local file to save logs to")
	flag.Float64Var(&errorTolerance, "errorTolerance", 0.05, "error tollerance rate")

	etcdFlags = AddETCDFlags(flag.CommandLine)
	encryptionFlags = encryption.AddEncryptionFlags(flag.CommandLine)

	flag.Parse()

	if bbsAddress == "" {
		log.Fatal("bbsAddress is required")
	}

	BenchmarkConvergenceGathering(numTrials)
	BenchmarkNsyncFetching(numTrials)
	BenchmarkRouteEmitterFetching(numTrials)
}

func TestBenchmarkBbs(t *testing.T) {
	var lagerLogLevel lager.LogLevel
	switch logLevel {
	case DEBUG:
		lagerLogLevel = lager.DEBUG
	case INFO:
		lagerLogLevel = lager.INFO
	case ERROR:
		lagerLogLevel = lager.ERROR
	case FATAL:
		lagerLogLevel = lager.FATAL
	default:
		panic(fmt.Errorf("unknown log level: %s", logLevel))
	}

	var logWriter io.Writer
	if logFilename == "" {
		logWriter = GinkgoWriter
	} else {
		logFile, err := os.Create(logFilename)
		if err != nil {
			panic(fmt.Errorf("Error opening file '%s': %s", logFilename, err.Error()))
		}
		defer logFile.Close()

		logWriter = logFile
	}

	logger = lager.NewLogger("bbs-benchmarks-test")
	logger.RegisterSink(lager.NewWriterSink(logWriter, lagerLogLevel))

	reporters = []Reporter{}

	if dataDogAPIKey != "" && dataDogAppKey != "" {
		dataDogClient = datadog.NewClient(dataDogAPIKey, dataDogAppKey)
		dataDogReporter = reporter.NewDataDogReporter(logger, metricPrefix, dataDogClient)
		reporters = append(reporters, &dataDogReporter)
	}

	if awsAccessKeyID != "" && awsSecretAccessKey != "" && awsBucketName != "" {
		creds := credentials.NewStaticCredentials(awsAccessKeyID, awsSecretAccessKey, "")
		s3Client := s3.New(&aws.Config{
			Region:      &awsRegion,
			Credentials: creds,
		})
		uploader := s3manager.NewUploader(&s3manager.UploadOptions{S3: s3Client})
		reporter := reporter.NewS3Reporter(logger, awsBucketName, uploader)
		reporters = append(reporters, &reporter)
	}

	RegisterFailHandler(Fail)
	RunSpecsWithDefaultAndCustomReporters(t, "Benchmark BBS Suite", reporters)
}

var _ = SynchronizedBeforeSuite(func() []byte {
	etcdOptions, err := etcdFlags.Validate()
	Expect(err).NotTo(HaveOccurred())

	etcdClient = initializeEtcdClient(logger, etcdOptions)
	bbsClient = initializeBBSClient(logger, bbsClientHTTPTimeout)
	etcdDB = initializeETCDDB(logger, etcdClient)

	purge("/v1/desired_lrp")
	purge("/v1/actual")

	_, err = bbsClient.Domains()
	Expect(err).NotTo(HaveOccurred())
	if desiredLRPs > 0 {
		desiredLRPGenerator := generator.NewDesiredLRPGenerator(errorTolerance, metricPrefix, numPopulateWorkers, bbsClient, dataDogClient)
		expectedLRPCount, err = desiredLRPGenerator.Generate(logger, desiredLRPs)
		Expect(err).NotTo(HaveOccurred())
		expectedLRPTolerance = float64(expectedLRPCount) * errorTolerance
	}
	values := fmt.Sprintf("%d %f", expectedLRPCount, expectedLRPTolerance)
	return []byte(values)
}, func(data []byte) {
	values := string(data)
	fmt.Sscanf(values, "%d %f", &expectedLRPCount, &expectedLRPTolerance)
	if etcdClient == nil {
		etcdOptions, err := etcdFlags.Validate()
		Expect(err).NotTo(HaveOccurred())
		etcdClient = initializeEtcdClient(logger, etcdOptions)
		bbsClient = initializeBBSClient(logger, bbsClientHTTPTimeout)
		etcdDB = initializeETCDDB(logger, etcdClient)
	}
})

var _ = SynchronizedAfterSuite(func() {
}, func() {
	purge("/v1/desired_lrp")
	purge("/v1/actual")
})

type ETCDFlags struct {
	etcdCertFile           string
	etcdKeyFile            string
	etcdCaFile             string
	clusterUrls            string
	clientSessionCacheSize int
	maxIdleConnsPerHost    int
}

func AddETCDFlags(flagSet *flag.FlagSet) *ETCDFlags {
	flags := &ETCDFlags{}

	flagSet.StringVar(
		&flags.clusterUrls,
		"etcdCluster",
		"http://127.0.0.1:4001",
		"comma-separated list of etcd URLs (scheme://ip:port)",
	)
	flagSet.StringVar(
		&flags.etcdCertFile,
		"etcdCertFile",
		"",
		"Location of the client certificate for mutual auth",
	)
	flagSet.StringVar(
		&flags.etcdKeyFile,
		"etcdKeyFile",
		"",
		"Location of the client key for mutual auth",
	)
	flagSet.StringVar(
		&flags.etcdCaFile,
		"etcdCaFile",
		"",
		"Location of the CA certificate for mutual auth",
	)

	flagSet.IntVar(
		&flags.clientSessionCacheSize,
		"etcdSessionCacheSize",
		0,
		"Capacity of the ClientSessionCache option on the TLS configuration. If zero, golang's default will be used",
	)
	flagSet.IntVar(
		&flags.maxIdleConnsPerHost,
		"etcdMaxIdleConnsPerHost",
		0,
		"Controls the maximum number of idle (keep-alive) connctions per host. If zero, golang's default will be used",
	)
	return flags
}

func (flags *ETCDFlags) Validate() (*etcddb.ETCDOptions, error) {
	scheme := ""
	clusterUrls := strings.Split(flags.clusterUrls, ",")
	for i, uString := range clusterUrls {
		uString = strings.TrimSpace(uString)
		clusterUrls[i] = uString
		u, err := url.Parse(uString)
		if err != nil {
			return nil, fmt.Errorf("Invalid cluster URL: '%s', error: [%s]", uString, err.Error())
		}
		if scheme == "" {
			if u.Scheme != "http" && u.Scheme != "https" {
				return nil, errors.New("Invalid scheme: " + uString)
			}
			scheme = u.Scheme
		} else if scheme != u.Scheme {
			return nil, fmt.Errorf("Multiple url schemes provided: %s", flags.clusterUrls)
		}
	}

	isSSL := false
	if scheme == "https" {
		isSSL = true
		if flags.etcdCertFile == "" {
			return nil, errors.New("Cert file must be provided for https connections")
		}
		if flags.etcdKeyFile == "" {
			return nil, errors.New("Key file must be provided for https connections")
		}
	}

	return &etcddb.ETCDOptions{
		CertFile:    flags.etcdCertFile,
		KeyFile:     flags.etcdKeyFile,
		CAFile:      flags.etcdCaFile,
		ClusterUrls: clusterUrls,
		IsSSL:       isSSL,
		ClientSessionCacheSize: flags.clientSessionCacheSize,
		MaxIdleConnsPerHost:    flags.maxIdleConnsPerHost,
	}, nil
}

func initializeEtcdClient(logger lager.Logger, etcdOptions *etcddb.ETCDOptions) *etcd.Client {
	var etcdClient *etcd.Client
	var tr *http.Transport

	if etcdOptions.IsSSL {
		if etcdOptions.CertFile == "" || etcdOptions.KeyFile == "" {
			logger.Fatal("failed-to-construct-etcd-tls-client", errors.New("Require both cert and key path"))
		}

		var err error
		etcdClient, err = etcd.NewTLSClient(etcdOptions.ClusterUrls, etcdOptions.CertFile, etcdOptions.KeyFile, etcdOptions.CAFile)
		if err != nil {
			logger.Fatal("failed-to-construct-etcd-tls-client", err)
		}

		tlsCert, err := tls.LoadX509KeyPair(etcdOptions.CertFile, etcdOptions.KeyFile)
		if err != nil {
			logger.Fatal("failed-to-construct-etcd-tls-client", err)
		}

		tlsConfig := &tls.Config{
			Certificates:       []tls.Certificate{tlsCert},
			InsecureSkipVerify: true,
			ClientSessionCache: tls.NewLRUClientSessionCache(etcdOptions.ClientSessionCacheSize),
		}
		tr = &http.Transport{
			TLSClientConfig:     tlsConfig,
			Dial:                etcdClient.DefaultDial,
			MaxIdleConnsPerHost: etcdOptions.MaxIdleConnsPerHost,
		}
		etcdClient.SetTransport(tr)
	} else {
		etcdClient = etcd.NewClient(etcdOptions.ClusterUrls)
	}
	etcdClient.SetConsistency(etcd.STRONG_CONSISTENCY)

	return etcdClient
}

func initializeETCDDB(logger lager.Logger, etcdClient *etcd.Client) *etcddb.ETCDDB {
	key, keys, err := encryptionFlags.Parse()
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	keyManager, err := encryption.NewKeyManager(key, keys)
	if err != nil {
		logger.Fatal("cannot-setup-encryption", err)
	}
	cryptor := encryption.NewCryptor(keyManager, rand.Reader)

	return etcddb.NewETCD(format.ENCRYPTED_PROTO, 1000, 1000, 1*time.Minute, cryptor, etcddb.NewStoreClient(etcdClient), nil, nil, clock.NewClock(), nil, nil)
}

func initializeBBSClient(logger lager.Logger, bbsClientHTTPTimeout time.Duration) bbs.Client {
	bbsURL, err := url.Parse(bbsAddress)
	if err != nil {
		logger.Fatal("Invalid BBS URL", err)
	}

	if bbsURL.Scheme != "https" {
		return bbs.NewClient(bbsAddress)
	}

	cf_http.Initialize(bbsClientHTTPTimeout)
	bbsClient, err := bbs.NewSecureSkipVerifyClient(bbsAddress, bbsClientCert, bbsClientKey, 1, 1)
	if err != nil {
		logger.Fatal("Failed to configure secure BBS client", err)
	}
	return bbsClient
}

func purge(key string) {
	_, err := etcdClient.Delete(key, true)
	if err != nil {
		matches, matchErr := regexp.Match(".*Key not found.*", []byte(err.Error()))
		if matchErr != nil {
			Fail(matchErr.Error())
		}
		if !matches {
			Fail(err.Error())
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/MG-RAST/AWE/lib/auth"
	"github.com/MG-RAST/AWE/lib/conf"
	"github.com/MG-RAST/AWE/lib/controller"
	"github.com/MG-RAST/AWE/lib/core"
	"github.com/MG-RAST/AWE/lib/db"
	"github.com/MG-RAST/AWE/lib/logger"
	"github.com/MG-RAST/AWE/lib/logger/event"
	"github.com/MG-RAST/AWE/lib/user"
	"github.com/MG-RAST/golib/goweb"
	"os"
	"runtime"
)

func launchSite(control chan int, port int) {
	goweb.ConfigureDefaultFormatters()
	r := &goweb.RouteManager{}
	r.MapFunc("*", controller.SiteDir)
	if conf.SSL_ENABLED {
		err := goweb.ListenAndServeRoutesTLS(fmt.Sprintf(":%d", conf.SITE_PORT), conf.SSL_CERT_FILE, conf.SSL_KEY_FILE, r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: site: %v\n", err)
			logger.Error("ERROR: site: " + err.Error())
		}
	} else {
		err := goweb.ListenAndServeRoutes(fmt.Sprintf(":%d", conf.SITE_PORT), r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: site: %v\n", err)
			logger.Error("ERROR: site: " + err.Error())
		}
	}
	control <- 1 //we are ending
}

func launchAPI(control chan int, port int) {
	c := controller.NewServerController()
	goweb.ConfigureDefaultFormatters()
	r := &goweb.RouteManager{}
	r.MapRest("/job", c.Job)
	r.MapRest("/work", c.Work)
	r.MapRest("/client", c.Client)
	r.MapRest("/queue", c.Queue)
	r.MapRest("/awf", c.Awf)
	r.MapRest("/user", c.User)
	r.MapFunc("*", controller.ResourceDescription, goweb.GetMethod)
	if conf.SSL_ENABLED {
		err := goweb.ListenAndServeRoutesTLS(fmt.Sprintf(":%d", conf.API_PORT), conf.SSL_CERT_FILE, conf.SSL_KEY_FILE, r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: api: %v\n", err)
			logger.Error("ERROR: api: " + err.Error())
		}
	} else {
		err := goweb.ListenAndServeRoutes(fmt.Sprintf(":%d", conf.API_PORT), r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: api: %v\n", err)
			logger.Error("ERROR: api: " + err.Error())
		}
	}
	control <- 1 //we are ending
}

func main() {

	if !conf.INIT_SUCCESS {
		conf.PrintServerUsage()
		os.Exit(1)
	}

	if _, err := os.Stat(conf.DATA_PATH); err != nil && os.IsNotExist(err) {
		if err := os.MkdirAll(conf.DATA_PATH, 0777); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR in creating data_path %s\n", err.Error())
			os.Exit(1)
		}
	}

	if _, err := os.Stat(conf.LOGS_PATH); err != nil && os.IsNotExist(err) {
		if err := os.MkdirAll(conf.LOGS_PATH, 0777); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR in creating log_path %s\n", err.Error())
			os.Exit(1)
		}
	}

	if _, err := os.Stat(conf.DATA_PATH + "/temp"); err != nil && os.IsNotExist(err) {
		if err := os.Mkdir(conf.DATA_PATH+"/temp", 0777); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
			os.Exit(1)
		}
	}

	//init logger
	logger.Initialize("server")

	//init db
	if err := db.Initialize(); err != nil {
		fmt.Printf("failed to initialize job db: %s\n", err.Error())
		os.Exit(1)
	}

	//init db collection for user
	user.Initialize()

	//init resource manager
	core.InitResMgr("server")
	core.InitAwfMgr()
	core.InitJobDB()

	//init auth
	auth.Initialize()

	controller.PrintLogo()
	conf.Print("server")

	// reload job directory
	if conf.RELOAD != "" {
		fmt.Println("####### Reloading #######")
		if err := reload(conf.RELOAD); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		}
		fmt.Println("Done")
	}

	//init max job number (jid)
	if err := core.QMgr.InitMaxJid(); err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		os.Exit(1)
	}

	//recover unfinished jobs before server went down last time
	if conf.RECOVER {
		fmt.Println("####### Recovering unfinished jobs #######")
		if err := core.QMgr.RecoverJobs(); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		}
		fmt.Println("Done")
	}

	//launch server
	control := make(chan int)
	go core.QMgr.Handle()
	go core.QMgr.Timer()
	go core.QMgr.ClientChecker()
	go launchSite(control, conf.SITE_PORT)
	go launchAPI(control, conf.API_PORT)

	if err := core.AwfMgr.LoadWorkflows(); err != nil {
		logger.Error("LoadWorkflows: " + err.Error())
	}

	var host string
	if hostname, err := os.Hostname(); err == nil {
		host = fmt.Sprintf("%s:%d", hostname, conf.API_PORT)
	}
	if conf.RECOVER {
		logger.Event(event.SERVER_RECOVER, "host="+host)
	} else {
		logger.Event(event.SERVER_START, "host="+host)
	}

	if conf.PID_FILE_PATH != "" {
		f, err := os.Create(conf.PID_FILE_PATH)
		if err != nil {
			err_msg := "Could not create pid file: " + conf.PID_FILE_PATH + "\n"
			fmt.Fprintf(os.Stderr, err_msg)
			logger.Error("ERROR: " + err_msg)
			os.Exit(1)
		}
		defer f.Close()

		pid := os.Getpid()
		fmt.Fprintln(f, pid)

		fmt.Println("##### pidfile #####")
		fmt.Printf("pid: %d saved to file: %s\n\n", pid, conf.PID_FILE_PATH)
	}

	// setting GOMAXPROCS
	var procs int
	avail := runtime.NumCPU()
	if avail <= 2 {
		procs = 1
	} else if avail == 3 {
		procs = 2
	} else {
		procs = avail - 2
	}
	fmt.Println("##### Procs #####")
	fmt.Printf("Number of available CPUs = %d\n", avail)
	if conf.GOMAXPROCS > 0 {
		procs = conf.GOMAXPROCS
	}
	if procs <= avail {
		fmt.Printf("Running AWE server with GOMAXPROCS = %d\n\n", procs)
		runtime.GOMAXPROCS(procs)
	} else {
		fmt.Println("GOMAXPROCS config value is greater than available number of CPUs.")
		fmt.Printf("Running Shock server with GOMAXPROCS = %d\n\n", avail)
		runtime.GOMAXPROCS(avail)
	}

	<-control //block till something dies
}
<commit_msg>change the order of job recovery when server starts<commit_after>package main

import (
	"fmt"
	"github.com/MG-RAST/AWE/lib/auth"
	"github.com/MG-RAST/AWE/lib/conf"
	"github.com/MG-RAST/AWE/lib/controller"
	"github.com/MG-RAST/AWE/lib/core"
	"github.com/MG-RAST/AWE/lib/db"
	"github.com/MG-RAST/AWE/lib/logger"
	"github.com/MG-RAST/AWE/lib/logger/event"
	"github.com/MG-RAST/AWE/lib/user"
	"github.com/MG-RAST/golib/goweb"
	"os"
	"runtime"
)

func launchSite(control chan int, port int) {
	goweb.ConfigureDefaultFormatters()
	r := &goweb.RouteManager{}
	r.MapFunc("*", controller.SiteDir)
	if conf.SSL_ENABLED {
		err := goweb.ListenAndServeRoutesTLS(fmt.Sprintf(":%d", conf.SITE_PORT), conf.SSL_CERT_FILE, conf.SSL_KEY_FILE, r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: site: %v\n", err)
			logger.Error("ERROR: site: " + err.Error())
		}
	} else {
		err := goweb.ListenAndServeRoutes(fmt.Sprintf(":%d", conf.SITE_PORT), r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: site: %v\n", err)
			logger.Error("ERROR: site: " + err.Error())
		}
	}
	control <- 1 //we are ending
}

func launchAPI(control chan int, port int) {
	c := controller.NewServerController()
	goweb.ConfigureDefaultFormatters()
	r := &goweb.RouteManager{}
	r.MapRest("/job", c.Job)
	r.MapRest("/work", c.Work)
	r.MapRest("/client", c.Client)
	r.MapRest("/queue", c.Queue)
	r.MapRest("/awf", c.Awf)
	r.MapRest("/user", c.User)
	r.MapFunc("*", controller.ResourceDescription, goweb.GetMethod)
	if conf.SSL_ENABLED {
		err := goweb.ListenAndServeRoutesTLS(fmt.Sprintf(":%d", conf.API_PORT), conf.SSL_CERT_FILE, conf.SSL_KEY_FILE, r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: api: %v\n", err)
			logger.Error("ERROR: api: " + err.Error())
		}
	} else {
		err := goweb.ListenAndServeRoutes(fmt.Sprintf(":%d", conf.API_PORT), r)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: api: %v\n", err)
			logger.Error("ERROR: api: " + err.Error())
		}
	}
	control <- 1 //we are ending
}

func main() {

	if !conf.INIT_SUCCESS {
		conf.PrintServerUsage()
		os.Exit(1)
	}

	if _, err := os.Stat(conf.DATA_PATH); err != nil && os.IsNotExist(err) {
		if err := os.MkdirAll(conf.DATA_PATH, 0777); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR in creating data_path %s\n", err.Error())
			os.Exit(1)
		}
	}

	if _, err := os.Stat(conf.LOGS_PATH); err != nil && os.IsNotExist(err) {
		if err := os.MkdirAll(conf.LOGS_PATH, 0777); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR in creating log_path %s\n", err.Error())
			os.Exit(1)
		}
	}

	if _, err := os.Stat(conf.DATA_PATH + "/temp"); err != nil && os.IsNotExist(err) {
		if err := os.Mkdir(conf.DATA_PATH+"/temp", 0777); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
			os.Exit(1)
		}
	}

	//init logger
	logger.Initialize("server")

	//init db
	if err := db.Initialize(); err != nil {
		fmt.Printf("failed to initialize job db: %s\n", err.Error())
		os.Exit(1)
	}

	//init db collection for user
	user.Initialize()

	//init resource manager
	core.InitResMgr("server")
	core.InitAwfMgr()
	core.InitJobDB()

	//init auth
	auth.Initialize()

	controller.PrintLogo()
	conf.Print("server")

	// reload job directory
	if conf.RELOAD != "" {
		fmt.Println("####### Reloading #######")
		if err := reload(conf.RELOAD); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		}
		fmt.Println("Done")
	}

	//init max job number (jid)
	if err := core.QMgr.InitMaxJid(); err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		os.Exit(1)
	}

	//launch server
	control := make(chan int)
	go core.QMgr.Handle()
	go core.QMgr.Timer()
	go core.QMgr.ClientChecker()
	go launchSite(control, conf.SITE_PORT)
	go launchAPI(control, conf.API_PORT)

	if err := core.AwfMgr.LoadWorkflows(); err != nil {
		logger.Error("LoadWorkflows: " + err.Error())
	}

	var host string
	if hostname, err := os.Hostname(); err == nil {
		host = fmt.Sprintf("%s:%d", hostname, conf.API_PORT)
	}

	//recover unfinished jobs before server went down last time
	if conf.RECOVER {
		fmt.Println("####### Recovering unfinished jobs #######")
		if err := core.QMgr.RecoverJobs(); err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %v\n", err)
		}
		fmt.Println("Done")
		logger.Event(event.SERVER_RECOVER, "host="+host)
	} else {
		logger.Event(event.SERVER_START, "host="+host)
	}

	if conf.PID_FILE_PATH != "" {
		f, err := os.Create(conf.PID_FILE_PATH)
		if err != nil {
			err_msg := "Could not create pid file: " + conf.PID_FILE_PATH + "\n"
			fmt.Fprintf(os.Stderr, err_msg)
			logger.Error("ERROR: " + err_msg)
			os.Exit(1)
		}
		defer f.Close()

		pid := os.Getpid()
		fmt.Fprintln(f, pid)

		fmt.Println("##### pidfile #####")
		fmt.Printf("pid: %d saved to file: %s\n\n", pid, conf.PID_FILE_PATH)
	}

	// setting GOMAXPROCS
	var procs int
	avail := runtime.NumCPU()
	if avail <= 2 {
		procs = 1
	} else if avail == 3 {
		procs = 2
	} else {
		procs = avail - 2
	}
	fmt.Println("##### Procs #####")
	fmt.Printf("Number of available CPUs = %d\n", avail)
	if conf.GOMAXPROCS > 0 {
		procs = conf.GOMAXPROCS
	}
	if procs <= avail {
		fmt.Printf("Running AWE server with GOMAXPROCS = %d\n\n", procs)
		runtime.GOMAXPROCS(procs)
	} else {
		fmt.Println("GOMAXPROCS config value is greater than available number of CPUs.")
		fmt.Printf("Running Shock server with GOMAXPROCS = %d\n\n", avail)
		runtime.GOMAXPROCS(avail)
	}

	<-control //block till something dies
}
<|endoftext|>


<commit_before>// Copyright 2017 The LUCI 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 venv

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"time"

	"go.chromium.org/luci/vpython/api/vpython"
	"go.chromium.org/luci/vpython/python"
	"go.chromium.org/luci/vpython/spec"
	"go.chromium.org/luci/vpython/venv/assets"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/system/environ"
	"go.chromium.org/luci/common/system/filesystem"
)

// EnvironmentVersion is an environment version string. It must advance each
// time the layout of a VirtualEnv environment changes.
const EnvironmentVersion = "v4"

const siteCustomizePy = "sitecustomize.py"

// Config is the configuration for a managed VirtualEnv.
//
// A VirtualEnv is specified based on its resolved vpython.Spec.
type Config struct {
	// MaxHashLen is the maximum number of hash characters to use in VirtualEnv
	// directory names.
	MaxHashLen int

	// BaseDir is the parent directory of all VirtualEnv.
	BaseDir string

	// SetupEnv if not nil, is the base environment that will be used during
	// VirtualEnv setup operations.
	SetupEnv environ.Env

	// OverrideName overrides the name of the specified VirtualEnv.
	//
	// Because the name is no longer derived from the specification, this will
	// force revalidation and deletion of any existing content if it is not a
	// fully defined and matching VirtualEnv
	OverrideName string

	// Package is the VirtualEnv package to install. It must be non-nil and
	// valid. It will be used if the environment specification doesn't supply an
	// overriding one.
	Package vpython.Spec_Package

	// Python is the Python interpreter to use. If empty, one will be resolved
	// based on the Spec and the current PATH.
	Python string

	// LookPathFunc, if not nil, will be used instead of exec.LookPath to find the
	// underlying Python interpreter.
	LookPathFunc python.LookPathFunc

	// Spec is the specification file to use to construct the VirtualEnv. If
	// nil, or if fields are missing, they will be filled in by probing the system
	// PATH.
	Spec *vpython.Spec

	// PruneThreshold, if >0, is the maximum age of a VirtualEnv before it should
	// be pruned. If <= 0, there is no maximum age, so no pruning will be
	// performed.
	PruneThreshold time.Duration
	// MaxPrunesPerSweep applies a limit to the number of items to prune per
	// execution.
	//
	// If <= 0, no limit will be applied.
	MaxPrunesPerSweep int

	// Loader is the PackageLoader instance to use for package resolution and
	// deployment.
	Loader PackageLoader

	// FailIfLocked, if true, means that if a lock is encountered during
	// VirtualEnv operation, the VirtualEnv will exit immediately with an error.
	// If false, the VirtualEnv will block pending the lock's availability.
	FailIfLocked bool

	// si is the system Python interpreter. It is resolved during
	// "resolvePythonInterpreter".
	si *python.Interpreter

	// rt is the resolved Python runtime.
	rt vpython.Runtime

	// testPreserveInstallationCapability is a testing parameter. If true, the
	// VirtualEnv's ability to install will be preserved after the setup. This is
	// used by the test whell generation bootstrap code.
	testPreserveInstallationCapability bool

	// testLeaveReadWrite, if true, instructs the VirtualEnv setup to leave the
	// directory read/write. This makes it easier to manage, and is safe since it
	// is not a production directory.
	testLeaveReadWrite bool
}

// WithoutWheels returns a clone of cfg that depends on no additional packages.
//
// If cfg is already an empty it will be returned directly.
func (cfg *Config) WithoutWheels() *Config {
	if !cfg.HasWheels() {
		return cfg
	}

	clone := *cfg
	clone.OverrideName = ""
	clone.Spec = clone.Spec.Clone()
	clone.Spec.Wheel = nil
	return &clone
}

// HasWheels returns true if this environment declares wheel dependencies.
func (cfg *Config) HasWheels() bool {
	return cfg.Spec != nil && len(cfg.Spec.Wheel) > 0
}

// makeEnv processes the config, validating and, where appropriate, populating
// any components. Upon success, it returns a configured Env instance.
//
// The supplied vpython.Environment is derived externally, and may be nil if
// this is a bootstrapped Environment.
//
// The returned Env instance may or may not actually exist. Setup must be called
// prior to using it.
func (cfg *Config) makeEnv(c context.Context, e *vpython.Environment) (*Env, error) {
	// We MUST have a package loader.
	if cfg.Loader == nil {
		return nil, errors.New("no package loader provided")
	}

	// Resolve our base directory, if one is not supplied.
	if cfg.BaseDir == "" {
		// Use one in a temporary directory.
		cfg.BaseDir = filepath.Join(os.TempDir(), "vpython")
		logging.Debugf(c, "Using tempdir-relative environment root: %s", cfg.BaseDir)
	}
	if err := filesystem.AbsPath(&cfg.BaseDir); err != nil {
		return nil, errors.Annotate(err, "failed to resolve absolute path of base directory").Err()
	}

	// Construct a new, independent Environment for this Env.
	e = e.Clone()
	if cfg.Spec != nil {
		e.Spec = cfg.Spec.Clone()
	}
	if err := spec.NormalizeEnvironment(e); err != nil {
		return nil, errors.Annotate(err, "invalid environment").Err()
	}

	// If the environment doesn't specify a VirtualEnv package (expected), use
	// our default.
	if e.Spec.Virtualenv == nil {
		e.Spec.Virtualenv = &cfg.Package
	}

	if err := cfg.Loader.Resolve(c, e); err != nil {
		return nil, errors.Annotate(err, "failed to resolve packages").Err()
	}

	if err := cfg.resolvePythonInterpreter(c, e.Spec); err != nil {
		return nil, errors.Annotate(err, "failed to resolve system Python interpreter").Err()
	}

	e.Runtime = &vpython.Runtime{
		Version: e.Spec.PythonVersion,
	}
	if err := fillRuntime(c, cfg.si, e.Runtime); err != nil {
		return nil, err
	}
	logging.Debugf(c, "Resolved system Python runtime: %#s", e.Runtime)

	// Ensure that our base directory exists.
	if err := filesystem.MakeDirs(cfg.BaseDir); err != nil {
		return nil, errors.Annotate(err, "could not create environment root: %s", cfg.BaseDir).Err()
	}

	// Generate our environment name based on the deterministic hash of its
	// fully-resolved specification.
	envName := cfg.OverrideName
	if envName == "" {
		envName = cfg.envNameForSpec(c, e.Spec, e.Runtime)
	}
	env := cfg.envForName(envName, e)
	return env, nil
}

// EnvName returns the VirtualEnv environment name for the environment that cfg
// describes.
func (cfg *Config) envNameForSpec(c context.Context, s *vpython.Spec, rt *vpython.Runtime) string {
	uid := "<unknown>"
	if user, err := user.Current(); err != nil {
		uid = user.Uid
	} else {
		logging.Warningf(c, "Unable to find current user ID: %s", err)
	}

	name := spec.Hash(s, rt, EnvironmentVersion,
		string(assets.GetAssetSHA256(siteCustomizePy)), uid)
	if cfg.MaxHashLen > 0 && len(name) > cfg.MaxHashLen {
		name = name[:cfg.MaxHashLen]
	}
	return name
}

// Prune performs a pruning round on the environment set described by this
// Config.
func (cfg *Config) Prune(c context.Context) error {
	if err := prune(c, cfg, nil); err != nil {
		return errors.Annotate(err, "").Err()
	}
	return nil
}

// envForName creates an Env for a named directory.
//
// The Environment, e, can be nil; however, code paths that require it may not
// be called.
func (cfg *Config) envForName(name string, e *vpython.Environment) *Env {
	// Env-specific root directory: <BaseDir>/<name>
	venvRoot := filepath.Join(cfg.BaseDir, name)
	binDir := venvBinDir(venvRoot)
	return &Env{
		Config:               cfg,
		Root:                 venvRoot,
		Name:                 name,
		Python:               filepath.Join(binDir, "python"),
		Environment:          e,
		BinDir:               binDir,
		EnvironmentStampPath: filepath.Join(venvRoot, fmt.Sprintf("environment.%s.pb.txt", vpython.Version)),
		lockPath:             filepath.Join(cfg.BaseDir, fmt.Sprintf(".%s.lock", name)),

		completeFlagPath: filepath.Join(venvRoot, "complete.flag"),
	}
}

func (cfg *Config) resolvePythonInterpreter(c context.Context, s *vpython.Spec) error {
	specVers, err := python.ParseVersion(s.PythonVersion)
	if err != nil {
		return errors.Annotate(err, "failed to parse Python version from: %q", s.PythonVersion).Err()
	}

	if cfg.Python == "" {
		// No explicitly-specified Python path. Determine one based on the
		// specification.
		if cfg.si, err = python.Find(c, specVers, cfg.LookPathFunc); err != nil {
			return errors.Annotate(err, "could not find Python for: %s", specVers).Err()
		}
		cfg.Python = cfg.si.Python
	} else {
		cfg.si = &python.Interpreter{
			Python: cfg.Python,
		}
	}
	if err := cfg.si.Normalize(); err != nil {
		return err
	}

	// Confirm that the version of the interpreter matches that which is
	// expected.
	interpreterVers, err := cfg.si.GetVersion(c)
	if err != nil {
		return errors.Annotate(err, "failed to determine Python version for: %s", cfg.Python).Err()
	}
	if !specVers.IsSatisfiedBy(interpreterVers) {
		return errors.Reason("supplied Python version (%s) doesn't match specification (%s)", interpreterVers, specVers).Err()
	}
	s.PythonVersion = interpreterVers.String()
	return nil
}

func (cfg *Config) systemInterpreter() *python.Interpreter { return cfg.si }

// fillRuntime returns the runtime information of the specified interpreter.
//
// Identifying this information requires the interpreter to be run.
func fillRuntime(c context.Context, i *python.Interpreter, r *vpython.Runtime) error {
	// JSON fields correspond to "vpython.Runtime" fields.
	const script = `` +
		`import json, sys;` +
		`json.dump({` +
		`'path': sys.executable,` +
		`'prefix': sys.prefix,` +
		`}, sys.stdout)`

	// Probe the runtime information.
	var stdout, stderr bytes.Buffer
	cmd := i.IsolatedCommand(c, python.CommandTarget{Command: script})
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		logging.WithError(err).Errorf(c, "Failed to get runtime information:\n%s", stderr.Bytes())
		return errors.Annotate(err, "failed to get runtime information").Err()
	}

	if err := json.Unmarshal(stdout.Bytes(), r); err != nil {
		return errors.Annotate(err, "could not unmarshal output: %q", stdout.Bytes()).Err()
	}

	// "sys.executable" is allowed to be None. If it is, use the Python
	// interpreter path.
	//
	// Resolve it to an absolute path. "sys.executable" says that it must be one,
	// but we'll enforce it just to be sure.
	if r.Path == "" {
		r.Path = i.Python
	}
	if err := filesystem.AbsPath(&r.Path); err != nil {
		return err
	}

	var err error
	if r.Hash, err = hashPath(r.Path); err != nil {
		return err
	}

	return nil
}

func hashPath(path string) (string, error) {
	fd, err := os.Open(path)
	if err != nil {
		return "", errors.Annotate(err, "failed to open interpreter").Err()
	}
	defer fd.Close()

	hash := sha256.New()
	if _, err := io.Copy(hash, fd); err != nil {
		return "", errors.Annotate(err, "failed to read [%s] for hashing", path).Err()
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}
<commit_msg>[vpython] Correctly handle error.<commit_after>// Copyright 2017 The LUCI 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 venv

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"os/user"
	"path/filepath"
	"time"

	"go.chromium.org/luci/vpython/api/vpython"
	"go.chromium.org/luci/vpython/python"
	"go.chromium.org/luci/vpython/spec"
	"go.chromium.org/luci/vpython/venv/assets"

	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/common/system/environ"
	"go.chromium.org/luci/common/system/filesystem"
)

// EnvironmentVersion is an environment version string. It must advance each
// time the layout of a VirtualEnv environment changes.
const EnvironmentVersion = "v4"

const siteCustomizePy = "sitecustomize.py"

// Config is the configuration for a managed VirtualEnv.
//
// A VirtualEnv is specified based on its resolved vpython.Spec.
type Config struct {
	// MaxHashLen is the maximum number of hash characters to use in VirtualEnv
	// directory names.
	MaxHashLen int

	// BaseDir is the parent directory of all VirtualEnv.
	BaseDir string

	// SetupEnv if not nil, is the base environment that will be used during
	// VirtualEnv setup operations.
	SetupEnv environ.Env

	// OverrideName overrides the name of the specified VirtualEnv.
	//
	// Because the name is no longer derived from the specification, this will
	// force revalidation and deletion of any existing content if it is not a
	// fully defined and matching VirtualEnv
	OverrideName string

	// Package is the VirtualEnv package to install. It must be non-nil and
	// valid. It will be used if the environment specification doesn't supply an
	// overriding one.
	Package vpython.Spec_Package

	// Python is the Python interpreter to use. If empty, one will be resolved
	// based on the Spec and the current PATH.
	Python string

	// LookPathFunc, if not nil, will be used instead of exec.LookPath to find the
	// underlying Python interpreter.
	LookPathFunc python.LookPathFunc

	// Spec is the specification file to use to construct the VirtualEnv. If
	// nil, or if fields are missing, they will be filled in by probing the system
	// PATH.
	Spec *vpython.Spec

	// PruneThreshold, if >0, is the maximum age of a VirtualEnv before it should
	// be pruned. If <= 0, there is no maximum age, so no pruning will be
	// performed.
	PruneThreshold time.Duration
	// MaxPrunesPerSweep applies a limit to the number of items to prune per
	// execution.
	//
	// If <= 0, no limit will be applied.
	MaxPrunesPerSweep int

	// Loader is the PackageLoader instance to use for package resolution and
	// deployment.
	Loader PackageLoader

	// FailIfLocked, if true, means that if a lock is encountered during
	// VirtualEnv operation, the VirtualEnv will exit immediately with an error.
	// If false, the VirtualEnv will block pending the lock's availability.
	FailIfLocked bool

	// si is the system Python interpreter. It is resolved during
	// "resolvePythonInterpreter".
	si *python.Interpreter

	// rt is the resolved Python runtime.
	rt vpython.Runtime

	// testPreserveInstallationCapability is a testing parameter. If true, the
	// VirtualEnv's ability to install will be preserved after the setup. This is
	// used by the test whell generation bootstrap code.
	testPreserveInstallationCapability bool

	// testLeaveReadWrite, if true, instructs the VirtualEnv setup to leave the
	// directory read/write. This makes it easier to manage, and is safe since it
	// is not a production directory.
	testLeaveReadWrite bool
}

// WithoutWheels returns a clone of cfg that depends on no additional packages.
//
// If cfg is already an empty it will be returned directly.
func (cfg *Config) WithoutWheels() *Config {
	if !cfg.HasWheels() {
		return cfg
	}

	clone := *cfg
	clone.OverrideName = ""
	clone.Spec = clone.Spec.Clone()
	clone.Spec.Wheel = nil
	return &clone
}

// HasWheels returns true if this environment declares wheel dependencies.
func (cfg *Config) HasWheels() bool {
	return cfg.Spec != nil && len(cfg.Spec.Wheel) > 0
}

// makeEnv processes the config, validating and, where appropriate, populating
// any components. Upon success, it returns a configured Env instance.
//
// The supplied vpython.Environment is derived externally, and may be nil if
// this is a bootstrapped Environment.
//
// The returned Env instance may or may not actually exist. Setup must be called
// prior to using it.
func (cfg *Config) makeEnv(c context.Context, e *vpython.Environment) (*Env, error) {
	// We MUST have a package loader.
	if cfg.Loader == nil {
		return nil, errors.New("no package loader provided")
	}

	// Resolve our base directory, if one is not supplied.
	if cfg.BaseDir == "" {
		// Use one in a temporary directory.
		cfg.BaseDir = filepath.Join(os.TempDir(), "vpython")
		logging.Debugf(c, "Using tempdir-relative environment root: %s", cfg.BaseDir)
	}
	if err := filesystem.AbsPath(&cfg.BaseDir); err != nil {
		return nil, errors.Annotate(err, "failed to resolve absolute path of base directory").Err()
	}

	// Construct a new, independent Environment for this Env.
	e = e.Clone()
	if cfg.Spec != nil {
		e.Spec = cfg.Spec.Clone()
	}
	if err := spec.NormalizeEnvironment(e); err != nil {
		return nil, errors.Annotate(err, "invalid environment").Err()
	}

	// If the environment doesn't specify a VirtualEnv package (expected), use
	// our default.
	if e.Spec.Virtualenv == nil {
		e.Spec.Virtualenv = &cfg.Package
	}

	if err := cfg.Loader.Resolve(c, e); err != nil {
		return nil, errors.Annotate(err, "failed to resolve packages").Err()
	}

	if err := cfg.resolvePythonInterpreter(c, e.Spec); err != nil {
		return nil, errors.Annotate(err, "failed to resolve system Python interpreter").Err()
	}

	e.Runtime = &vpython.Runtime{
		Version: e.Spec.PythonVersion,
	}
	if err := fillRuntime(c, cfg.si, e.Runtime); err != nil {
		return nil, err
	}
	logging.Debugf(c, "Resolved system Python runtime: %#s", e.Runtime)

	// Ensure that our base directory exists.
	if err := filesystem.MakeDirs(cfg.BaseDir); err != nil {
		return nil, errors.Annotate(err, "could not create environment root: %s", cfg.BaseDir).Err()
	}

	// Generate our environment name based on the deterministic hash of its
	// fully-resolved specification.
	envName := cfg.OverrideName
	if envName == "" {
		envName = cfg.envNameForSpec(c, e.Spec, e.Runtime)
	}
	env := cfg.envForName(envName, e)
	return env, nil
}

// EnvName returns the VirtualEnv environment name for the environment that cfg
// describes.
func (cfg *Config) envNameForSpec(c context.Context, s *vpython.Spec, rt *vpython.Runtime) string {
	uid := "<unknown>"
	if user, err := user.Current(); err == nil {
		uid = user.Uid
	} else {
		logging.Warningf(c, "Unable to find current user ID: %s", err)
	}

	name := spec.Hash(s, rt, EnvironmentVersion,
		string(assets.GetAssetSHA256(siteCustomizePy)), uid)
	if cfg.MaxHashLen > 0 && len(name) > cfg.MaxHashLen {
		name = name[:cfg.MaxHashLen]
	}
	return name
}

// Prune performs a pruning round on the environment set described by this
// Config.
func (cfg *Config) Prune(c context.Context) error {
	if err := prune(c, cfg, nil); err != nil {
		return errors.Annotate(err, "").Err()
	}
	return nil
}

// envForName creates an Env for a named directory.
//
// The Environment, e, can be nil; however, code paths that require it may not
// be called.
func (cfg *Config) envForName(name string, e *vpython.Environment) *Env {
	// Env-specific root directory: <BaseDir>/<name>
	venvRoot := filepath.Join(cfg.BaseDir, name)
	binDir := venvBinDir(venvRoot)
	return &Env{
		Config:               cfg,
		Root:                 venvRoot,
		Name:                 name,
		Python:               filepath.Join(binDir, "python"),
		Environment:          e,
		BinDir:               binDir,
		EnvironmentStampPath: filepath.Join(venvRoot, fmt.Sprintf("environment.%s.pb.txt", vpython.Version)),
		lockPath:             filepath.Join(cfg.BaseDir, fmt.Sprintf(".%s.lock", name)),

		completeFlagPath: filepath.Join(venvRoot, "complete.flag"),
	}
}

func (cfg *Config) resolvePythonInterpreter(c context.Context, s *vpython.Spec) error {
	specVers, err := python.ParseVersion(s.PythonVersion)
	if err != nil {
		return errors.Annotate(err, "failed to parse Python version from: %q", s.PythonVersion).Err()
	}

	if cfg.Python == "" {
		// No explicitly-specified Python path. Determine one based on the
		// specification.
		if cfg.si, err = python.Find(c, specVers, cfg.LookPathFunc); err != nil {
			return errors.Annotate(err, "could not find Python for: %s", specVers).Err()
		}
		cfg.Python = cfg.si.Python
	} else {
		cfg.si = &python.Interpreter{
			Python: cfg.Python,
		}
	}
	if err := cfg.si.Normalize(); err != nil {
		return err
	}

	// Confirm that the version of the interpreter matches that which is
	// expected.
	interpreterVers, err := cfg.si.GetVersion(c)
	if err != nil {
		return errors.Annotate(err, "failed to determine Python version for: %s", cfg.Python).Err()
	}
	if !specVers.IsSatisfiedBy(interpreterVers) {
		return errors.Reason("supplied Python version (%s) doesn't match specification (%s)", interpreterVers, specVers).Err()
	}
	s.PythonVersion = interpreterVers.String()
	return nil
}

func (cfg *Config) systemInterpreter() *python.Interpreter { return cfg.si }

// fillRuntime returns the runtime information of the specified interpreter.
//
// Identifying this information requires the interpreter to be run.
func fillRuntime(c context.Context, i *python.Interpreter, r *vpython.Runtime) error {
	// JSON fields correspond to "vpython.Runtime" fields.
	const script = `` +
		`import json, sys;` +
		`json.dump({` +
		`'path': sys.executable,` +
		`'prefix': sys.prefix,` +
		`}, sys.stdout)`

	// Probe the runtime information.
	var stdout, stderr bytes.Buffer
	cmd := i.IsolatedCommand(c, python.CommandTarget{Command: script})
	cmd.Stdout = &stdout
	if err := cmd.Run(); err != nil {
		logging.WithError(err).Errorf(c, "Failed to get runtime information:\n%s", stderr.Bytes())
		return errors.Annotate(err, "failed to get runtime information").Err()
	}

	if err := json.Unmarshal(stdout.Bytes(), r); err != nil {
		return errors.Annotate(err, "could not unmarshal output: %q", stdout.Bytes()).Err()
	}

	// "sys.executable" is allowed to be None. If it is, use the Python
	// interpreter path.
	//
	// Resolve it to an absolute path. "sys.executable" says that it must be one,
	// but we'll enforce it just to be sure.
	if r.Path == "" {
		r.Path = i.Python
	}
	if err := filesystem.AbsPath(&r.Path); err != nil {
		return err
	}

	var err error
	if r.Hash, err = hashPath(r.Path); err != nil {
		return err
	}

	return nil
}

func hashPath(path string) (string, error) {
	fd, err := os.Open(path)
	if err != nil {
		return "", errors.Annotate(err, "failed to open interpreter").Err()
	}
	defer fd.Close()

	hash := sha256.New()
	if _, err := io.Copy(hash, fd); err != nil {
		return "", errors.Annotate(err, "failed to read [%s] for hashing", path).Err()
	}
	return hex.EncodeToString(hash.Sum(nil)), nil
}
<|endoftext|>


<commit_before>// Copyright 2022 Google 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 telemetry_port_speed_test implements tests that cover port-speed related
// telemetry variables.
package telemetry_port_speed_test

import (
	"fmt"
	"sort"
	"testing"
	"time"

	"github.com/openconfig/featureprofiles/internal/attrs"
	"github.com/openconfig/featureprofiles/internal/deviations"
	"github.com/openconfig/featureprofiles/internal/fptest"
	"github.com/openconfig/ondatra"
	"github.com/openconfig/ondatra/netutil"
	"github.com/openconfig/ondatra/telemetry"
	"github.com/openconfig/ygot/ygot"
)

func TestMain(m *testing.M) {
	fptest.RunTests(m)
}

// Test cases:
//   - Validate that the port speed reported in telemetry is the expected port speed.
//   - Turn port down at ATE, validate that operational status of the port is reported as down.
//   - Connect N ports between ATE and DUT configured as part of a LACP bundle.
//     Validate that the effective speed of the LAG is reported as N*port speed.
//   - Disable each port at ATE and determine that the effective speed is reduced by the expected amount.
//   - Turn ports sequentially up at the ATE, and determine that the effective speed is increased as expected.
//
// Topology:
//
//	dut:port1 <--> ate:port1
//	dut:portN <--> ate:portN
const (
	plen4 = 30
	plen6 = 126
)

var (
	dutIPs = attrs.Attributes{
		Name:    "dutip",
		IPv4:    "192.0.2.5",
		IPv6:    "2001:db8::5",
		IPv4Len: plen4,
		IPv6Len: plen6,
	}

	ateIPs = attrs.Attributes{
		Name:    "ateip",
		IPv4:    "192.0.2.6",
		IPv6:    "2001:db8::6",
		IPv4Len: plen4,
		IPv6Len: plen6,
	}
)

const (
	lagTypeLACP   = telemetry.IfAggregate_AggregationType_LACP
	lagTypeSTATIC = telemetry.IfAggregate_AggregationType_STATIC
	minLink       = 1
)

type testCase struct {
	minlinks uint16
	lagType  telemetry.E_IfAggregate_AggregationType

	dut *ondatra.DUTDevice
	ate *ondatra.ATEDevice
	top *ondatra.ATETopology

	dutPorts []*ondatra.Port
	atePorts []*ondatra.Port
	aggID    string
}

func (*testCase) configDUT(i *telemetry.Interface, a *attrs.Attributes) {
	i.Description = ygot.String(a.Desc)
	if *deviations.InterfaceEnabled {
		i.Enabled = ygot.Bool(true)
	}

	s := i.GetOrCreateSubinterface(0)
	s4 := s.GetOrCreateIpv4()
	if *deviations.InterfaceEnabled {
		s4.Enabled = ygot.Bool(true)
	}
	s4.GetOrCreateAddress(a.IPv4).PrefixLength = ygot.Uint8((plen4))

	s6 := s.GetOrCreateIpv6()
	if *deviations.InterfaceEnabled {
		s6.Enabled = ygot.Bool(true)
	}
	s6.GetOrCreateAddress(a.IPv6).PrefixLength = ygot.Uint8(plen6)
}

func (tc *testCase) configAggregateDUT(i *telemetry.Interface, a *attrs.Attributes) {
	tc.configDUT(i, a)
	g := i.GetOrCreateAggregation()
	g.LagType = tc.lagType
	g.MinLinks = ygot.Uint16(tc.minlinks)
}

var portSpeed = map[ondatra.Speed]telemetry.E_IfEthernet_ETHERNET_SPEED{
	ondatra.Speed10Gb:  telemetry.IfEthernet_ETHERNET_SPEED_SPEED_10GB,
	ondatra.Speed100Gb: telemetry.IfEthernet_ETHERNET_SPEED_SPEED_100GB,
	ondatra.Speed400Gb: telemetry.IfEthernet_ETHERNET_SPEED_SPEED_400GB,
}

func (tc *testCase) configMemberDUT(i *telemetry.Interface, p *ondatra.Port) {
	i.Description = ygot.String(p.String())
	if *deviations.InterfaceEnabled {
		i.Enabled = ygot.Bool(true)
	}
	e := i.GetOrCreateEthernet()
	e.AggregateId = ygot.String(tc.aggID)
}

func (tc *testCase) setupAggregateAtomically(t *testing.T) {
	d := &telemetry.Device{}

	if tc.lagType == lagTypeLACP {
		d.GetOrCreateLacp().GetOrCreateInterface(tc.aggID)
	}

	agg := d.GetOrCreateInterface(tc.aggID)
	agg.GetOrCreateAggregation().LagType = tc.lagType

	for _, port := range tc.dutPorts {
		i := d.GetOrCreateInterface(port.Name())
		i.GetOrCreateEthernet().AggregateId = ygot.String(tc.aggID)
	}

	p := tc.dut.Config()
	fptest.LogYgot(t, fmt.Sprintf("%s to Update()", tc.dut), p, d)
	p.Update(t, d)
}

func (tc *testCase) clearAggregateMembers(t *testing.T) {
	for _, port := range tc.dutPorts {
		tc.dut.Config().Interface(port.Name()).Ethernet().AggregateId().Delete(t)
	}
}

// sortPorts sorts the ports by the testbed port ID.
func sortPorts(ports []*ondatra.Port) []*ondatra.Port {
	sort.SliceStable(ports, func(i, j int) bool {
		return ports[i].ID() < ports[j].ID()
	})
	return ports
}

func (tc *testCase) configureDUT(t *testing.T) {
	t.Logf("dut ports = %v", tc.dutPorts)
	if len(tc.dutPorts) < 2 {
		t.Fatalf("Testbed requires at least 2 ports, got %d", len(tc.dutPorts))
	}

	d := tc.dut.Config()

	if *deviations.AggregateAtomicUpdate {
		tc.clearAggregateMembers(t)
		tc.setupAggregateAtomically(t)
	}

	if tc.lagType == lagTypeLACP {
		lacp := &telemetry.Lacp_Interface{Name: ygot.String(tc.aggID)}
		lacp.LacpMode = telemetry.Lacp_LacpActivityType_ACTIVE

		lacpPath := d.Lacp().Interface(tc.aggID)
		fptest.LogYgot(t, "LACP", lacpPath, lacp)
		lacpPath.Replace(t, lacp)
		t.Cleanup(func() {
			lacpPath.Delete(t)
		})
	}

	agg := &telemetry.Interface{Name: ygot.String(tc.aggID)}
	tc.configAggregateDUT(agg, &dutIPs)
	aggPath := d.Interface(tc.aggID)
	fptest.LogYgot(t, tc.aggID, aggPath, agg)
	aggPath.Replace(t, agg)
	t.Cleanup(func() {
		aggPath.Delete(t)
	})

	for _, port := range tc.dutPorts {
		iName := port.Name()
		i := &telemetry.Interface{Name: ygot.String(iName)}
		tc.configMemberDUT(i, port)
		iPath := d.Interface(iName)
		fptest.LogYgot(t, port.String(), iPath, i)
		iPath.Replace(t, i)
		t.Cleanup(func() {
			iPath.Replace(t, &telemetry.Interface{Name: ygot.String(iName)})
		})
	}
}

func (tc *testCase) configureATE(t *testing.T) {
	if len(tc.atePorts) < 2 {
		t.Fatalf("Testbed requires at least 2 ports, got: %v", tc.atePorts)
	}

	// Don't use WithLACPEnabled which is for emulated Ixia LACP.
	agg := tc.top.AddInterface(ateIPs.Name)
	lag := tc.top.AddLAG("lag").WithPorts(tc.atePorts...)
	lag.LACP().WithEnabled(tc.lagType == lagTypeLACP)
	agg.WithLAG(lag)

	// Disable FEC for 100G-FR ports because Novus does not support it.
	is100gfr := false
	for _, p := range tc.atePorts {
		if p.PMD() == ondatra.PMD100GFR {
			is100gfr = true
		}
	}
	if is100gfr {
		agg.Ethernet().FEC().WithEnabled(false)
	}

	agg.IPv4().
		WithAddress(ateIPs.IPv4CIDR()).
		WithDefaultGateway(dutIPs.IPv4)
	agg.IPv6().
		WithAddress(ateIPs.IPv6CIDR()).
		WithDefaultGateway(dutIPs.IPv6)

	tc.top.Push(t).StartProtocols(t)
}

func (tc *testCase) verifyDUT(t *testing.T, numPort int) {
	dutPort := tc.dut.Port(t, "port1")
	want := int(dutPort.Speed()) * numPort * 1000
	val, _ := tc.dut.Telemetry().Interface(tc.aggID).Aggregation().LagSpeed().Watch(t, 30*time.Second, func(val *telemetry.QualifiedUint32) bool { return val.IsPresent() }).Await(t)
	if got := int(val.Val(t)); got != want {
		t.Errorf("Get(DUT port status): got %v, want %v", got, want)
	}
}

func TestGNMIPortSpeed(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dutPort := dut.Port(t, "port1")
	if got, want := dut.Telemetry().Interface(dutPort.Name()).Ethernet().PortSpeed().Get(t), portSpeed[dutPort.Speed()]; got != want {
		t.Errorf("Get(DUT port1 status): got %v, want %v", got, want)
	}
}

func TestGNMIPortDown(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	ate := ondatra.ATE(t, "ate")
	dutPort := dut.Port(t, "port1")
	atePort := ate.Port(t, "port1")
	top := ate.Topology().New()
	intf := top.AddInterface(ateIPs.Name).WithPort(atePort)
	intf.IPv4().
		WithAddress(ateIPs.IPv4CIDR()).
		WithDefaultGateway(dutIPs.IPv4)
	intf.IPv6().
		WithAddress(ateIPs.IPv6CIDR()).
		WithDefaultGateway(dutIPs.IPv6)
	top.Push(t)
	ate.Actions().NewSetPortState().WithPort(atePort).WithEnabled(false).Send(t)
	dutPortStatus := dut.Telemetry().Interface(dutPort.Name()).OperStatus().Get(t)

	if want := telemetry.Interface_OperStatus_DOWN; dutPortStatus != want {
		t.Errorf("Get(DUT port1 status): got %v, want %v", dutPortStatus, want)
	}
	ate.Actions().NewSetPortState().WithPort(atePort).WithEnabled(true).Send(t)
}

func TestGNMICombinedLACPSpeed(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	ate := ondatra.ATE(t, "ate")

	for _, lagType := range []telemetry.E_IfAggregate_AggregationType{lagTypeLACP, lagTypeSTATIC} {
		t.Run(lagType.String(), func(t *testing.T) {
			top := ate.Topology().New()
			tc := &testCase{
				minlinks: minLink,
				lagType:  lagType,

				dut: dut,
				ate: ate,
				top: top,

				dutPorts: sortPorts(dut.Ports()),
				atePorts: sortPorts(ate.Ports()),
				aggID:    netutil.NextBundleInterface(t, dut),
			}
			tc.configureDUT(t)
			tc.configureATE(t)
			tc.verifyDUT(t, len(tc.dutPorts))
		})
	}
}

func TestGNMIReducedLACPSpeed(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	ate := ondatra.ATE(t, "ate")
	totalPort := len(ate.Ports())

	for _, lagType := range []telemetry.E_IfAggregate_AggregationType{lagTypeLACP, lagTypeSTATIC} {
		t.Run(lagType.String(), func(t *testing.T) {
			top := ate.Topology().New()
			tc := &testCase{
				minlinks: minLink,
				lagType:  lagType,
				dut:      dut,
				ate:      ate,
				top:      top,

				dutPorts: sortPorts(dut.Ports()),
				atePorts: sortPorts(ate.Ports()),
				aggID:    netutil.NextBundleInterface(t, dut),
			}
			tc.configureDUT(t)
			tc.configureATE(t)
			for _, port := range tc.atePorts {
				totalPort--
				if totalPort < 1 {
					break
				}
				ate.Actions().NewSetPortState().WithPort(port).WithEnabled(false).Send(t)
				time.Sleep(10 * time.Second)
				tc.verifyDUT(t, totalPort)
			}
			for _, port := range tc.atePorts {
				totalPort++
				if totalPort > len(tc.atePorts)-1 {
					break
				}
				ate.Actions().NewSetPortState().WithPort(port).WithEnabled(true).Send(t)
				time.Sleep(10 * time.Second)
				tc.verifyDUT(t, totalPort+1)
			}
		})
	}
}
<commit_msg>Adding code changes for https://github.com/openconfig/featureprofiles/issues/493 (#494)<commit_after>// Copyright 2022 Google 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 telemetry_port_speed_test implements tests that cover port-speed related
// telemetry variables.
package telemetry_port_speed_test

import (
	"fmt"
	"sort"
	"testing"
	"time"

	"github.com/openconfig/featureprofiles/internal/attrs"
	"github.com/openconfig/featureprofiles/internal/deviations"
	"github.com/openconfig/featureprofiles/internal/fptest"
	"github.com/openconfig/ondatra"
	"github.com/openconfig/ondatra/netutil"
	"github.com/openconfig/ondatra/telemetry"
	"github.com/openconfig/ygot/ygot"
)

func TestMain(m *testing.M) {
	fptest.RunTests(m)
}

// Test cases:
//   - Validate that the port speed reported in telemetry is the expected port speed.
//   - Turn port down at ATE, validate that operational status of the port is reported as down.
//   - Connect N ports between ATE and DUT configured as part of a LACP bundle.
//     Validate that the effective speed of the LAG is reported as N*port speed.
//   - Disable each port at ATE and determine that the effective speed is reduced by the expected amount.
//   - Turn ports sequentially up at the ATE, and determine that the effective speed is increased as expected.
//
// Topology:
//
//	dut:port1 <--> ate:port1
//	dut:portN <--> ate:portN
const (
	plen4 = 30
	plen6 = 126
)

var (
	dutIPs = attrs.Attributes{
		Name:    "dutip",
		IPv4:    "192.0.2.5",
		IPv6:    "2001:db8::5",
		IPv4Len: plen4,
		IPv6Len: plen6,
	}

	ateIPs = attrs.Attributes{
		Name:    "ateip",
		IPv4:    "192.0.2.6",
		IPv6:    "2001:db8::6",
		IPv4Len: plen4,
		IPv6Len: plen6,
	}
)

const (
	ethernetCsmacd = telemetry.IETFInterfaces_InterfaceType_ethernetCsmacd
	ieee8023adLag  = telemetry.IETFInterfaces_InterfaceType_ieee8023adLag
	lagTypeLACP    = telemetry.IfAggregate_AggregationType_LACP
	lagTypeSTATIC  = telemetry.IfAggregate_AggregationType_STATIC
	minLink        = 1
)

type testCase struct {
	minlinks uint16
	lagType  telemetry.E_IfAggregate_AggregationType

	dut *ondatra.DUTDevice
	ate *ondatra.ATEDevice
	top *ondatra.ATETopology

	dutPorts []*ondatra.Port
	atePorts []*ondatra.Port
	aggID    string
}

func (*testCase) configDUT(i *telemetry.Interface, a *attrs.Attributes) {
	i.Description = ygot.String(a.Desc)
	if *deviations.InterfaceEnabled {
		i.Enabled = ygot.Bool(true)
	}

	s := i.GetOrCreateSubinterface(0)
	s4 := s.GetOrCreateIpv4()
	if *deviations.InterfaceEnabled {
		s4.Enabled = ygot.Bool(true)
	}
	s4.GetOrCreateAddress(a.IPv4).PrefixLength = ygot.Uint8((plen4))

	s6 := s.GetOrCreateIpv6()
	if *deviations.InterfaceEnabled {
		s6.Enabled = ygot.Bool(true)
	}
	s6.GetOrCreateAddress(a.IPv6).PrefixLength = ygot.Uint8(plen6)
}

func (tc *testCase) configAggregateDUT(i *telemetry.Interface, a *attrs.Attributes) {
	tc.configDUT(i, a)
	i.Type = ieee8023adLag
	g := i.GetOrCreateAggregation()
	g.LagType = tc.lagType
	g.MinLinks = ygot.Uint16(tc.minlinks)
}

var portSpeed = map[ondatra.Speed]telemetry.E_IfEthernet_ETHERNET_SPEED{
	ondatra.Speed10Gb:  telemetry.IfEthernet_ETHERNET_SPEED_SPEED_10GB,
	ondatra.Speed100Gb: telemetry.IfEthernet_ETHERNET_SPEED_SPEED_100GB,
	ondatra.Speed400Gb: telemetry.IfEthernet_ETHERNET_SPEED_SPEED_400GB,
}

func (tc *testCase) configMemberDUT(i *telemetry.Interface, p *ondatra.Port) {
	i.Description = ygot.String(p.String())
	i.Type = ethernetCsmacd
	if *deviations.InterfaceEnabled {
		i.Enabled = ygot.Bool(true)
	}
	e := i.GetOrCreateEthernet()
	e.AggregateId = ygot.String(tc.aggID)
}

func (tc *testCase) setupAggregateAtomically(t *testing.T) {
	d := &telemetry.Device{}

	if tc.lagType == lagTypeLACP {
		d.GetOrCreateLacp().GetOrCreateInterface(tc.aggID)
	}

	agg := d.GetOrCreateInterface(tc.aggID)
	agg.GetOrCreateAggregation().LagType = tc.lagType

	for _, port := range tc.dutPorts {
		i := d.GetOrCreateInterface(port.Name())
		i.GetOrCreateEthernet().AggregateId = ygot.String(tc.aggID)
		i.Type = ethernetCsmacd
	}

	p := tc.dut.Config()
	fptest.LogYgot(t, fmt.Sprintf("%s to Update()", tc.dut), p, d)
	p.Update(t, d)
}

func (tc *testCase) clearAggregateMembers(t *testing.T) {
	for _, port := range tc.dutPorts {
		tc.dut.Config().Interface(port.Name()).Ethernet().AggregateId().Delete(t)
	}
}

// sortPorts sorts the ports by the testbed port ID.
func sortPorts(ports []*ondatra.Port) []*ondatra.Port {
	sort.SliceStable(ports, func(i, j int) bool {
		return ports[i].ID() < ports[j].ID()
	})
	return ports
}

func (tc *testCase) configureDUT(t *testing.T) {
	t.Logf("dut ports = %v", tc.dutPorts)
	if len(tc.dutPorts) < 2 {
		t.Fatalf("Testbed requires at least 2 ports, got %d", len(tc.dutPorts))
	}

	d := tc.dut.Config()

	if *deviations.AggregateAtomicUpdate {
		tc.clearAggregateMembers(t)
		tc.setupAggregateAtomically(t)
	}

	if tc.lagType == lagTypeLACP {
		lacp := &telemetry.Lacp_Interface{Name: ygot.String(tc.aggID)}
		lacp.LacpMode = telemetry.Lacp_LacpActivityType_ACTIVE

		lacpPath := d.Lacp().Interface(tc.aggID)
		fptest.LogYgot(t, "LACP", lacpPath, lacp)
		lacpPath.Replace(t, lacp)
		t.Cleanup(func() {
			lacpPath.Delete(t)
		})
	}

	agg := &telemetry.Interface{Name: ygot.String(tc.aggID)}
	tc.configAggregateDUT(agg, &dutIPs)
	aggPath := d.Interface(tc.aggID)
	fptest.LogYgot(t, tc.aggID, aggPath, agg)
	aggPath.Replace(t, agg)
	t.Cleanup(func() {
		aggPath.Delete(t)
	})

	for _, port := range tc.dutPorts {
		iName := port.Name()
		i := &telemetry.Interface{Name: ygot.String(iName)}
		tc.configMemberDUT(i, port)
		iPath := d.Interface(iName)
		fptest.LogYgot(t, port.String(), iPath, i)
		iPath.Replace(t, i)
		t.Cleanup(func() {
			iPath.Replace(t, &telemetry.Interface{Name: ygot.String(iName)})
		})
	}
}

func (tc *testCase) configureATE(t *testing.T) {
	if len(tc.atePorts) < 2 {
		t.Fatalf("Testbed requires at least 2 ports, got: %v", tc.atePorts)
	}

	// Don't use WithLACPEnabled which is for emulated Ixia LACP.
	agg := tc.top.AddInterface(ateIPs.Name)
	lag := tc.top.AddLAG("lag").WithPorts(tc.atePorts...)
	lag.LACP().WithEnabled(tc.lagType == lagTypeLACP)
	agg.WithLAG(lag)

	// Disable FEC for 100G-FR ports because Novus does not support it.
	is100gfr := false
	for _, p := range tc.atePorts {
		if p.PMD() == ondatra.PMD100GFR {
			is100gfr = true
		}
	}
	if is100gfr {
		agg.Ethernet().FEC().WithEnabled(false)
	}

	agg.IPv4().
		WithAddress(ateIPs.IPv4CIDR()).
		WithDefaultGateway(dutIPs.IPv4)
	agg.IPv6().
		WithAddress(ateIPs.IPv6CIDR()).
		WithDefaultGateway(dutIPs.IPv6)

	tc.top.Push(t).StartProtocols(t)
}

func (tc *testCase) verifyDUT(t *testing.T, numPort int) {
	dutPort := tc.dut.Port(t, "port1")
	want := int(dutPort.Speed()) * numPort * 1000
	val, _ := tc.dut.Telemetry().Interface(tc.aggID).Aggregation().LagSpeed().Watch(t, 30*time.Second, func(val *telemetry.QualifiedUint32) bool { return val.IsPresent() }).Await(t)
	if got := int(val.Val(t)); got != want {
		t.Errorf("Get(DUT port status): got %v, want %v", got, want)
	}
}

func TestGNMIPortSpeed(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	dutPort := dut.Port(t, "port1")
	if got, want := dut.Telemetry().Interface(dutPort.Name()).Ethernet().PortSpeed().Get(t), portSpeed[dutPort.Speed()]; got != want {
		t.Errorf("Get(DUT port1 status): got %v, want %v", got, want)
	}
}

func TestGNMIPortDown(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	ate := ondatra.ATE(t, "ate")
	dutPort := dut.Port(t, "port1")
	atePort := ate.Port(t, "port1")
	top := ate.Topology().New()
	intf := top.AddInterface(ateIPs.Name).WithPort(atePort)
	intf.IPv4().
		WithAddress(ateIPs.IPv4CIDR()).
		WithDefaultGateway(dutIPs.IPv4)
	intf.IPv6().
		WithAddress(ateIPs.IPv6CIDR()).
		WithDefaultGateway(dutIPs.IPv6)
	top.Push(t)
	ate.Actions().NewSetPortState().WithPort(atePort).WithEnabled(false).Send(t)
	dutPortStatus := dut.Telemetry().Interface(dutPort.Name()).OperStatus().Get(t)

	if want := telemetry.Interface_OperStatus_DOWN; dutPortStatus != want {
		t.Errorf("Get(DUT port1 status): got %v, want %v", dutPortStatus, want)
	}
	ate.Actions().NewSetPortState().WithPort(atePort).WithEnabled(true).Send(t)
}

func TestGNMICombinedLACPSpeed(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	ate := ondatra.ATE(t, "ate")

	for _, lagType := range []telemetry.E_IfAggregate_AggregationType{lagTypeLACP, lagTypeSTATIC} {
		t.Run(lagType.String(), func(t *testing.T) {
			top := ate.Topology().New()
			tc := &testCase{
				minlinks: minLink,
				lagType:  lagType,

				dut: dut,
				ate: ate,
				top: top,

				dutPorts: sortPorts(dut.Ports()),
				atePorts: sortPorts(ate.Ports()),
				aggID:    netutil.NextBundleInterface(t, dut),
			}
			tc.configureDUT(t)
			tc.configureATE(t)
			tc.verifyDUT(t, len(tc.dutPorts))
		})
	}
}

func TestGNMIReducedLACPSpeed(t *testing.T) {
	dut := ondatra.DUT(t, "dut")
	ate := ondatra.ATE(t, "ate")
	totalPort := len(ate.Ports())

	for _, lagType := range []telemetry.E_IfAggregate_AggregationType{lagTypeLACP, lagTypeSTATIC} {
		t.Run(lagType.String(), func(t *testing.T) {
			top := ate.Topology().New()
			tc := &testCase{
				minlinks: minLink,
				lagType:  lagType,
				dut:      dut,
				ate:      ate,
				top:      top,

				dutPorts: sortPorts(dut.Ports()),
				atePorts: sortPorts(ate.Ports()),
				aggID:    netutil.NextBundleInterface(t, dut),
			}
			tc.configureDUT(t)
			tc.configureATE(t)
			for _, port := range tc.atePorts {
				totalPort--
				if totalPort < 1 {
					break
				}
				ate.Actions().NewSetPortState().WithPort(port).WithEnabled(false).Send(t)
				time.Sleep(10 * time.Second)
				tc.verifyDUT(t, totalPort)
			}
			for _, port := range tc.atePorts {
				totalPort++
				if totalPort > len(tc.atePorts)-1 {
					break
				}
				ate.Actions().NewSetPortState().WithPort(port).WithEnabled(true).Send(t)
				time.Sleep(10 * time.Second)
				tc.verifyDUT(t, totalPort+1)
			}
		})
	}
}
<|endoftext|>


<commit_before>package text

import (
	"testing"
)

func TestWord_Match(t *testing.T) {
	cases := []struct {
		s   string
		q   word
		exp int
	}{
		{
			s:   "foo",
			q:   Word("foo"),
			exp: 0,
		},
		{
			s:   "foo bar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo barbaz",
			q:   Word("bar"),
			exp: -1,
		},
		{
			s:   "foobar baz",
			q:   Word("bar"),
			exp: -1,
		},
		{
			s:   "foo bar",
			q:   Word("FOO"),
			exp: -1,
		},
		{
			s:   "FOO BAR",
			q:   Word("foo"),
			exp: -1,
		},
		{
			s:   "foo\tbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\nbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\vbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\fbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\rbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "å bar",
			q:   Word("bar"),
			exp: 2,
		},
	}

	for _, c := range cases {
		act := c.q.Match(c.s)

		if act != c.exp {
			t.Errorf("Expected Word(%q).Match(%q) to return %v, got %v", c.q.W, c.s, c.exp, act)
		}
	}
}

// BenchmarkWord_MatchNotExist6_587 benchmarks a single Word query of length
// 6 against a search text of length 587, where the sought string does not
// exist in the search text.
//
// This benchmark gives a good indication of the average performance of an
// unsuccessful search.
func BenchmarkWord_MatchNotExist6_587(b *testing.B) {
	w := Word("foobar")

	txt := `Lorem ipsum dolor sit amet, an cum vero soleat concludaturque, te purto vero reprimique vis.
	Ignota mediocritatem ut sea. Cetero deserunt pericula te vel. Omnis legendos no per.
	Sale illum pertinax no sed, est posse putent minimum no. Pri et vitae mentitum eligendi,
	no ius reque fugit libris, eos ad quaeque pericula mediocrem. Habemus corpora an mea,
	inermis partiendo per et, at nemore dolorem iudicabit eos. At est mucius docendi. Sed et nisl facilisi.
	Idque suavitate argumentum eu eam, vis putant insolens dissentiunt id. Dictas labitur in mei, duo omnium assentior scripserit cu omnium`

	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		w.Match(txt)
	}
}

// BenchmarkWord_MatchPartials5_587 benchmarks a single Word query of length
// 6 against a search text of length 587 with partials of length 5.
//
// This benchmark gives a good indication of the worst case performance
// of an unsuccessful search.
func BenchmarkWord_MatchPartials5_587(b *testing.B) {
	w := Word("foobar")

	txt := `fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba`

	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		w.Match(txt)
	}
}

// BenchmarkWord_MatchExist6_587 benchmarks a single Word query of length
// 6 against a search text of length 587, where the sought string is at the
// very end of the search text.
func BenchmarkWord_MatchExist6_587(b *testing.B) {
	w := Word("foobar")

	txt := `Lorem ipsum dolor sit amet, an cum vero soleat concludaturque, te purto vero reprimique vis.
	Ignota mediocritatem ut sea. Cetero deserunt pericula te vel. Omnis legendos no per.
	Sale illum pertinax no sed, est posse putent minimum no. Pri et vitae mentitum eligendi,
	no ius reque fugit libris, eos ad quaeque pericula mediocrem. Habemus corpora an mea,
	inermis partiendo per et, at nemore dolorem iudicabit eos. At est mucius docendi. Sed et nisl facilisi.
	Idque suavitate argumentum eu eam, vis putant insolens dissentiunt id. Dictas labitur in mei, duo omnium assentior scripserit cu foobar`

	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		w.Match(txt)
	}
}
<commit_msg>Improve UTF-8 test<commit_after>package text

import (
	"testing"
)

func TestWord_Match(t *testing.T) {
	cases := []struct {
		s   string
		q   word
		exp int
	}{
		{
			s:   "foo",
			q:   Word("foo"),
			exp: 0,
		},
		{
			s:   "foo bar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo barbaz",
			q:   Word("bar"),
			exp: -1,
		},
		{
			s:   "foobar baz",
			q:   Word("bar"),
			exp: -1,
		},
		{
			s:   "foo bar",
			q:   Word("FOO"),
			exp: -1,
		},
		{
			s:   "FOO BAR",
			q:   Word("foo"),
			exp: -1,
		},
		{
			s:   "foo\tbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\nbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\vbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\fbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "foo\rbar",
			q:   Word("bar"),
			exp: 4,
		},
		{
			s:   "åäö bar",
			q:   Word("bar"),
			exp: 4,
		},
	}

	for _, c := range cases {
		act := c.q.Match(c.s)

		if act != c.exp {
			t.Errorf("Expected Word(%q).Match(%q) to return %v, got %v", c.q.W, c.s, c.exp, act)
		}
	}
}

// BenchmarkWord_MatchNotExist6_587 benchmarks a single Word query of length
// 6 against a search text of length 587, where the sought string does not
// exist in the search text.
//
// This benchmark gives a good indication of the average performance of an
// unsuccessful search.
func BenchmarkWord_MatchNotExist6_587(b *testing.B) {
	w := Word("foobar")

	txt := `Lorem ipsum dolor sit amet, an cum vero soleat concludaturque, te purto vero reprimique vis.
	Ignota mediocritatem ut sea. Cetero deserunt pericula te vel. Omnis legendos no per.
	Sale illum pertinax no sed, est posse putent minimum no. Pri et vitae mentitum eligendi,
	no ius reque fugit libris, eos ad quaeque pericula mediocrem. Habemus corpora an mea,
	inermis partiendo per et, at nemore dolorem iudicabit eos. At est mucius docendi. Sed et nisl facilisi.
	Idque suavitate argumentum eu eam, vis putant insolens dissentiunt id. Dictas labitur in mei, duo omnium assentior scripserit cu omnium`

	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		w.Match(txt)
	}
}

// BenchmarkWord_MatchPartials5_587 benchmarks a single Word query of length
// 6 against a search text of length 587 with partials of length 5.
//
// This benchmark gives a good indication of the worst case performance
// of an unsuccessful search.
func BenchmarkWord_MatchPartials5_587(b *testing.B) {
	w := Word("foobar")

	txt := `fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba fooba fooba
		fooba fooba fooba fooba fooba fooba fooba fooba`

	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		w.Match(txt)
	}
}

// BenchmarkWord_MatchExist6_587 benchmarks a single Word query of length
// 6 against a search text of length 587, where the sought string is at the
// very end of the search text.
func BenchmarkWord_MatchExist6_587(b *testing.B) {
	w := Word("foobar")

	txt := `Lorem ipsum dolor sit amet, an cum vero soleat concludaturque, te purto vero reprimique vis.
	Ignota mediocritatem ut sea. Cetero deserunt pericula te vel. Omnis legendos no per.
	Sale illum pertinax no sed, est posse putent minimum no. Pri et vitae mentitum eligendi,
	no ius reque fugit libris, eos ad quaeque pericula mediocrem. Habemus corpora an mea,
	inermis partiendo per et, at nemore dolorem iudicabit eos. At est mucius docendi. Sed et nisl facilisi.
	Idque suavitate argumentum eu eam, vis putant insolens dissentiunt id. Dictas labitur in mei, duo omnium assentior scripserit cu foobar`

	b.ResetTimer()

	for n := 0; n < b.N; n++ {
		w.Match(txt)
	}
}
<|endoftext|>