<commit_before>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

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

	"github.com/google/syzkaller/pkg/cover"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/rpctype"
	"github.com/google/syzkaller/pkg/signal"
	"github.com/google/syzkaller/prog"
)

type RPCServer struct {
	mgr                   RPCManagerView
	port                  int
	target                *prog.Target
	configEnabledSyscalls []int
	targetEnabledSyscalls map[*prog.Syscall]bool
	stats                 *Stats
	sandbox               string
	batchSize             int

	mu           sync.Mutex
	fuzzers      map[string]*Fuzzer
	checkResult  *rpctype.CheckArgs
	maxSignal    signal.Signal
	corpusSignal signal.Signal
	corpusCover  cover.Cover
	rotator      *prog.Rotator
	rnd          *rand.Rand
}

type Fuzzer struct {
	name          string
	inputs        []rpctype.RPCInput
	newMaxSignal  signal.Signal
	rotatedSignal signal.Signal
	machineInfo   []byte
}

type BugFrames struct {
	memoryLeaks []string
	dataRaces   []string
}

// RPCManagerView restricts interface between RPCServer and Manager.
type RPCManagerView interface {
	fuzzerConnect() ([]rpctype.RPCInput, BugFrames)
	machineChecked(result *rpctype.CheckArgs, enabledSyscalls map[*prog.Syscall]bool)
	newInput(inp rpctype.RPCInput, sign signal.Signal) bool
	candidateBatch(size int) []rpctype.RPCCandidate
	rotateCorpus() bool
}

func startRPCServer(mgr *Manager) (*RPCServer, error) {
	serv := &RPCServer{
		mgr:                   mgr,
		target:                mgr.target,
		configEnabledSyscalls: mgr.configEnabledSyscalls,
		stats:                 mgr.stats,
		sandbox:               mgr.cfg.Sandbox,
		fuzzers:               make(map[string]*Fuzzer),
		rnd:                   rand.New(rand.NewSource(time.Now().UnixNano())),
	}
	serv.batchSize = 5
	if serv.batchSize < mgr.cfg.Procs {
		serv.batchSize = mgr.cfg.Procs
	}
	s, err := rpctype.NewRPCServer(mgr.cfg.RPC, "Manager", serv)
	if err != nil {
		return nil, err
	}
	log.Logf(0, "serving rpc on tcp://%v", s.Addr())
	serv.port = s.Addr().(*net.TCPAddr).Port
	go s.Serve()
	return serv, nil
}

func (serv *RPCServer) Connect(a *rpctype.ConnectArgs, r *rpctype.ConnectRes) error {
	log.Logf(1, "fuzzer %v connected", a.Name)
	serv.stats.vmRestarts.inc()

	corpus, bugFrames := serv.mgr.fuzzerConnect()

	serv.mu.Lock()
	defer serv.mu.Unlock()

	f := &Fuzzer{
		name:        a.Name,
		machineInfo: a.MachineInfo,
	}
	serv.fuzzers[a.Name] = f
	r.MemoryLeakFrames = bugFrames.memoryLeaks
	r.DataRaceFrames = bugFrames.dataRaces
	r.EnabledCalls = serv.configEnabledSyscalls
	r.GitRevision = prog.GitRevision
	r.TargetRevision = serv.target.Revision
	// TODO: temporary disabled b/c we suspect this negatively affects fuzzing.
	if false && serv.mgr.rotateCorpus() && serv.rnd.Intn(3) != 0 {
		// We do rotation every other time because there are no objective
		// proofs regarding its efficiency either way.
		// Also, rotation gives significantly skewed syscall selection
		// (run prog.TestRotationCoverage), it may or may not be OK.
		r.CheckResult = serv.rotateCorpus(f, corpus)
	} else {
		r.CheckResult = serv.checkResult
		f.inputs = corpus
		f.newMaxSignal = serv.maxSignal.Copy()
	}
	return nil
}

func (serv *RPCServer) rotateCorpus(f *Fuzzer, corpus []rpctype.RPCInput) *rpctype.CheckArgs {
	// Fuzzing tends to stuck in some local optimum and then it fails to cover
	// other state space points since code coverage is only a very approximate
	// measure of logic coverage. To overcome this we introduce some variation
	// into the process which should cause steady corpus rotation over time
	// (the same coverage is achieved in different ways).
	//
	// First, we select a subset of all syscalls for each VM run (result.EnabledCalls).
	// This serves 2 goals: (1) target fuzzer at a particular area of state space,
	// (2) disable syscalls that cause frequent crashes at least in some runs
	// to allow it to do actual fuzzing.
	//
	// Then, we remove programs that contain disabled syscalls from corpus
	// that will be sent to the VM (f.inputs). We also remove 10% of remaining
	// programs at random to allow to rediscover different variations of these programs.
	//
	// Then, we drop signal provided by the removed programs and also 10%
	// of the remaining signal at random (f.newMaxSignal). This again allows
	// rediscovery of this signal by different programs.
	//
	// Finally, we adjust criteria for accepting new programs from this VM (f.rotatedSignal).
	// This allows to accept rediscovered varied programs even if they don't
	// increase overall coverage. As the result we have multiple programs
	// providing the same duplicate coverage, these are removed during periodic
	// corpus minimization process. The minimization process is specifically
	// non-deterministic to allow the corpus rotation.
	//
	// Note: at no point we drop anything globally and permanently.
	// Everything we remove during this process is temporal and specific to a single VM.
	calls := serv.rotator.Select()

	var callIDs []int
	callNames := make(map[string]bool)
	for call := range calls {
		callNames[call.Name] = true
		callIDs = append(callIDs, call.ID)
	}

	f.inputs, f.newMaxSignal = serv.selectInputs(callNames, corpus, serv.maxSignal)
	// Remove the corresponding signal from rotatedSignal which will
	// be used to accept new inputs from this manager.
	f.rotatedSignal = serv.corpusSignal.Intersection(f.newMaxSignal)

	result := *serv.checkResult
	result.EnabledCalls = map[string][]int{serv.sandbox: callIDs}
	return &result
}

func (serv *RPCServer) selectInputs(enabled map[string]bool, inputs0 []rpctype.RPCInput, signal0 signal.Signal) (
	inputs []rpctype.RPCInput, signal signal.Signal) {
	signal = signal0.Copy()
	for _, inp := range inputs0 {
		calls, _, err := prog.CallSet(inp.Prog)
		if err != nil {
			panic(fmt.Sprintf("rotateInputs: CallSet failed: %v\n%s", err, inp.Prog))
		}
		for call := range calls {
			if !enabled[call] {
				goto drop
			}
		}
		if serv.rnd.Float64() > 0.9 {
			goto drop
		}
		inputs = append(inputs, inp)
		continue
	drop:
		for _, sig := range inp.Signal.Elems {
			delete(signal, sig)
		}
	}
	signal.Split(len(signal) / 10)
	return inputs, signal
}

func (serv *RPCServer) Check(a *rpctype.CheckArgs, r *int) error {
	serv.mu.Lock()
	defer serv.mu.Unlock()

	if serv.checkResult != nil {
		return nil
	}
	serv.targetEnabledSyscalls = make(map[*prog.Syscall]bool)
	for _, call := range a.EnabledCalls[serv.sandbox] {
		serv.targetEnabledSyscalls[serv.target.Syscalls[call]] = true
	}
	serv.mgr.machineChecked(a, serv.targetEnabledSyscalls)
	a.DisabledCalls = nil
	serv.checkResult = a
	serv.rotator = prog.MakeRotator(serv.target, serv.targetEnabledSyscalls, serv.rnd)
	return nil
}

func (serv *RPCServer) NewInput(a *rpctype.NewInputArgs, r *int) error {
	inputSignal := a.Signal.Deserialize()
	log.Logf(4, "new input from %v for syscall %v (signal=%v, cover=%v)",
		a.Name, a.Call, inputSignal.Len(), len(a.Cover))
	bad, disabled := checkProgram(serv.target, serv.targetEnabledSyscalls, a.RPCInput.Prog)
	if bad || disabled {
		log.Logf(0, "rejecting program from fuzzer (bad=%v, disabled=%v):\n%s", bad, disabled, a.RPCInput.Prog)
		return nil
	}
	serv.mu.Lock()
	defer serv.mu.Unlock()

	f := serv.fuzzers[a.Name]
	// Note: f may be nil if we called shutdownInstance,
	// but this request is already in-flight.
	genuine := !serv.corpusSignal.Diff(inputSignal).Empty()
	rotated := false
	if !genuine && f != nil && f.rotatedSignal != nil {
		rotated = !f.rotatedSignal.Diff(inputSignal).Empty()
	}
	if !genuine && !rotated {
		return nil
	}
	if !serv.mgr.newInput(a.RPCInput, inputSignal) {
		return nil
	}

	if f != nil && f.rotatedSignal != nil {
		f.rotatedSignal.Merge(inputSignal)
	}
	serv.corpusCover.Merge(a.Cover)
	serv.stats.corpusCover.set(len(serv.corpusCover))
	serv.stats.newInputs.inc()
	if rotated {
		serv.stats.rotatedInputs.inc()
	}

	if genuine {
		serv.corpusSignal.Merge(inputSignal)
		serv.stats.corpusSignal.set(serv.corpusSignal.Len())

		a.RPCInput.Cover = nil // Don't send coverage back to all fuzzers.
		for _, other := range serv.fuzzers {
			if other == f {
				continue
			}
			other.inputs = append(other.inputs, a.RPCInput)
		}
	}
	return nil
}

func (serv *RPCServer) Poll(a *rpctype.PollArgs, r *rpctype.PollRes) error {
	serv.stats.mergeNamed(a.Stats)

	serv.mu.Lock()
	defer serv.mu.Unlock()

	f := serv.fuzzers[a.Name]
	if f == nil {
		// This is possible if we called shutdownInstance,
		// but already have a pending request from this instance in-flight.
		log.Logf(1, "poll: fuzzer %v is not connected", a.Name)
		return nil
	}
	newMaxSignal := serv.maxSignal.Diff(a.MaxSignal.Deserialize())
	if !newMaxSignal.Empty() {
		serv.maxSignal.Merge(newMaxSignal)
		serv.stats.maxSignal.set(len(serv.maxSignal))
		for _, f1 := range serv.fuzzers {
			if f1 == f {
				continue
			}
			f1.newMaxSignal.Merge(newMaxSignal)
		}
	}
	r.MaxSignal = f.newMaxSignal.Split(500).Serialize()
	if a.NeedCandidates {
		r.Candidates = serv.mgr.candidateBatch(serv.batchSize)
	}
	if len(r.Candidates) == 0 {
		batchSize := serv.batchSize
		// When the fuzzer starts, it pumps the whole corpus.
		// If we do it using the final batchSize, it can be very slow
		// (batch of size 6 can take more than 10 mins for 50K corpus and slow kernel).
		// So use a larger batch initially (we use no stats as approximation of initial pump).
		const initialBatch = 30
		if len(a.Stats) == 0 && batchSize < initialBatch {
			batchSize = initialBatch
		}
		for i := 0; i < batchSize && len(f.inputs) > 0; i++ {
			last := len(f.inputs) - 1
			r.NewInputs = append(r.NewInputs, f.inputs[last])
			f.inputs[last] = rpctype.RPCInput{}
			f.inputs = f.inputs[:last]
		}
		if len(f.inputs) == 0 {
			f.inputs = nil
		}
	}
	log.Logf(4, "poll from %v: candidates=%v inputs=%v maxsignal=%v",
		a.Name, len(r.Candidates), len(r.NewInputs), len(r.MaxSignal.Elems))
	return nil
}

func (serv *RPCServer) shutdownInstance(name string) []byte {
	serv.mu.Lock()
	defer serv.mu.Unlock()

	fuzzer := serv.fuzzers[name]
	if fuzzer == nil {
		return nil
	}
	delete(serv.fuzzers, name)
	return fuzzer.machineInfo
}
<commit_msg>syz-manager: re-enable corpus rotation<commit_after>// Copyright 2018 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package main

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

	"github.com/google/syzkaller/pkg/cover"
	"github.com/google/syzkaller/pkg/log"
	"github.com/google/syzkaller/pkg/rpctype"
	"github.com/google/syzkaller/pkg/signal"
	"github.com/google/syzkaller/prog"
)

type RPCServer struct {
	mgr                   RPCManagerView
	port                  int
	target                *prog.Target
	configEnabledSyscalls []int
	targetEnabledSyscalls map[*prog.Syscall]bool
	stats                 *Stats
	sandbox               string
	batchSize             int

	mu           sync.Mutex
	fuzzers      map[string]*Fuzzer
	checkResult  *rpctype.CheckArgs
	maxSignal    signal.Signal
	corpusSignal signal.Signal
	corpusCover  cover.Cover
	rotator      *prog.Rotator
	rnd          *rand.Rand
}

type Fuzzer struct {
	name          string
	inputs        []rpctype.RPCInput
	newMaxSignal  signal.Signal
	rotatedSignal signal.Signal
	machineInfo   []byte
}

type BugFrames struct {
	memoryLeaks []string
	dataRaces   []string
}

// RPCManagerView restricts interface between RPCServer and Manager.
type RPCManagerView interface {
	fuzzerConnect() ([]rpctype.RPCInput, BugFrames)
	machineChecked(result *rpctype.CheckArgs, enabledSyscalls map[*prog.Syscall]bool)
	newInput(inp rpctype.RPCInput, sign signal.Signal) bool
	candidateBatch(size int) []rpctype.RPCCandidate
	rotateCorpus() bool
}

func startRPCServer(mgr *Manager) (*RPCServer, error) {
	serv := &RPCServer{
		mgr:                   mgr,
		target:                mgr.target,
		configEnabledSyscalls: mgr.configEnabledSyscalls,
		stats:                 mgr.stats,
		sandbox:               mgr.cfg.Sandbox,
		fuzzers:               make(map[string]*Fuzzer),
		rnd:                   rand.New(rand.NewSource(time.Now().UnixNano())),
	}
	serv.batchSize = 5
	if serv.batchSize < mgr.cfg.Procs {
		serv.batchSize = mgr.cfg.Procs
	}
	s, err := rpctype.NewRPCServer(mgr.cfg.RPC, "Manager", serv)
	if err != nil {
		return nil, err
	}
	log.Logf(0, "serving rpc on tcp://%v", s.Addr())
	serv.port = s.Addr().(*net.TCPAddr).Port
	go s.Serve()
	return serv, nil
}

func (serv *RPCServer) Connect(a *rpctype.ConnectArgs, r *rpctype.ConnectRes) error {
	log.Logf(1, "fuzzer %v connected", a.Name)
	serv.stats.vmRestarts.inc()

	corpus, bugFrames := serv.mgr.fuzzerConnect()

	serv.mu.Lock()
	defer serv.mu.Unlock()

	f := &Fuzzer{
		name:        a.Name,
		machineInfo: a.MachineInfo,
	}
	serv.fuzzers[a.Name] = f
	r.MemoryLeakFrames = bugFrames.memoryLeaks
	r.DataRaceFrames = bugFrames.dataRaces
	r.EnabledCalls = serv.configEnabledSyscalls
	r.GitRevision = prog.GitRevision
	r.TargetRevision = serv.target.Revision
	if serv.mgr.rotateCorpus() && serv.rnd.Intn(5) == 0 {
		// We do rotation every other time because there are no objective
		// proofs regarding its efficiency either way.
		// Also, rotation gives significantly skewed syscall selection
		// (run prog.TestRotationCoverage), it may or may not be OK.
		r.CheckResult = serv.rotateCorpus(f, corpus)
	} else {
		r.CheckResult = serv.checkResult
		f.inputs = corpus
		f.newMaxSignal = serv.maxSignal.Copy()
	}
	return nil
}

func (serv *RPCServer) rotateCorpus(f *Fuzzer, corpus []rpctype.RPCInput) *rpctype.CheckArgs {
	// Fuzzing tends to stuck in some local optimum and then it fails to cover
	// other state space points since code coverage is only a very approximate
	// measure of logic coverage. To overcome this we introduce some variation
	// into the process which should cause steady corpus rotation over time
	// (the same coverage is achieved in different ways).
	//
	// First, we select a subset of all syscalls for each VM run (result.EnabledCalls).
	// This serves 2 goals: (1) target fuzzer at a particular area of state space,
	// (2) disable syscalls that cause frequent crashes at least in some runs
	// to allow it to do actual fuzzing.
	//
	// Then, we remove programs that contain disabled syscalls from corpus
	// that will be sent to the VM (f.inputs). We also remove 10% of remaining
	// programs at random to allow to rediscover different variations of these programs.
	//
	// Then, we drop signal provided by the removed programs and also 10%
	// of the remaining signal at random (f.newMaxSignal). This again allows
	// rediscovery of this signal by different programs.
	//
	// Finally, we adjust criteria for accepting new programs from this VM (f.rotatedSignal).
	// This allows to accept rediscovered varied programs even if they don't
	// increase overall coverage. As the result we have multiple programs
	// providing the same duplicate coverage, these are removed during periodic
	// corpus minimization process. The minimization process is specifically
	// non-deterministic to allow the corpus rotation.
	//
	// Note: at no point we drop anything globally and permanently.
	// Everything we remove during this process is temporal and specific to a single VM.
	calls := serv.rotator.Select()

	var callIDs []int
	callNames := make(map[string]bool)
	for call := range calls {
		callNames[call.Name] = true
		callIDs = append(callIDs, call.ID)
	}

	f.inputs, f.newMaxSignal = serv.selectInputs(callNames, corpus, serv.maxSignal)
	// Remove the corresponding signal from rotatedSignal which will
	// be used to accept new inputs from this manager.
	f.rotatedSignal = serv.corpusSignal.Intersection(f.newMaxSignal)

	result := *serv.checkResult
	result.EnabledCalls = map[string][]int{serv.sandbox: callIDs}
	return &result
}

func (serv *RPCServer) selectInputs(enabled map[string]bool, inputs0 []rpctype.RPCInput, signal0 signal.Signal) (
	inputs []rpctype.RPCInput, signal signal.Signal) {
	signal = signal0.Copy()
	for _, inp := range inputs0 {
		calls, _, err := prog.CallSet(inp.Prog)
		if err != nil {
			panic(fmt.Sprintf("rotateInputs: CallSet failed: %v\n%s", err, inp.Prog))
		}
		for call := range calls {
			if !enabled[call] {
				goto drop
			}
		}
		if serv.rnd.Float64() > 0.9 {
			goto drop
		}
		inputs = append(inputs, inp)
		continue
	drop:
		for _, sig := range inp.Signal.Elems {
			delete(signal, sig)
		}
	}
	signal.Split(len(signal) / 10)
	return inputs, signal
}

func (serv *RPCServer) Check(a *rpctype.CheckArgs, r *int) error {
	serv.mu.Lock()
	defer serv.mu.Unlock()

	if serv.checkResult != nil {
		return nil
	}
	serv.targetEnabledSyscalls = make(map[*prog.Syscall]bool)
	for _, call := range a.EnabledCalls[serv.sandbox] {
		serv.targetEnabledSyscalls[serv.target.Syscalls[call]] = true
	}
	serv.mgr.machineChecked(a, serv.targetEnabledSyscalls)
	a.DisabledCalls = nil
	serv.checkResult = a
	serv.rotator = prog.MakeRotator(serv.target, serv.targetEnabledSyscalls, serv.rnd)
	return nil
}

func (serv *RPCServer) NewInput(a *rpctype.NewInputArgs, r *int) error {
	inputSignal := a.Signal.Deserialize()
	log.Logf(4, "new input from %v for syscall %v (signal=%v, cover=%v)",
		a.Name, a.Call, inputSignal.Len(), len(a.Cover))
	bad, disabled := checkProgram(serv.target, serv.targetEnabledSyscalls, a.RPCInput.Prog)
	if bad || disabled {
		log.Logf(0, "rejecting program from fuzzer (bad=%v, disabled=%v):\n%s", bad, disabled, a.RPCInput.Prog)
		return nil
	}
	serv.mu.Lock()
	defer serv.mu.Unlock()

	f := serv.fuzzers[a.Name]
	// Note: f may be nil if we called shutdownInstance,
	// but this request is already in-flight.
	genuine := !serv.corpusSignal.Diff(inputSignal).Empty()
	rotated := false
	if !genuine && f != nil && f.rotatedSignal != nil {
		rotated = !f.rotatedSignal.Diff(inputSignal).Empty()
	}
	if !genuine && !rotated {
		return nil
	}
	if !serv.mgr.newInput(a.RPCInput, inputSignal) {
		return nil
	}

	if f != nil && f.rotatedSignal != nil {
		f.rotatedSignal.Merge(inputSignal)
	}
	serv.corpusCover.Merge(a.Cover)
	serv.stats.corpusCover.set(len(serv.corpusCover))
	serv.stats.newInputs.inc()
	if rotated {
		serv.stats.rotatedInputs.inc()
	}

	if genuine {
		serv.corpusSignal.Merge(inputSignal)
		serv.stats.corpusSignal.set(serv.corpusSignal.Len())

		a.RPCInput.Cover = nil // Don't send coverage back to all fuzzers.
		for _, other := range serv.fuzzers {
			if other == f || other.rotatedSignal != nil {
				continue
			}
			other.inputs = append(other.inputs, a.RPCInput)
		}
	}
	return nil
}

func (serv *RPCServer) Poll(a *rpctype.PollArgs, r *rpctype.PollRes) error {
	serv.stats.mergeNamed(a.Stats)

	serv.mu.Lock()
	defer serv.mu.Unlock()

	f := serv.fuzzers[a.Name]
	if f == nil {
		// This is possible if we called shutdownInstance,
		// but already have a pending request from this instance in-flight.
		log.Logf(1, "poll: fuzzer %v is not connected", a.Name)
		return nil
	}
	newMaxSignal := serv.maxSignal.Diff(a.MaxSignal.Deserialize())
	if !newMaxSignal.Empty() {
		serv.maxSignal.Merge(newMaxSignal)
		serv.stats.maxSignal.set(len(serv.maxSignal))
		for _, f1 := range serv.fuzzers {
			if f1 == f || f1.rotatedSignal != nil {
				continue
			}
			f1.newMaxSignal.Merge(newMaxSignal)
		}
	}
	if f.rotatedSignal != nil {
		// Let rotated VMs run in isolation, don't send them anything.
		return nil
	}
	r.MaxSignal = f.newMaxSignal.Split(500).Serialize()
	if a.NeedCandidates {
		r.Candidates = serv.mgr.candidateBatch(serv.batchSize)
	}
	if len(r.Candidates) == 0 {
		batchSize := serv.batchSize
		// When the fuzzer starts, it pumps the whole corpus.
		// If we do it using the final batchSize, it can be very slow
		// (batch of size 6 can take more than 10 mins for 50K corpus and slow kernel).
		// So use a larger batch initially (we use no stats as approximation of initial pump).
		const initialBatch = 30
		if len(a.Stats) == 0 && batchSize < initialBatch {
			batchSize = initialBatch
		}
		for i := 0; i < batchSize && len(f.inputs) > 0; i++ {
			last := len(f.inputs) - 1
			r.NewInputs = append(r.NewInputs, f.inputs[last])
			f.inputs[last] = rpctype.RPCInput{}
			f.inputs = f.inputs[:last]
		}
		if len(f.inputs) == 0 {
			f.inputs = nil
		}
	}
	log.Logf(4, "poll from %v: candidates=%v inputs=%v maxsignal=%v",
		a.Name, len(r.Candidates), len(r.NewInputs), len(r.MaxSignal.Elems))
	return nil
}

func (serv *RPCServer) shutdownInstance(name string) []byte {
	serv.mu.Lock()
	defer serv.mu.Unlock()

	fuzzer := serv.fuzzers[name]
	if fuzzer == nil {
		return nil
	}
	delete(serv.fuzzers, name)
	return fuzzer.machineInfo
}
<|endoftext|>


<commit_before>// //+build ignore

// NOTE (Jeff Mitchell): I'm disabling this for now as this package never
// changes, but leaving this file active causes go mod to pull in some very
// ugly deps. This can be run manually if further changes are ever made to this
// lib.

// Copyright (c) 2012, 2013 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a BSD-style license found in the LICENSE file.

package codec

/*

// This file includes benchmarks which have dependencies on 3rdparty
// packages (bson and vmihailenco/msgpack) which must be installed locally.
//
// To run the benchmarks including these 3rdparty packages, first
//   - Uncomment first line in this file (put // // in front of it)
//   - Get those packages:
//       go get github.com/vmihailenco/msgpack
//       go get labix.org/v2/mgo/bson
//   - Run:
//       go test -bi -bench=.

import (
	"testing"

	vmsgpack "gopkg.in/vmihailenco/msgpack.v2"
	"labix.org/v2/mgo/bson"
)

func init() {
	benchCheckers = append(benchCheckers,
		benchChecker{"v-msgpack", fnVMsgpackEncodeFn, fnVMsgpackDecodeFn},
		benchChecker{"bson", fnBsonEncodeFn, fnBsonDecodeFn},
	)
}

func fnVMsgpackEncodeFn(ts interface{}) ([]byte, error) {
	return vmsgpack.Marshal(ts)
}

func fnVMsgpackDecodeFn(buf []byte, ts interface{}) error {
	return vmsgpack.Unmarshal(buf, ts)
}

func fnBsonEncodeFn(ts interface{}) ([]byte, error) {
	return bson.Marshal(ts)
}

func fnBsonDecodeFn(buf []byte, ts interface{}) error {
	return bson.Unmarshal(buf, ts)
}

func Benchmark__Bson_______Encode(b *testing.B) {
	fnBenchmarkEncode(b, "bson", benchTs, fnBsonEncodeFn)
}

func Benchmark__Bson_______Decode(b *testing.B) {
	fnBenchmarkDecode(b, "bson", benchTs, fnBsonEncodeFn, fnBsonDecodeFn, fnBenchNewTs)
}

func Benchmark__VMsgpack___Encode(b *testing.B) {
	fnBenchmarkEncode(b, "v-msgpack", benchTs, fnVMsgpackEncodeFn)
}

func Benchmark__VMsgpack___Decode(b *testing.B) {
	fnBenchmarkDecode(b, "v-msgpack", benchTs, fnVMsgpackEncodeFn, fnVMsgpackDecodeFn, fnBenchNewTs)
}

func TestMsgpackPythonGenStreams(t *testing.T) {
	doTestMsgpackPythonGenStreams(t)
}

func TestMsgpackRpcSpecGoClientToPythonSvc(t *testing.T) {
	doTestMsgpackRpcSpecGoClientToPythonSvc(t)
}

func TestMsgpackRpcSpecPythonClientToGoSvc(t *testing.T) {
	doTestMsgpackRpcSpecPythonClientToGoSvc(t)
}
*/
<commit_msg>fix +build directive<commit_after>//+build ignore

// NOTE (Jeff Mitchell): I'm disabling this for now as this package never
// changes, but leaving this file active causes go mod to pull in some very
// ugly deps. This can be run manually if further changes are ever made to this
// lib.

// Copyright (c) 2012, 2013 Ugorji Nwoke. All rights reserved.
// Use of this source code is governed by a BSD-style license found in the LICENSE file.

package codec

/*

// This file includes benchmarks which have dependencies on 3rdparty
// packages (bson and vmihailenco/msgpack) which must be installed locally.
//
// To run the benchmarks including these 3rdparty packages, first
//   - Uncomment first line in this file (put // // in front of it)
//   - Get those packages:
//       go get github.com/vmihailenco/msgpack
//       go get labix.org/v2/mgo/bson
//   - Run:
//       go test -bi -bench=.

import (
	"testing"

	vmsgpack "gopkg.in/vmihailenco/msgpack.v2"
	"labix.org/v2/mgo/bson"
)

func init() {
	benchCheckers = append(benchCheckers,
		benchChecker{"v-msgpack", fnVMsgpackEncodeFn, fnVMsgpackDecodeFn},
		benchChecker{"bson", fnBsonEncodeFn, fnBsonDecodeFn},
	)
}

func fnVMsgpackEncodeFn(ts interface{}) ([]byte, error) {
	return vmsgpack.Marshal(ts)
}

func fnVMsgpackDecodeFn(buf []byte, ts interface{}) error {
	return vmsgpack.Unmarshal(buf, ts)
}

func fnBsonEncodeFn(ts interface{}) ([]byte, error) {
	return bson.Marshal(ts)
}

func fnBsonDecodeFn(buf []byte, ts interface{}) error {
	return bson.Unmarshal(buf, ts)
}

func Benchmark__Bson_______Encode(b *testing.B) {
	fnBenchmarkEncode(b, "bson", benchTs, fnBsonEncodeFn)
}

func Benchmark__Bson_______Decode(b *testing.B) {
	fnBenchmarkDecode(b, "bson", benchTs, fnBsonEncodeFn, fnBsonDecodeFn, fnBenchNewTs)
}

func Benchmark__VMsgpack___Encode(b *testing.B) {
	fnBenchmarkEncode(b, "v-msgpack", benchTs, fnVMsgpackEncodeFn)
}

func Benchmark__VMsgpack___Decode(b *testing.B) {
	fnBenchmarkDecode(b, "v-msgpack", benchTs, fnVMsgpackEncodeFn, fnVMsgpackDecodeFn, fnBenchNewTs)
}

func TestMsgpackPythonGenStreams(t *testing.T) {
	doTestMsgpackPythonGenStreams(t)
}

func TestMsgpackRpcSpecGoClientToPythonSvc(t *testing.T) {
	doTestMsgpackRpcSpecGoClientToPythonSvc(t)
}

func TestMsgpackRpcSpecPythonClientToGoSvc(t *testing.T) {
	doTestMsgpackRpcSpecPythonClientToGoSvc(t)
}
*/
<|endoftext|>


<commit_before>package cpu

import (
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	"github.com/shirou/gopsutil/internal/common"
)

// sys/resource.h
const (
	CPUser    = 0
	CPNice    = 1
	CPSys     = 2
	CPIntr    = 3
	CPIdle    = 4
	CPUStates = 5
)

var ClocksPerSec = float64(128)

func init() {
	getconf, err := exec.LookPath("/usr/bin/getconf")
	if err != nil {
		return
	}
	out, err := exec.Command(getconf, "CLK_TCK").Output()
	// ignore errors
	if err == nil {
		i, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
		if err == nil {
			ClocksPerSec = float64(i)
		}
	}
}

func Times(percpu bool) ([]TimesStat, error) {
	var ret []TimesStat

	var sysctlCall string
	var ncpu int
	if percpu {
		sysctlCall = "kern.cp_times"
		ncpu, _ = Counts(true)
	} else {
		sysctlCall = "kern.cp_time"
		ncpu = 1
	}

	cpuTimes, err := common.DoSysctrl(sysctlCall)
	if err != nil {
		return ret, err
	}

	for i := 0; i < ncpu; i++ {
		offset := CPUStates * i
		user, err := strconv.ParseFloat(cpuTimes[CPUser+offset], 64)
		if err != nil {
			return ret, err
		}
		nice, err := strconv.ParseFloat(cpuTimes[CPNice+offset], 64)
		if err != nil {
			return ret, err
		}
		sys, err := strconv.ParseFloat(cpuTimes[CPSys+offset], 64)
		if err != nil {
			return ret, err
		}
		idle, err := strconv.ParseFloat(cpuTimes[CPIdle+offset], 64)
		if err != nil {
			return ret, err
		}
		intr, err := strconv.ParseFloat(cpuTimes[CPIntr+offset], 64)
		if err != nil {
			return ret, err
		}

		c := TimesStat{
			User:   float64(user / ClocksPerSec),
			Nice:   float64(nice / ClocksPerSec),
			System: float64(sys / ClocksPerSec),
			Idle:   float64(idle / ClocksPerSec),
			Irq:    float64(intr / ClocksPerSec),
		}
		if !percpu {
			c.CPU = "cpu-total"
		} else {
			c.CPU = fmt.Sprintf("cpu%d", i)
		}

		ret = append(ret, c)
	}

	return ret, nil
}

// Returns only one CPUInfoStat on FreeBSD
func Info() ([]InfoStat, error) {
	filename := "/var/run/dmesg.boot"
	lines, _ := common.ReadLines(filename)

	var ret []InfoStat

	c := InfoStat{}
	for _, line := range lines {
		if matches := regexp.MustCompile(`CPU:\s+(.+) \(([\d.]+).+\)`).FindStringSubmatch(line); matches != nil {
			c.ModelName = matches[1]
			t, err := strconv.ParseFloat(matches[2], 64)
			if err != nil {
				return ret, nil
			}
			c.Mhz = t
		} else if matches := regexp.MustCompile(`Origin = "(.+)"  Id = (.+)  Family = (.+)  Model = (.+)  Stepping = (.+)`).FindStringSubmatch(line); matches != nil {
			c.VendorID = matches[1]
			c.Family = matches[3]
			c.Model = matches[4]
			t, err := strconv.ParseInt(matches[5], 10, 32)
			if err != nil {
				return ret, nil
			}
			c.Stepping = int32(t)
		} else if matches := regexp.MustCompile(`Features=.+<(.+)>`).FindStringSubmatch(line); matches != nil {
			for _, v := range strings.Split(matches[1], ",") {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if matches := regexp.MustCompile(`Features2=[a-f\dx]+<(.+)>`).FindStringSubmatch(line); matches != nil {
			for _, v := range strings.Split(matches[1], ",") {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if matches := regexp.MustCompile(`Logical CPUs per core: (\d+)`).FindStringSubmatch(line); matches != nil {
			// FIXME: no this line?
			t, err := strconv.ParseInt(matches[1], 10, 32)
			if err != nil {
				return ret, nil
			}
			c.Cores = int32(t)
		}

	}

	return append(ret, c), nil
}
<commit_msg>Improve the accuracy of CPU detection on FreeBSD.<commit_after>package cpu

import (
	"fmt"
	"os/exec"
	"regexp"
	"strconv"
	"strings"

	"github.com/shirou/gopsutil/internal/common"
)

// sys/resource.h
const (
	CPUser    = 0
	CPNice    = 1
	CPSys     = 2
	CPIntr    = 3
	CPIdle    = 4
	CPUStates = 5
)

var ClocksPerSec = float64(128)

func init() {
	getconf, err := exec.LookPath("/usr/bin/getconf")
	if err != nil {
		return
	}
	out, err := exec.Command(getconf, "CLK_TCK").Output()
	// ignore errors
	if err == nil {
		i, err := strconv.ParseFloat(strings.TrimSpace(string(out)), 64)
		if err == nil {
			ClocksPerSec = float64(i)
		}
	}
}

func Times(percpu bool) ([]TimesStat, error) {
	var ret []TimesStat

	var sysctlCall string
	var ncpu int
	if percpu {
		sysctlCall = "kern.cp_times"
		ncpu, _ = Counts(true)
	} else {
		sysctlCall = "kern.cp_time"
		ncpu = 1
	}

	cpuTimes, err := common.DoSysctrl(sysctlCall)
	if err != nil {
		return ret, err
	}

	for i := 0; i < ncpu; i++ {
		offset := CPUStates * i
		user, err := strconv.ParseFloat(cpuTimes[CPUser+offset], 64)
		if err != nil {
			return ret, err
		}
		nice, err := strconv.ParseFloat(cpuTimes[CPNice+offset], 64)
		if err != nil {
			return ret, err
		}
		sys, err := strconv.ParseFloat(cpuTimes[CPSys+offset], 64)
		if err != nil {
			return ret, err
		}
		idle, err := strconv.ParseFloat(cpuTimes[CPIdle+offset], 64)
		if err != nil {
			return ret, err
		}
		intr, err := strconv.ParseFloat(cpuTimes[CPIntr+offset], 64)
		if err != nil {
			return ret, err
		}

		c := TimesStat{
			User:   float64(user / ClocksPerSec),
			Nice:   float64(nice / ClocksPerSec),
			System: float64(sys / ClocksPerSec),
			Idle:   float64(idle / ClocksPerSec),
			Irq:    float64(intr / ClocksPerSec),
		}
		if !percpu {
			c.CPU = "cpu-total"
		} else {
			c.CPU = fmt.Sprintf("cpu%d", i)
		}

		ret = append(ret, c)
	}

	return ret, nil
}

// Returns only one InfoStat on FreeBSD.  The information regarding core
// count, however is accurate and it is assumed that all InfoStat attributes
// are the same across CPUs.
func Info() ([]InfoStat, error) {
	const dmesgBoot = "/var/run/dmesg.boot"
	lines, _ := common.ReadLines(dmesgBoot)

	c := InfoStat{}
	var vals []string
	var err error
	if vals, err = common.DoSysctrl("hw.clockrate"); err != nil {
		return nil, err
	}
	if c.Mhz, err = strconv.ParseFloat(vals[0], 64); err != nil {
		return nil, fmt.Errorf("Unable to parse FreeBSD CPU clock rate: %v", err)
	}
	c.CPU = int32(c.Mhz)

	if vals, err = common.DoSysctrl("hw.ncpu"); err != nil {
		return nil, err
	}
	var i64 int64
	if i64, err = strconv.ParseInt(vals[0], 10, 32); err != nil {
		return nil, fmt.Errorf("Unable to parse FreeBSD cores: %v", err)
	}
	c.Cores = int32(i64)

	if vals, err = common.DoSysctrl("hw.model"); err != nil {
		return nil, err
	}
	c.ModelName = vals[0]

	for _, line := range lines {
		if matches := regexp.MustCompile(`Origin\s*=\s*"(.+)"\s+Id\s*=\s*(.+)\s+Family\s*=\s*(.+)\s+Model\s*=\s*(.+)\s+Stepping\s*=\s*(.+)`).FindStringSubmatch(line); matches != nil {
			c.VendorID = matches[1]
			c.Family = matches[3]
			c.Model = matches[4]
			t, err := strconv.ParseInt(matches[5], 10, 32)
			if err != nil {
				return nil, fmt.Errorf("Unable to parse FreeBSD CPU stepping information from %q: %v", line, err)
			}
			c.Stepping = int32(t)
		} else if matches := regexp.MustCompile(`Features=.+<(.+)>`).FindStringSubmatch(line); matches != nil {
			for _, v := range strings.Split(matches[1], ",") {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		} else if matches := regexp.MustCompile(`Features2=[a-f\dx]+<(.+)>`).FindStringSubmatch(line); matches != nil {
			for _, v := range strings.Split(matches[1], ",") {
				c.Flags = append(c.Flags, strings.ToLower(v))
			}
		}
	}

	return []InfoStat{c}, nil
}
<|endoftext|>


<commit_before>package aws

import (
	mach "github.com/jeffjen/machine/lib/machine"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/codegangsta/cli"

	"fmt"
	"math/rand"
	"net"
	"os"
	"time"
)

func init() {
	rand.Seed(time.Now().Unix())
}

var (
	// AWS EC2 client object for establishing command
	svc *ec2.EC2

	// Instance Roster
	instList = make(mach.RegisteredInstances)
)

func NewCommand() cli.Command {
	return cli.Command{
		Name:  "aws",
		Usage: "Create and Manage AWS machine",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "region", EnvVar: "AWS_REGION", Usage: "AWS Region"},
			cli.StringFlag{Name: "key", EnvVar: "AWS_ACCESS_KEY_ID", Usage: "AWS access key"},
			cli.StringFlag{Name: "secret", EnvVar: "AWS_SECRET_ACCESS_KEY", Usage: "AWS secret key"},
			cli.StringFlag{Name: "token", EnvVar: "AWS_SESSION_TOKEN", Usage: "session token for temporary credentials"},
			cli.StringFlag{Name: "user", EnvVar: "MACHINE_USER", Usage: "Run command as user"},
			cli.StringFlag{Name: "cert", EnvVar: "MACHINE_CERT_FILE", Usage: "Private key to use in Authentication"},
		},
		Before: func(c *cli.Context) error {
			// bootstrap EC2 client with command line args
			cfg := aws.NewConfig()
			if region := c.String("region"); region != "" {
				cfg = cfg.WithRegion(region)
			}
			if id, secret, token := c.String("key"), c.String("secret"), c.String("token"); id != "" && secret != "" {
				cfg = cfg.WithCredentials(credentials.NewStaticCredentials(id, secret, token))
			}
			svc = ec2.New(session.New(cfg))
			return nil
		},
		Subcommands: []cli.Command{
			newConfigCommand(),
			newCreateCommand(),
			newStartCommand(),
			newRmCommand(),
			newImageCommand(),
		},
	}
}

func newStartCommand() cli.Command {
	return cli.Command{
		Name:  "start",
		Usage: "Start instance",
		Action: func(c *cli.Context) error {
			var name = c.Args().First()

			// Load from Instance Roster
			instList.Load()

			info, ok := instList[name]
			if !ok {
				fmt.Fprintln(os.Stderr, "Target machine not found")
				os.Exit(1)
			}

			_, err := svc.StartInstances(&ec2.StartInstancesInput{
				InstanceIds: []*string{
					aws.String(info.Id),
				},
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			return nil
		},
	}
}

func newRmCommand() cli.Command {
	return cli.Command{
		Name:  "rm",
		Usage: "Remove and Terminate instance",
		Action: func(c *cli.Context) error {
			var name = c.Args().First()

			// Load from Instance Roster
			instList.Load()

			info, ok := instList[name]
			if !ok {
				fmt.Fprintln(os.Stderr, "Target machine not found")
				os.Exit(1)
			}

			_, err := svc.TerminateInstances(&ec2.TerminateInstancesInput{
				InstanceIds: []*string{
					aws.String(info.Id),
				},
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			return nil
		},
	}
}

func newConfigCommand() cli.Command {
	return cli.Command{
		Name:  "config",
		Usage: "Configure AWS environment",
		Subcommands: []cli.Command{
			syncFromAWS(),
			getFromAWSConfig(),
		},
	}
}

func newCreateCommand() cli.Command {
	return cli.Command{
		Name:  "create",
		Usage: "Create a new EC2 instance",
		Flags: []cli.Flag{
			cli.BoolTFlag{Name: "use-docker", Usage: "Opt in to use Docker Engine"},
			cli.StringFlag{Name: "ami-id", Usage: "EC2 instance AMI ID"},
			cli.IntFlag{Name: "count", Value: 1, Usage: "EC2 instances to launch in this request"},
			cli.StringSliceFlag{Name: "group", Usage: "Network security group for user"},
			cli.StringFlag{Name: "iam-role", Usage: "EC2 IAM Role to apply"},
			cli.StringFlag{Name: "profile", Value: "default", Usage: "Name of the profile"},
			cli.IntFlag{Name: "root-size", Value: 8, Usage: "EC2 root volume size"},
			cli.StringFlag{Name: "ssh-key", Usage: "EC2 instance SSH KeyPair"},
			cli.BoolFlag{Name: "subnet-private", Usage: "Launch EC2 instance to internal subnet"},
			cli.StringFlag{Name: "subnet-id", Usage: "Launch EC2 instance to the specified subnet"},
			cli.StringSliceFlag{Name: "tag", Usage: "EC2 instance tag in the form field=value"},
			cli.StringFlag{Name: "type", Value: "t2.micro", Usage: "EC2 instance type"},
			cli.IntSliceFlag{Name: "volume-size", Usage: "EC2 EBS volume size"},
		},
		Action: func(c *cli.Context) error {
			var (
				profile = make(AWSProfile)

				user = c.GlobalString("user")
				cert = c.GlobalString("cert")

				name = c.Args().First()

				useDocker = c.Bool("use-docker")
			)

			if user == "" || cert == "" {
				fmt.Fprintln(os.Stderr, "Missing required remote auth info")
				os.Exit(1)
			}

			// Load from AWS configuration from last sync
			profile.Load()

			// Load from Instance Roster to register and defer write back
			defer instList.Load().Dump()

			if name == "" {
				fmt.Fprintln(os.Stderr, "Required argument `name` missing")
				os.Exit(1)
			} else if _, ok := instList[name]; ok {
				fmt.Fprintln(os.Stderr, "Machine exist")
				os.Exit(1)
			}

			region, ok := profile[c.GlobalString("region")]
			if !ok {
				fmt.Fprintln(os.Stderr, "Please run sync in the region of choice")
				os.Exit(1)
			}
			p, ok := region[c.String("profile")]
			if !ok {
				fmt.Fprintln(os.Stderr, "Unable to find matching VPC profile")
				os.Exit(1)
			}

			// Invoke EC2 launch procedure
			for state := range newEC2Inst(c, p, user, cert, name, useDocker) {
				if addr, _ := net.ResolveTCPAddr("tcp", *state.PublicIpAddress+":2376"); state.err == nil {
					fmt.Printf("%s - %s - Instance ID: %s\n", *state.PublicIpAddress, *state.PrivateIpAddress, *state.InstanceId)
					if useDocker {
						instList[state.name] = &mach.Instance{
							Id:         *state.InstanceId,
							Driver:     "aws",
							DockerHost: addr,
							State:      "running",
						}
					}
				} else {
					fmt.Fprintln(os.Stderr, state.err)
				}
			}

			return nil
		},
	}
}

func newImageCommand() cli.Command {
	return cli.Command{
		Name:  "register-ami",
		Usage: "Register an AMI from specification",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "instance-id", Usage: "EC2 instance ID"},
			cli.StringFlag{Name: "name", Usage: "EC2 AMI Name"},
			cli.StringFlag{Name: "desc", Usage: "EC2 AMI Description"},
		},
		Action: func(c *cli.Context) error {
			var (
				instId = c.String("instance-id")
				name   = c.String("name")
				desc   = c.String("desc")
			)

			resp, err := svc.CreateImage(&ec2.CreateImageInput{
				InstanceId:  aws.String(instId),
				Name:        aws.String(name),
				Description: aws.String(desc),
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			} else {
				fmt.Println(*resp.ImageId)
			}

			return nil
		},
	}
}
<commit_msg>UPDATE: aws stop instance command<commit_after>package aws

import (
	mach "github.com/jeffjen/machine/lib/machine"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/ec2"
	"github.com/codegangsta/cli"

	"fmt"
	"math/rand"
	"net"
	"os"
	"time"
)

func init() {
	rand.Seed(time.Now().Unix())
}

var (
	// AWS EC2 client object for establishing command
	svc *ec2.EC2

	// Instance Roster
	instList = make(mach.RegisteredInstances)
)

func NewCommand() cli.Command {
	return cli.Command{
		Name:  "aws",
		Usage: "Create and Manage AWS machine",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "region", EnvVar: "AWS_REGION", Usage: "AWS Region"},
			cli.StringFlag{Name: "key", EnvVar: "AWS_ACCESS_KEY_ID", Usage: "AWS access key"},
			cli.StringFlag{Name: "secret", EnvVar: "AWS_SECRET_ACCESS_KEY", Usage: "AWS secret key"},
			cli.StringFlag{Name: "token", EnvVar: "AWS_SESSION_TOKEN", Usage: "session token for temporary credentials"},
			cli.StringFlag{Name: "user", EnvVar: "MACHINE_USER", Usage: "Run command as user"},
			cli.StringFlag{Name: "cert", EnvVar: "MACHINE_CERT_FILE", Usage: "Private key to use in Authentication"},
		},
		Before: func(c *cli.Context) error {
			// bootstrap EC2 client with command line args
			cfg := aws.NewConfig()
			if region := c.String("region"); region != "" {
				cfg = cfg.WithRegion(region)
			}
			if id, secret, token := c.String("key"), c.String("secret"), c.String("token"); id != "" && secret != "" {
				cfg = cfg.WithCredentials(credentials.NewStaticCredentials(id, secret, token))
			}
			svc = ec2.New(session.New(cfg))
			return nil
		},
		Subcommands: []cli.Command{
			newConfigCommand(),
			newCreateCommand(),
			newStartCommand(),
			newStopCommand(),
			newRmCommand(),
			newImageCommand(),
		},
	}
}

func newStartCommand() cli.Command {
	return cli.Command{
		Name:  "start",
		Usage: "Start instance",
		Action: func(c *cli.Context) error {
			var name = c.Args().First()

			// Load from Instance Roster
			instList.Load()

			info, ok := instList[name]
			if !ok {
				fmt.Fprintln(os.Stderr, "Target machine not found")
				os.Exit(1)
			}

			_, err := svc.StartInstances(&ec2.StartInstancesInput{
				InstanceIds: []*string{
					aws.String(info.Id),
				},
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			return nil
		},
	}
}

func newStopCommand() cli.Command {
	return cli.Command{
		Name:  "stop",
		Usage: "Stop instance",
		Action: func(c *cli.Context) error {
			var name = c.Args().First()

			// Load from Instance Roster
			instList.Load()

			info, ok := instList[name]
			if !ok {
				fmt.Fprintln(os.Stderr, "Target machine not found")
				os.Exit(1)
			}

			_, err := svc.StopInstances(&ec2.StopInstancesInput{
				InstanceIds: []*string{
					aws.String(info.Id),
				},
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}

			return nil
		},
	}
}

func newRmCommand() cli.Command {
	return cli.Command{
		Name:  "rm",
		Usage: "Remove and Terminate instance",
		Action: func(c *cli.Context) error {
			var name = c.Args().First()

			// Load from Instance Roster
			instList.Load()

			info, ok := instList[name]
			if !ok {
				fmt.Fprintln(os.Stderr, "Target machine not found")
				os.Exit(1)
			}

			_, err := svc.TerminateInstances(&ec2.TerminateInstancesInput{
				InstanceIds: []*string{
					aws.String(info.Id),
				},
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			}
			return nil
		},
	}
}

func newConfigCommand() cli.Command {
	return cli.Command{
		Name:  "config",
		Usage: "Configure AWS environment",
		Subcommands: []cli.Command{
			syncFromAWS(),
			getFromAWSConfig(),
		},
	}
}

func newCreateCommand() cli.Command {
	return cli.Command{
		Name:  "create",
		Usage: "Create a new EC2 instance",
		Flags: []cli.Flag{
			cli.BoolTFlag{Name: "use-docker", Usage: "Opt in to use Docker Engine"},
			cli.StringFlag{Name: "ami-id", Usage: "EC2 instance AMI ID"},
			cli.IntFlag{Name: "count", Value: 1, Usage: "EC2 instances to launch in this request"},
			cli.StringSliceFlag{Name: "group", Usage: "Network security group for user"},
			cli.StringFlag{Name: "iam-role", Usage: "EC2 IAM Role to apply"},
			cli.StringFlag{Name: "profile", Value: "default", Usage: "Name of the profile"},
			cli.IntFlag{Name: "root-size", Value: 8, Usage: "EC2 root volume size"},
			cli.StringFlag{Name: "ssh-key", Usage: "EC2 instance SSH KeyPair"},
			cli.BoolFlag{Name: "subnet-private", Usage: "Launch EC2 instance to internal subnet"},
			cli.StringFlag{Name: "subnet-id", Usage: "Launch EC2 instance to the specified subnet"},
			cli.StringSliceFlag{Name: "tag", Usage: "EC2 instance tag in the form field=value"},
			cli.StringFlag{Name: "type", Value: "t2.micro", Usage: "EC2 instance type"},
			cli.IntSliceFlag{Name: "volume-size", Usage: "EC2 EBS volume size"},
		},
		Action: func(c *cli.Context) error {
			var (
				profile = make(AWSProfile)

				user = c.GlobalString("user")
				cert = c.GlobalString("cert")

				name = c.Args().First()

				useDocker = c.Bool("use-docker")
			)

			if user == "" || cert == "" {
				fmt.Fprintln(os.Stderr, "Missing required remote auth info")
				os.Exit(1)
			}

			// Load from AWS configuration from last sync
			profile.Load()

			// Load from Instance Roster to register and defer write back
			defer instList.Load().Dump()

			if name == "" {
				fmt.Fprintln(os.Stderr, "Required argument `name` missing")
				os.Exit(1)
			} else if _, ok := instList[name]; ok {
				fmt.Fprintln(os.Stderr, "Machine exist")
				os.Exit(1)
			}

			region, ok := profile[c.GlobalString("region")]
			if !ok {
				fmt.Fprintln(os.Stderr, "Please run sync in the region of choice")
				os.Exit(1)
			}
			p, ok := region[c.String("profile")]
			if !ok {
				fmt.Fprintln(os.Stderr, "Unable to find matching VPC profile")
				os.Exit(1)
			}

			// Invoke EC2 launch procedure
			for state := range newEC2Inst(c, p, user, cert, name, useDocker) {
				if addr, _ := net.ResolveTCPAddr("tcp", *state.PublicIpAddress+":2376"); state.err == nil {
					fmt.Printf("%s - %s - Instance ID: %s\n", *state.PublicIpAddress, *state.PrivateIpAddress, *state.InstanceId)
					if useDocker {
						instList[state.name] = &mach.Instance{
							Id:         *state.InstanceId,
							Driver:     "aws",
							DockerHost: addr,
							State:      "running",
						}
					}
				} else {
					fmt.Fprintln(os.Stderr, state.err)
				}
			}

			return nil
		},
	}
}

func newImageCommand() cli.Command {
	return cli.Command{
		Name:  "register-ami",
		Usage: "Register an AMI from specification",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "instance-id", Usage: "EC2 instance ID"},
			cli.StringFlag{Name: "name", Usage: "EC2 AMI Name"},
			cli.StringFlag{Name: "desc", Usage: "EC2 AMI Description"},
		},
		Action: func(c *cli.Context) error {
			var (
				instId = c.String("instance-id")
				name   = c.String("name")
				desc   = c.String("desc")
			)

			resp, err := svc.CreateImage(&ec2.CreateImageInput{
				InstanceId:  aws.String(instId),
				Name:        aws.String(name),
				Description: aws.String(desc),
			})
			if err != nil {
				fmt.Fprintln(os.Stderr, err)
				os.Exit(1)
			} else {
				fmt.Println(*resp.ImageId)
			}

			return nil
		},
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2020 gRPC 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 controllers

import (
	"context"

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

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	grpcv1 "github.com/grpc/test-infra/api/v1"
	"github.com/grpc/test-infra/config"
	"github.com/grpc/test-infra/podbuilder"
	"github.com/grpc/test-infra/status"
)

// createPod creates a pod resource, given a pod pointer and a test pointer.
func createPod(pod *corev1.Pod, test *grpcv1.LoadTest) error {
	// TODO: Get the controllerRef to work here.
	// kind := reflect.TypeOf(grpcv1.LoadTest{}).Name()
	// gvk := grpcv1.GroupVersion.WithKind(kind)
	// controllerRef := metav1.NewControllerRef(test, gvk)
	// pod.SetOwnerReferences([]metav1.OwnerReference{*controllerRef})
	return k8sClient.Create(context.Background(), pod)
}

// updatePodWithContainerState changes the container state in the status of a
// pod resource that already exists on the cluster. This is useful for testing
// different failing, running and succeeding states.
func updatePodWithContainerState(pod *corev1.Pod, containerState corev1.ContainerState) error {
	status := &pod.Status
	status.ContainerStatuses = []corev1.ContainerStatus{
		{
			State: containerState,
		},
	}
	return k8sClient.Status().Update(context.Background(), pod)
}

var _ = Describe("LoadTest controller", func() {
	var test *grpcv1.LoadTest
	var namespacedName types.NamespacedName

	BeforeEach(func() {
		test = newLoadTest()
		namespacedName = types.NamespacedName{
			Name:      test.Name,
			Namespace: test.Namespace,
		}
	})

	It("does not change the status after termination", func() {
		now := metav1.Now()
		test.Status = grpcv1.LoadTestStatus{
			State:     grpcv1.Succeeded,
			StartTime: &now,
			StopTime:  &now,
		}
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())
		Expect(k8sClient.Status().Update(context.Background(), test)).To(Succeed())

		getTestStatus := func() (grpcv1.LoadTestStatus, error) {
			fetchedTest := new(grpcv1.LoadTest)
			err := k8sClient.Get(context.Background(), namespacedName, fetchedTest)
			if err != nil {
				return grpcv1.LoadTestStatus{}, err
			}
			return fetchedTest.Status, nil
		}

		By("ensuring we can eventually get the created status")
		Eventually(getTestStatus).Should(Equal(test.Status))

		By("checking that the expected status remains unchanged")
		Consistently(getTestStatus).Should(Equal(test.Status))
	})

	It("creates a scenarios ConfigMap", func() {
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		type expectedFields struct {
			name          string
			namespace     string
			scenariosJSON string
			owner         string
		}
		getConfigMapFields := func() (expectedFields, error) {
			cfgMap := new(corev1.ConfigMap)
			err := k8sClient.Get(context.Background(), namespacedName, cfgMap)

			var owner string
			if len(cfgMap.OwnerReferences) > 0 {
				owner = cfgMap.OwnerReferences[0].Name
			}
			return expectedFields{
				name:          cfgMap.Name,
				namespace:     cfgMap.Namespace,
				scenariosJSON: cfgMap.Data["scenarios.json"],
				owner:         owner,
			}, err
		}

		By("checking that the ConfigMap was created correctly")
		Eventually(getConfigMapFields).Should(Equal(expectedFields{
			name:          test.Name,
			namespace:     test.Namespace,
			scenariosJSON: test.Spec.ScenariosJSON,
			owner:         test.Name,
		}))
	})

	It("creates correct number of pods when all are missing", func() {
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		expectedPodCount := 0
		missingPods := status.CheckMissingPods(test, []*corev1.Pod{})
		for range missingPods.Servers {
			expectedPodCount++
		}
		for range missingPods.Clients {
			expectedPodCount++
		}
		if missingPods.Driver != nil {
			expectedPodCount++
		}

		Eventually(func() (int, error) {
			foundPodCount := 0

			list := new(corev1.PodList)
			if err := k8sClient.List(context.Background(), list, client.InNamespace(test.Namespace)); err != nil {
				return 0, err
			}

			for i := range list.Items {
				item := &list.Items[i]
				if item.Labels[config.LoadTestLabel] == test.Name {
					foundPodCount++
				}
			}

			return foundPodCount, nil
		}).Should(Equal(expectedPodCount))
	})

	It("updates the status when pods terminate with errors", func() {
		By("creating a fake environment with errored pods")
		errorState := corev1.ContainerState{
			Terminated: &corev1.ContainerStateTerminated{
				ExitCode: 1,
			},
		}
		builder := podbuilder.New(newDefaults(), test)
		testSpec := &test.Spec
		var pod *corev1.Pod
		for i := range testSpec.Servers {
			pod = builder.PodForServer(&testSpec.Servers[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, errorState)).To(Succeed())
		}
		for i := range testSpec.Clients {
			pod = builder.PodForClient(&testSpec.Clients[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, errorState)).To(Succeed())

		}
		if testSpec.Driver != nil {
			pod = builder.PodForDriver(testSpec.Driver)
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, errorState)).To(Succeed())
		}

		By("waiting for one of the pods to eventually be fetchable")
		Eventually(func() (*corev1.Pod, error) {
			podNamespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
			fetchedPod := new(corev1.Pod)
			if err := k8sClient.Get(context.Background(), podNamespacedName, fetchedPod); err != nil {
				return nil, err
			}
			return fetchedPod, nil
		}).ShouldNot(BeNil())

		By("creating the load test")
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		By("ensuring the test state becomes errored")
		Eventually(func() (grpcv1.LoadTestState, error) {
			fetchedTest := new(grpcv1.LoadTest)
			if err := k8sClient.Get(context.Background(), namespacedName, fetchedTest); err != nil {
				return grpcv1.Unknown, err
			}
			return fetchedTest.Status.State, nil
		}).Should(Equal(grpcv1.Errored))
	})

	It("updates the status when pods are running", func() {
		By("creating a fake environment with running pods")
		runningState := corev1.ContainerState{
			Running: &corev1.ContainerStateRunning{},
		}
		builder := podbuilder.New(newDefaults(), test)
		testSpec := &test.Spec
		var pod *corev1.Pod
		for i := range testSpec.Servers {
			pod = builder.PodForServer(&testSpec.Servers[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, runningState)).To(Succeed())
		}
		for i := range testSpec.Clients {
			pod = builder.PodForClient(&testSpec.Clients[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, runningState)).To(Succeed())

		}
		if testSpec.Driver != nil {
			pod = builder.PodForDriver(testSpec.Driver)
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, runningState)).To(Succeed())
		}

		By("waiting for one of the pods to eventually be fetchable")
		Eventually(func() (*corev1.Pod, error) {
			podNamespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
			fetchedPod := new(corev1.Pod)
			if err := k8sClient.Get(context.Background(), podNamespacedName, fetchedPod); err != nil {
				return nil, err
			}
			return fetchedPod, nil
		}).ShouldNot(BeNil())

		By("creating the load test")
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		By("ensuring the test state becomes running")
		Eventually(func() (grpcv1.LoadTestState, error) {
			fetchedTest := new(grpcv1.LoadTest)
			if err := k8sClient.Get(context.Background(), namespacedName, fetchedTest); err != nil {
				return grpcv1.Unknown, err
			}
			return fetchedTest.Status.State, nil
		}).Should(Equal(grpcv1.Running))
	})

	It("updates the status when pods terminate successfully", func() {
		By("creating a fake environment with finished pods")
		successState := corev1.ContainerState{
			Terminated: &corev1.ContainerStateTerminated{
				ExitCode: 0,
			},
		}
		builder := podbuilder.New(newDefaults(), test)
		testSpec := &test.Spec
		var pod *corev1.Pod
		for i := range testSpec.Servers {
			pod = builder.PodForServer(&testSpec.Servers[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, successState)).To(Succeed())
		}
		for i := range testSpec.Clients {
			pod = builder.PodForClient(&testSpec.Clients[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, successState)).To(Succeed())

		}
		if testSpec.Driver != nil {
			pod = builder.PodForDriver(testSpec.Driver)
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, successState)).To(Succeed())
		}

		By("waiting for one of the pods to eventually be fetchable")
		Eventually(func() (*corev1.Pod, error) {
			podNamespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
			fetchedPod := new(corev1.Pod)
			if err := k8sClient.Get(context.Background(), podNamespacedName, fetchedPod); err != nil {
				return nil, err
			}
			return fetchedPod, nil
		}).ShouldNot(BeNil())

		By("creating the load test")
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		By("ensuring the test state becomes succeeded")
		Eventually(func() (grpcv1.LoadTestState, error) {
			fetchedTest := new(grpcv1.LoadTest)
			if err := k8sClient.Get(context.Background(), namespacedName, fetchedTest); err != nil {
				return grpcv1.Unknown, err
			}
			return fetchedTest.Status.State, nil
		}).Should(Equal(grpcv1.Succeeded))
	})
})
<commit_msg>Add word "test" to disambiguate test cases<commit_after>/*
Copyright 2020 gRPC 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 controllers

import (
	"context"

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

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"

	grpcv1 "github.com/grpc/test-infra/api/v1"
	"github.com/grpc/test-infra/config"
	"github.com/grpc/test-infra/podbuilder"
	"github.com/grpc/test-infra/status"
)

// createPod creates a pod resource, given a pod pointer and a test pointer.
func createPod(pod *corev1.Pod, test *grpcv1.LoadTest) error {
	// TODO: Get the controllerRef to work here.
	// kind := reflect.TypeOf(grpcv1.LoadTest{}).Name()
	// gvk := grpcv1.GroupVersion.WithKind(kind)
	// controllerRef := metav1.NewControllerRef(test, gvk)
	// pod.SetOwnerReferences([]metav1.OwnerReference{*controllerRef})
	return k8sClient.Create(context.Background(), pod)
}

// updatePodWithContainerState changes the container state in the status of a
// pod resource that already exists on the cluster. This is useful for testing
// different failing, running and succeeding states.
func updatePodWithContainerState(pod *corev1.Pod, containerState corev1.ContainerState) error {
	status := &pod.Status
	status.ContainerStatuses = []corev1.ContainerStatus{
		{
			State: containerState,
		},
	}
	return k8sClient.Status().Update(context.Background(), pod)
}

var _ = Describe("LoadTest controller", func() {
	var test *grpcv1.LoadTest
	var namespacedName types.NamespacedName

	BeforeEach(func() {
		test = newLoadTest()
		namespacedName = types.NamespacedName{
			Name:      test.Name,
			Namespace: test.Namespace,
		}
	})

	It("does not change the test status after termination", func() {
		now := metav1.Now()
		test.Status = grpcv1.LoadTestStatus{
			State:     grpcv1.Succeeded,
			StartTime: &now,
			StopTime:  &now,
		}
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())
		Expect(k8sClient.Status().Update(context.Background(), test)).To(Succeed())

		getTestStatus := func() (grpcv1.LoadTestStatus, error) {
			fetchedTest := new(grpcv1.LoadTest)
			err := k8sClient.Get(context.Background(), namespacedName, fetchedTest)
			if err != nil {
				return grpcv1.LoadTestStatus{}, err
			}
			return fetchedTest.Status, nil
		}

		By("ensuring we can eventually get the created status")
		Eventually(getTestStatus).Should(Equal(test.Status))

		By("checking that the expected status remains unchanged")
		Consistently(getTestStatus).Should(Equal(test.Status))
	})

	It("creates a scenarios ConfigMap", func() {
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		type expectedFields struct {
			name          string
			namespace     string
			scenariosJSON string
			owner         string
		}
		getConfigMapFields := func() (expectedFields, error) {
			cfgMap := new(corev1.ConfigMap)
			err := k8sClient.Get(context.Background(), namespacedName, cfgMap)

			var owner string
			if len(cfgMap.OwnerReferences) > 0 {
				owner = cfgMap.OwnerReferences[0].Name
			}
			return expectedFields{
				name:          cfgMap.Name,
				namespace:     cfgMap.Namespace,
				scenariosJSON: cfgMap.Data["scenarios.json"],
				owner:         owner,
			}, err
		}

		By("checking that the ConfigMap was created correctly")
		Eventually(getConfigMapFields).Should(Equal(expectedFields{
			name:          test.Name,
			namespace:     test.Namespace,
			scenariosJSON: test.Spec.ScenariosJSON,
			owner:         test.Name,
		}))
	})

	It("creates correct number of pods when all are missing", func() {
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		expectedPodCount := 0
		missingPods := status.CheckMissingPods(test, []*corev1.Pod{})
		for range missingPods.Servers {
			expectedPodCount++
		}
		for range missingPods.Clients {
			expectedPodCount++
		}
		if missingPods.Driver != nil {
			expectedPodCount++
		}

		Eventually(func() (int, error) {
			foundPodCount := 0

			list := new(corev1.PodList)
			if err := k8sClient.List(context.Background(), list, client.InNamespace(test.Namespace)); err != nil {
				return 0, err
			}

			for i := range list.Items {
				item := &list.Items[i]
				if item.Labels[config.LoadTestLabel] == test.Name {
					foundPodCount++
				}
			}

			return foundPodCount, nil
		}).Should(Equal(expectedPodCount))
	})

	It("updates the test status when pods terminate with errors", func() {
		By("creating a fake environment with errored pods")
		errorState := corev1.ContainerState{
			Terminated: &corev1.ContainerStateTerminated{
				ExitCode: 1,
			},
		}
		builder := podbuilder.New(newDefaults(), test)
		testSpec := &test.Spec
		var pod *corev1.Pod
		for i := range testSpec.Servers {
			pod = builder.PodForServer(&testSpec.Servers[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, errorState)).To(Succeed())
		}
		for i := range testSpec.Clients {
			pod = builder.PodForClient(&testSpec.Clients[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, errorState)).To(Succeed())

		}
		if testSpec.Driver != nil {
			pod = builder.PodForDriver(testSpec.Driver)
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, errorState)).To(Succeed())
		}

		By("waiting for one of the pods to eventually be fetchable")
		Eventually(func() (*corev1.Pod, error) {
			podNamespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
			fetchedPod := new(corev1.Pod)
			if err := k8sClient.Get(context.Background(), podNamespacedName, fetchedPod); err != nil {
				return nil, err
			}
			return fetchedPod, nil
		}).ShouldNot(BeNil())

		By("creating the load test")
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		By("ensuring the test state becomes errored")
		Eventually(func() (grpcv1.LoadTestState, error) {
			fetchedTest := new(grpcv1.LoadTest)
			if err := k8sClient.Get(context.Background(), namespacedName, fetchedTest); err != nil {
				return grpcv1.Unknown, err
			}
			return fetchedTest.Status.State, nil
		}).Should(Equal(grpcv1.Errored))
	})

	It("updates the test status when pods are running", func() {
		By("creating a fake environment with running pods")
		runningState := corev1.ContainerState{
			Running: &corev1.ContainerStateRunning{},
		}
		builder := podbuilder.New(newDefaults(), test)
		testSpec := &test.Spec
		var pod *corev1.Pod
		for i := range testSpec.Servers {
			pod = builder.PodForServer(&testSpec.Servers[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, runningState)).To(Succeed())
		}
		for i := range testSpec.Clients {
			pod = builder.PodForClient(&testSpec.Clients[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, runningState)).To(Succeed())

		}
		if testSpec.Driver != nil {
			pod = builder.PodForDriver(testSpec.Driver)
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, runningState)).To(Succeed())
		}

		By("waiting for one of the pods to eventually be fetchable")
		Eventually(func() (*corev1.Pod, error) {
			podNamespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
			fetchedPod := new(corev1.Pod)
			if err := k8sClient.Get(context.Background(), podNamespacedName, fetchedPod); err != nil {
				return nil, err
			}
			return fetchedPod, nil
		}).ShouldNot(BeNil())

		By("creating the load test")
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		By("ensuring the test state becomes running")
		Eventually(func() (grpcv1.LoadTestState, error) {
			fetchedTest := new(grpcv1.LoadTest)
			if err := k8sClient.Get(context.Background(), namespacedName, fetchedTest); err != nil {
				return grpcv1.Unknown, err
			}
			return fetchedTest.Status.State, nil
		}).Should(Equal(grpcv1.Running))
	})

	It("updates the test status when pods terminate successfully", func() {
		By("creating a fake environment with finished pods")
		successState := corev1.ContainerState{
			Terminated: &corev1.ContainerStateTerminated{
				ExitCode: 0,
			},
		}
		builder := podbuilder.New(newDefaults(), test)
		testSpec := &test.Spec
		var pod *corev1.Pod
		for i := range testSpec.Servers {
			pod = builder.PodForServer(&testSpec.Servers[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, successState)).To(Succeed())
		}
		for i := range testSpec.Clients {
			pod = builder.PodForClient(&testSpec.Clients[i])
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, successState)).To(Succeed())

		}
		if testSpec.Driver != nil {
			pod = builder.PodForDriver(testSpec.Driver)
			Expect(createPod(pod, test)).To(Succeed())
			Expect(updatePodWithContainerState(pod, successState)).To(Succeed())
		}

		By("waiting for one of the pods to eventually be fetchable")
		Eventually(func() (*corev1.Pod, error) {
			podNamespacedName := types.NamespacedName{Name: pod.Name, Namespace: pod.Namespace}
			fetchedPod := new(corev1.Pod)
			if err := k8sClient.Get(context.Background(), podNamespacedName, fetchedPod); err != nil {
				return nil, err
			}
			return fetchedPod, nil
		}).ShouldNot(BeNil())

		By("creating the load test")
		Expect(k8sClient.Create(context.Background(), test)).To(Succeed())

		By("ensuring the test state becomes succeeded")
		Eventually(func() (grpcv1.LoadTestState, error) {
			fetchedTest := new(grpcv1.LoadTest)
			if err := k8sClient.Get(context.Background(), namespacedName, fetchedTest); err != nil {
				return grpcv1.Unknown, err
			}
			return fetchedTest.Status.State, nil
		}).Should(Equal(grpcv1.Succeeded))
	})
})
<|endoftext|>


<commit_before>package collector

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"strconv"
	"strings"
	"github.com/ulricqin/goutils/slicetool"
)

var protocol_search = map[string]string{
	"tcp":  " 00000000:0000 0A",
	"tcp6": " 00000000000000000000000000000000:0000 0A",
	"udp":  " 00000000:0000 07",
	"udp6": " 00000000000000000000000000000000:0000 07",
}

// protocol: ['tcp', 'tcp6', 'udp', 'udp6']
func ListenPorts(protocol string) []int64 {
	ret := []int64{}

	procFile := "/proc/net/" + protocol
	contents, err := ioutil.ReadFile(procFile)
	if err != nil {
		return ret
	}

	reader := bufio.NewReader(bytes.NewBuffer(contents))
	search := protocol_search[protocol]

	for {
		lineBytes, _, err := reader.ReadLine()
		if err == io.EOF {
			break
		}

		rawLine := string(lineBytes)
		idx := strings.Index(rawLine, search)
		if idx < 0 {
			continue
		}

		portStr := rawLine[idx-4 : idx]
		i, err := strconv.ParseInt(portStr, 16, 64)
		if err != nil {
			continue
		}

		ret = append(ret, i)

	}
	return slicetool.SliceUniqueInt64(ret)
}

/*
ulric@ubuntu:~$ cat /proc/net/tcp
  sl  local_address rem_address   st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode
   0: 0100007F:0CEA 00000000:0000 0A 00000000:00000000 00:00000000 00000000   115        0 14563 1 0000000000000000 100 0 0 10 0
   1: 0100007F:0035 00000000:0000 0A 00000000:00000000 00:00000000 00000000     0        0 13321 1 0000000000000000 100 0 0 10 0
   2: 0100007F:12D6 00000000:0000 0A 00000000:00000000 00:00000000 00000000   999        0 9794 1 0000000000000000 100 0 0 10 0
   3: 0100007F:0277 00000000:0000 0A 00000000:00000000 00:00000000 00000000     0        0 39066 1 0000000000000000 100 0 0 10 0
   4: 8284A8C0:89C3 195EBD5B:0050 08 00000000:00000001 00:00000000 00000000  1000        0 19408 1 0000000000000000 20 4 18 2 -1
   5: 8284A8C0:CEDA 017AA8C0:0C38 01 00000000:00000000 02:00001031 00000000  1000        0 119115 2 0000000000000000 21 4 1 10 -1
   6: 8284A8C0:B151 7D807D4A:1466 01 00000000:00000000 02:0000057E 00000000  1000        0 17091 2 0000000000000000 20 4 28 10 -1
ulric@ubuntu:~$ cat /proc/net/tcp6
  sl  local_address                         remote_address                        st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode
   0: 0000000000000000FFFF00000100007F:1F45 00000000000000000000000000000000:0000 0A 00000000:00000000 00:00000000 00000000   116        0 12621 1 0000000000000000 100 0 0 10 -1
   1: 00000000000000000000000000000000:1F90 00000000000000000000000000000000:0000 0A 00000000:00000000 00:00000000 00000000   116        0 10657 1 0000000000000000 100 0 0 10 -1
   2: 00000000000000000000000001000000:0277 00000000000000000000000000000000:0000 0A 00000000:00000000 00:00000000 00000000     0        0 39065 1 0000000000000000 100 0 0 10 -1
ulric@ubuntu:~$ cat /proc/net/udp
  sl  local_address rem_address   st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode ref pointer drops
  944: 00000000:14E9 00000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8549 2 0000000000000000 0
  962: 00000000:BCFB 00000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8551 2 0000000000000000 0
 1788: 0100007F:0035 00000000:0000 07 00000000:00000000 00:00000000 00000000     0        0 13320 2 0000000000000000 0
 1803: 00000000:0044 00000000:0000 07 00000000:00000000 00:00000000 00000000     0        0 7961 2 0000000000000000 0
ulric@ubuntu:~$ cat /proc/net/udp6
  sl  local_address                         remote_address                        st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode ref pointer drops
  329: 00000000000000000000000000000000:9282 00000000000000000000000000000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8552 2 0000000000000000 0
  944: 00000000000000000000000000000000:14E9 00000000000000000000000000000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8550 2 0000000000000000 0
ulric@ubuntu:~$
*/
<commit_msg>break when not found listen port<commit_after>package collector

import (
	"bufio"
	"bytes"
	"github.com/ulricqin/goutils/slicetool"
	"io"
	"io/ioutil"
	"strconv"
	"strings"
)

var protocol_search = map[string]string{
	"tcp":  " 00000000:0000 0A",
	"tcp6": " 00000000000000000000000000000000:0000 0A",
	"udp":  " 00000000:0000 07",
	"udp6": " 00000000000000000000000000000000:0000 07",
}

// protocol: ['tcp', 'tcp6', 'udp', 'udp6']
func ListenPorts(protocol string) []int64 {
	ret := []int64{}

	procFile := "/proc/net/" + protocol
	contents, err := ioutil.ReadFile(procFile)
	if err != nil {
		return ret
	}

	reader := bufio.NewReader(bytes.NewBuffer(contents))
	search := protocol_search[protocol]

	for {
		lineBytes, _, err := reader.ReadLine()
		if err == io.EOF {
			break
		}

		rawLine := string(lineBytes)
		idx := strings.Index(rawLine, search)
		if idx < 0 {
			break
		}

		portStr := rawLine[idx-4 : idx]
		i, err := strconv.ParseInt(portStr, 16, 64)
		if err != nil {
			continue
		}

		ret = append(ret, i)

	}
	return slicetool.SliceUniqueInt64(ret)
}

/*
ulric@ubuntu:~$ cat /proc/net/tcp
  sl  local_address rem_address   st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode
   0: 0100007F:0CEA 00000000:0000 0A 00000000:00000000 00:00000000 00000000   115        0 14563 1 0000000000000000 100 0 0 10 0
   1: 0100007F:0035 00000000:0000 0A 00000000:00000000 00:00000000 00000000     0        0 13321 1 0000000000000000 100 0 0 10 0
   2: 0100007F:12D6 00000000:0000 0A 00000000:00000000 00:00000000 00000000   999        0 9794 1 0000000000000000 100 0 0 10 0
   3: 0100007F:0277 00000000:0000 0A 00000000:00000000 00:00000000 00000000     0        0 39066 1 0000000000000000 100 0 0 10 0
   4: 8284A8C0:89C3 195EBD5B:0050 08 00000000:00000001 00:00000000 00000000  1000        0 19408 1 0000000000000000 20 4 18 2 -1
   5: 8284A8C0:CEDA 017AA8C0:0C38 01 00000000:00000000 02:00001031 00000000  1000        0 119115 2 0000000000000000 21 4 1 10 -1
   6: 8284A8C0:B151 7D807D4A:1466 01 00000000:00000000 02:0000057E 00000000  1000        0 17091 2 0000000000000000 20 4 28 10 -1
ulric@ubuntu:~$ cat /proc/net/tcp6
  sl  local_address                         remote_address                        st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode
   0: 0000000000000000FFFF00000100007F:1F45 00000000000000000000000000000000:0000 0A 00000000:00000000 00:00000000 00000000   116        0 12621 1 0000000000000000 100 0 0 10 -1
   1: 00000000000000000000000000000000:1F90 00000000000000000000000000000000:0000 0A 00000000:00000000 00:00000000 00000000   116        0 10657 1 0000000000000000 100 0 0 10 -1
   2: 00000000000000000000000001000000:0277 00000000000000000000000000000000:0000 0A 00000000:00000000 00:00000000 00000000     0        0 39065 1 0000000000000000 100 0 0 10 -1
ulric@ubuntu:~$ cat /proc/net/udp
  sl  local_address rem_address   st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode ref pointer drops
  944: 00000000:14E9 00000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8549 2 0000000000000000 0
  962: 00000000:BCFB 00000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8551 2 0000000000000000 0
 1788: 0100007F:0035 00000000:0000 07 00000000:00000000 00:00000000 00000000     0        0 13320 2 0000000000000000 0
 1803: 00000000:0044 00000000:0000 07 00000000:00000000 00:00000000 00000000     0        0 7961 2 0000000000000000 0
ulric@ubuntu:~$ cat /proc/net/udp6
  sl  local_address                         remote_address                        st tx_queue rx_queue tr tm->when retrnsmt   uid  timeout inode ref pointer drops
  329: 00000000000000000000000000000000:9282 00000000000000000000000000000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8552 2 0000000000000000 0
  944: 00000000000000000000000000000000:14E9 00000000000000000000000000000000:0000 07 00000000:00000000 00:00000000 00000000   107        0 8550 2 0000000000000000 0
ulric@ubuntu:~$
*/
<|endoftext|>


<commit_before>// Copyright 2018 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 encoding_test

import (
	"flag"
	"fmt"
	"testing"

	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/encoding/prototext"
	pref "google.golang.org/protobuf/reflect/protoreflect"

	tpb "google.golang.org/protobuf/internal/testprotos/test"
)

// The results of these microbenchmarks are unlikely to correspond well
// to real world peformance. They are mainly useful as a quick check to
// detect unexpected regressions and for profiling specific cases.

var benchV1 = flag.Bool("v1", false, "benchmark the v1 implementation")

const maxRecurseLevel = 3

func makeProto() *tpb.TestAllTypes {
	m := &tpb.TestAllTypes{}
	fillMessage(m.ProtoReflect(), 0)
	return m
}

func fillMessage(m pref.Message, level int) {
	if level > maxRecurseLevel {
		return
	}

	fieldDescs := m.Descriptor().Fields()
	for i := 0; i < fieldDescs.Len(); i++ {
		fd := fieldDescs.Get(i)
		switch {
		case fd.IsList():
			setList(m.Mutable(fd).List(), fd, level)
		case fd.IsMap():
			setMap(m.Mutable(fd).Map(), fd, level)
		default:
			setScalarField(m, fd, level)
		}
	}
}

func setScalarField(m pref.Message, fd pref.FieldDescriptor, level int) {
	switch fd.Kind() {
	case pref.MessageKind, pref.GroupKind:
		val := m.NewField(fd)
		fillMessage(val.Message(), level+1)
		m.Set(fd, val)
	default:
		m.Set(fd, scalarField(fd.Kind()))
	}
}

func scalarField(kind pref.Kind) pref.Value {
	switch kind {
	case pref.BoolKind:
		return pref.ValueOfBool(true)

	case pref.Int32Kind, pref.Sint32Kind, pref.Sfixed32Kind:
		return pref.ValueOfInt32(1 << 30)

	case pref.Int64Kind, pref.Sint64Kind, pref.Sfixed64Kind:
		return pref.ValueOfInt64(1 << 30)

	case pref.Uint32Kind, pref.Fixed32Kind:
		return pref.ValueOfUint32(1 << 30)

	case pref.Uint64Kind, pref.Fixed64Kind:
		return pref.ValueOfUint64(1 << 30)

	case pref.FloatKind:
		return pref.ValueOfFloat32(3.14159265)

	case pref.DoubleKind:
		return pref.ValueOfFloat64(3.14159265)

	case pref.BytesKind:
		return pref.ValueOfBytes([]byte("hello world"))

	case pref.StringKind:
		return pref.ValueOfString("hello world")

	case pref.EnumKind:
		return pref.ValueOfEnum(42)
	}

	panic(fmt.Sprintf("FieldDescriptor.Kind %v is not valid", kind))
}

func setList(list pref.List, fd pref.FieldDescriptor, level int) {
	switch fd.Kind() {
	case pref.MessageKind, pref.GroupKind:
		for i := 0; i < 10; i++ {
			val := list.NewElement()
			fillMessage(val.Message(), level+1)
			list.Append(val)
		}
	default:
		for i := 0; i < 100; i++ {
			list.Append(scalarField(fd.Kind()))
		}
	}
}

func setMap(mmap pref.Map, fd pref.FieldDescriptor, level int) {
	fields := fd.Message().Fields()
	keyDesc := fields.ByNumber(1)
	valDesc := fields.ByNumber(2)

	pkey := scalarField(keyDesc.Kind())
	switch kind := valDesc.Kind(); kind {
	case pref.MessageKind, pref.GroupKind:
		val := mmap.NewValue()
		fillMessage(val.Message(), level+1)
		mmap.Set(pkey.MapKey(), val)
	default:
		mmap.Set(pkey.MapKey(), scalarField(kind))
	}
}

func BenchmarkTextEncode(b *testing.B) {
	m := makeProto()
	for i := 0; i < b.N; i++ {
		_, err := prototext.MarshalOptions{Indent: "  "}.Marshal(m)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkTextDecode(b *testing.B) {
	m := makeProto()
	in, err := prototext.MarshalOptions{Indent: "  "}.Marshal(m)
	if err != nil {
		b.Fatal(err)
	}

	for i := 0; i < b.N; i++ {
		m := &tpb.TestAllTypes{}
		if err := prototext.Unmarshal(in, m); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkJSONEncode(b *testing.B) {
	m := makeProto()
	for i := 0; i < b.N; i++ {
		_, err := protojson.MarshalOptions{Indent: "  "}.Marshal(m)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkJSONDecode(b *testing.B) {
	m := makeProto()
	out, err := protojson.MarshalOptions{Indent: "  "}.Marshal(m)
	if err != nil {
		b.Fatal(err)
	}

	for i := 0; i < b.N; i++ {
		m := &tpb.TestAllTypes{}
		if err := protojson.Unmarshal(out, m); err != nil {
			b.Fatal(err)
		}
	}
}
<commit_msg>encoding: remove unused benchV1 flag<commit_after>// Copyright 2018 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 encoding_test

import (
	"fmt"
	"testing"

	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/encoding/prototext"
	pref "google.golang.org/protobuf/reflect/protoreflect"

	tpb "google.golang.org/protobuf/internal/testprotos/test"
)

// The results of these microbenchmarks are unlikely to correspond well
// to real world peformance. They are mainly useful as a quick check to
// detect unexpected regressions and for profiling specific cases.

const maxRecurseLevel = 3

func makeProto() *tpb.TestAllTypes {
	m := &tpb.TestAllTypes{}
	fillMessage(m.ProtoReflect(), 0)
	return m
}

func fillMessage(m pref.Message, level int) {
	if level > maxRecurseLevel {
		return
	}

	fieldDescs := m.Descriptor().Fields()
	for i := 0; i < fieldDescs.Len(); i++ {
		fd := fieldDescs.Get(i)
		switch {
		case fd.IsList():
			setList(m.Mutable(fd).List(), fd, level)
		case fd.IsMap():
			setMap(m.Mutable(fd).Map(), fd, level)
		default:
			setScalarField(m, fd, level)
		}
	}
}

func setScalarField(m pref.Message, fd pref.FieldDescriptor, level int) {
	switch fd.Kind() {
	case pref.MessageKind, pref.GroupKind:
		val := m.NewField(fd)
		fillMessage(val.Message(), level+1)
		m.Set(fd, val)
	default:
		m.Set(fd, scalarField(fd.Kind()))
	}
}

func scalarField(kind pref.Kind) pref.Value {
	switch kind {
	case pref.BoolKind:
		return pref.ValueOfBool(true)

	case pref.Int32Kind, pref.Sint32Kind, pref.Sfixed32Kind:
		return pref.ValueOfInt32(1 << 30)

	case pref.Int64Kind, pref.Sint64Kind, pref.Sfixed64Kind:
		return pref.ValueOfInt64(1 << 30)

	case pref.Uint32Kind, pref.Fixed32Kind:
		return pref.ValueOfUint32(1 << 30)

	case pref.Uint64Kind, pref.Fixed64Kind:
		return pref.ValueOfUint64(1 << 30)

	case pref.FloatKind:
		return pref.ValueOfFloat32(3.14159265)

	case pref.DoubleKind:
		return pref.ValueOfFloat64(3.14159265)

	case pref.BytesKind:
		return pref.ValueOfBytes([]byte("hello world"))

	case pref.StringKind:
		return pref.ValueOfString("hello world")

	case pref.EnumKind:
		return pref.ValueOfEnum(42)
	}

	panic(fmt.Sprintf("FieldDescriptor.Kind %v is not valid", kind))
}

func setList(list pref.List, fd pref.FieldDescriptor, level int) {
	switch fd.Kind() {
	case pref.MessageKind, pref.GroupKind:
		for i := 0; i < 10; i++ {
			val := list.NewElement()
			fillMessage(val.Message(), level+1)
			list.Append(val)
		}
	default:
		for i := 0; i < 100; i++ {
			list.Append(scalarField(fd.Kind()))
		}
	}
}

func setMap(mmap pref.Map, fd pref.FieldDescriptor, level int) {
	fields := fd.Message().Fields()
	keyDesc := fields.ByNumber(1)
	valDesc := fields.ByNumber(2)

	pkey := scalarField(keyDesc.Kind())
	switch kind := valDesc.Kind(); kind {
	case pref.MessageKind, pref.GroupKind:
		val := mmap.NewValue()
		fillMessage(val.Message(), level+1)
		mmap.Set(pkey.MapKey(), val)
	default:
		mmap.Set(pkey.MapKey(), scalarField(kind))
	}
}

func BenchmarkTextEncode(b *testing.B) {
	m := makeProto()
	for i := 0; i < b.N; i++ {
		_, err := prototext.MarshalOptions{Indent: "  "}.Marshal(m)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkTextDecode(b *testing.B) {
	m := makeProto()
	in, err := prototext.MarshalOptions{Indent: "  "}.Marshal(m)
	if err != nil {
		b.Fatal(err)
	}

	for i := 0; i < b.N; i++ {
		m := &tpb.TestAllTypes{}
		if err := prototext.Unmarshal(in, m); err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkJSONEncode(b *testing.B) {
	m := makeProto()
	for i := 0; i < b.N; i++ {
		_, err := protojson.MarshalOptions{Indent: "  "}.Marshal(m)
		if err != nil {
			b.Fatal(err)
		}
	}
}

func BenchmarkJSONDecode(b *testing.B) {
	m := makeProto()
	out, err := protojson.MarshalOptions{Indent: "  "}.Marshal(m)
	if err != nil {
		b.Fatal(err)
	}

	for i := 0; i < b.N; i++ {
		m := &tpb.TestAllTypes{}
		if err := protojson.Unmarshal(out, m); err != nil {
			b.Fatal(err)
		}
	}
}
<|endoftext|>