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

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/options"
	"v.io/v23/rpc"
	"v.io/v23/security"
	"v.io/v23/security/access"
	"v.io/x/ref"
	"v.io/x/ref/runtime/factories/generic"
	"v.io/x/ref/services/identity/identitylib"
	"v.io/x/ref/services/mounttable/mounttablelib"
	"v.io/x/ref/services/xproxy/xproxy"
	"v.io/x/ref/test/expect"
	"v.io/x/ref/test/modules"
)

const (
	stdoutLog = "tmp/runner.stdout.log" // Used as stdout drain when shutting down.
	stderrLog = "tmp/runner.stderr.log" // Used as stderr drain when shutting down.
)

var (
	runTestsWatch bool
)

func init() {
	flag.BoolVar(&runTestsWatch, "runTestsWatch", false, "if true runs the tests in watch mode")
}

var runMT = modules.Register(func(env *modules.Env, args ...string) error {
	ctx, shutdown := v23.Init()
	defer shutdown()
	mp := args[0]

	mt, err := mounttablelib.NewMountTableDispatcher(ctx, "", "", "mounttable")
	if err != nil {
		return fmt.Errorf("mounttablelib.NewMountTableDispatcher failed: %s", err)
	}

	_, server, err := v23.WithNewDispatchingServer(ctx, mp, mt, options.ServesMountTable(true))
	if err != nil {
		return fmt.Errorf("root failed: %v", err)
	}
	fmt.Fprintf(env.Stdout, "PID=%d\n", os.Getpid())
	for _, ep := range server.Status().Endpoints {
		fmt.Fprintf(env.Stdout, "MT_NAME=%s\n", ep.Name())
	}
	modules.WaitForEOF(env.Stdin)
	return nil
}, "runMT")

// Helper function to simply print an error and then exit.
func exitOnError(err error, desc string) {
	if err != nil {
		fmt.Fprintln(os.Stderr, desc, err)
		os.Exit(1)
	}
}

// updateVars captures the vars from the given Handle's stdout and adds them to
// the given vars map, overwriting existing entries.
func updateVars(h modules.Handle, vars map[string]string, varNames ...string) error {
	varsToAdd := map[string]bool{}
	for _, v := range varNames {
		varsToAdd[v] = true
	}
	numLeft := len(varsToAdd)

	s := expect.NewSession(nil, h.Stdout(), 30*time.Second)
	for {
		l := s.ReadLine()
		if err := s.OriginalError(); err != nil {
			return err // EOF or otherwise
		}
		parts := strings.Split(l, "=")
		if len(parts) != 2 {
			return fmt.Errorf("Unexpected line: %s", l)
		}
		if _, ok := varsToAdd[parts[0]]; ok {
			numLeft--
			vars[parts[0]] = parts[1]
			if numLeft == 0 {
				break
			}
		}
	}
	return nil
}

func main() {
	if modules.IsChildProcess() {
		exitOnError(modules.Dispatch(), "Failed to dispatch module")
		return
	}

	// If we ever get a SIGHUP (terminal closes), then end the program.
	signalChannel := make(chan os.Signal)
	signal.Notify(signalChannel, syscall.SIGHUP)
	go func() {
		sig := <-signalChannel
		switch sig {
		case syscall.SIGHUP:
			os.Exit(1)
		}
	}()

	// Try running the program; on failure, exit with error status code.
	if !run() {
		os.Exit(1)
	}
}

// Runs the services and cleans up afterwards.
// Returns true if the run was successful.
func run() bool {
	ctx, shutdown := v23.Init()
	defer shutdown()

	// In order to prevent conflicts, tests and webapp use different mounttable ports.
	port := 8884
	cottagePort := 8885
	housePort := 8886
	host := "localhost"

	// Start a new shell module.
	vars := map[string]string{}
	sh, err := modules.NewShell(ctx, nil, false, nil)
	if err != nil {
		panic(fmt.Sprintf("modules.NewShell: %s", err))
	}

	// Collect the output of this shell on termination.
	err = os.MkdirAll("tmp", 0750)
	exitOnError(err, "Could not make temp directory")
	outFile, err := os.Create(stdoutLog)
	exitOnError(err, "Could not open stdout log file")
	defer outFile.Close()
	errFile, err := os.Create(stderrLog)
	exitOnError(err, "Could not open stderr log file")
	defer errFile.Close()
	defer sh.Cleanup(outFile, errFile)

	// Run a mounttable for tests
	hRoot, err := sh.Start(nil, runMT, "--v23.tcp.protocol=wsh", fmt.Sprintf("--v23.tcp.address=%s:%d", host, port), "root")
	exitOnError(err, "Failed to start root mount table")
	exitOnError(updateVars(hRoot, vars, "MT_NAME"), "Failed to get MT_NAME")
	defer hRoot.Shutdown(outFile, errFile)

	// Set ref.EnvNamespacePrefix env var, consumed downstream.
	sh.SetVar(ref.EnvNamespacePrefix, vars["MT_NAME"])
	v23.GetNamespace(ctx).SetRoots(vars["MT_NAME"])

	// Run the cottage mounttable at host/cottage.
	hCottage, err := sh.Start(nil, runMT, "--v23.tcp.protocol=wsh", fmt.Sprintf("--v23.tcp.address=%s:%d", host, cottagePort), "cottage")
	exitOnError(err, "Failed to start cottage mount table")
	expect.NewSession(nil, hCottage.Stdout(), 30*time.Second)
	defer hCottage.Shutdown(outFile, errFile)

	// run the house mounttable at host/house.
	hHouse, err := sh.Start(nil, runMT, "--v23.tcp.protocol=wsh", fmt.Sprintf("--v23.tcp.address=%s:%d", host, housePort), "house")
	exitOnError(err, "Failed to start house mount table")
	expect.NewSession(nil, hHouse.Stdout(), 30*time.Second)
	defer hHouse.Shutdown(outFile, errFile)

	// Just print out the collected variables. This is for debugging purposes.
	bytes, err := json.Marshal(vars)
	exitOnError(err, "Failed to marshal the collected variables")
	fmt.Println(string(bytes))

	// Also set HOUSE_MOUNTTABLE (used in the tests)
	os.Setenv("HOUSE_MOUNTTABLE", fmt.Sprintf("/%s:%d", host, housePort))

	lspec := v23.GetListenSpec(ctx)
	lspec.Addrs = rpc.ListenAddrs{{"wsh", ":0"}}
	// Allow all processes started by this runner to use the proxy.
	if ref.RPCTransitionState() >= ref.XServers {
		ctx, cancel := context.WithCancel(ctx)
		proxy, err := xproxy.New(ctx, "test/proxy", security.AllowEveryone())
		exitOnError(err, "Failed to start proxy")
		vars["PROXY_NAME"] = proxy.ListeningEndpoints()[0].Name()
		defer func() {
			cancel()
			<-proxy.Closed()
		}()
	} else {
		proxyACL := access.AccessList{In: security.DefaultBlessingPatterns(v23.GetPrincipal(ctx))}
		proxyShutdown, proxyEndpoint, err := generic.NewProxy(ctx, lspec, proxyACL, "test/proxy")
		exitOnError(err, "Failed to start proxy")
		defer proxyShutdown()
		vars["PROXY_NAME"] = proxyEndpoint.Name()
	}

	hIdentityd, err := sh.Start(nil, identitylib.TestIdentityd, "--v23.tcp.protocol=wsh", "--v23.tcp.address=:0", "--v23.proxy=test/proxy", "--http-addr=localhost:0")
	exitOnError(err, "Failed to start identityd")
	exitOnError(updateVars(hIdentityd, vars, "TEST_IDENTITYD_NAME", "TEST_IDENTITYD_HTTP_ADDR"), "Failed to obtain identityd address")
	defer hIdentityd.Shutdown(outFile, errFile)

	// Setup a lot of environment variables; these are used for the tests and building the test extension.
	os.Setenv(ref.EnvNamespacePrefix, vars["MT_NAME"])
	os.Setenv("PROXY_ADDR", vars["PROXY_NAME"])
	os.Setenv("IDENTITYD", fmt.Sprintf("%s/google", vars["TEST_IDENTITYD_NAME"]))
	os.Setenv("IDENTITYD_BLESSING_URL", fmt.Sprintf("%s/auth/blessing-root", vars["TEST_IDENTITYD_HTTP_ADDR"]))
	os.Setenv("DEBUG", "false")

	testsOk := runProva()

	fmt.Println("Cleaning up launched services...")
	return testsOk
}

// Run the prova tests and convert its tap output to xunit.
func runProva() bool {
	// This is also useful information for routing the test output.
	JIRI_ROOT := os.Getenv("JIRI_ROOT")
	VANADIUM_JS := fmt.Sprintf("%s/release/javascript/core", JIRI_ROOT)
	VANADIUM_BROWSER := fmt.Sprintf("%s/release/projects/browser", JIRI_ROOT)

	TAP_XUNIT := fmt.Sprintf("%s/node_modules/.bin/tap-xunit", VANADIUM_BROWSER)
	XUNIT_OUTPUT_FILE := os.Getenv("XUNIT_OUTPUT_FILE")
	if XUNIT_OUTPUT_FILE == "" {
		XUNIT_OUTPUT_FILE = fmt.Sprintf("%s/test_output.xml", os.Getenv("TMPDIR"))
	}
	TAP_XUNIT_OPTIONS := " --package=namespace-browser"

	// Make sure we're in the right folder when we run make test-extension.
	vbroot, err := os.Open(VANADIUM_BROWSER)
	exitOnError(err, "Failed to open vanadium browser dir")
	err = vbroot.Chdir()
	exitOnError(err, "Failed to change to vanadium browser dir")

	// Make the test-extension, this should also remove the old one.
	fmt.Println("Rebuilding test extension...")
	cmdExtensionClean := exec.Command("rm", "-fr", fmt.Sprintf("%s/extension/build-test", VANADIUM_JS))
	err = cmdExtensionClean.Run()
	exitOnError(err, "Failed to clean test extension")
	cmdExtensionBuild := exec.Command("make", "-C", fmt.Sprintf("%s/extension", VANADIUM_JS), "build-test")
	out, err := cmdExtensionBuild.CombinedOutput()
	exitOnError(err, fmt.Sprintf("Failed to build test extension: make -C %s/extension build-test [[\n%s\n]]", VANADIUM_JS, out))

	// These are the basic prova options.
	options := []string{
		"test/**/*.js",
		"--browser",
		"--includeFilenameAsPackage",
		"--launch",
		"chrome",
		"--plugin",
		"proxyquireify/plugin",
		"--transform",
		"envify,./main-transform",
		"--log",
		"tmp/chrome.log",
		fmt.Sprintf("--options=--load-extension=%s/extension/build-test/,--ignore-certificate-errors,--enable-logging=stderr", VANADIUM_JS),
	}

	// Normal tests have a few more options and a different port from the watch tests.
	var PROVA_PORT int
	if !runTestsWatch {
		PROVA_PORT = 8893
		options = append(options, "--headless", "--quit", "--progress", "--tap")
		fmt.Printf("\033[34m-Executing tests. See %s for test xunit output.\033[0m\n", XUNIT_OUTPUT_FILE)
	} else {
		PROVA_PORT = 8894
		fmt.Println("\033[34m-Running tests in watch mode.\033[0m")
	}
	options = append(options, "--port", fmt.Sprintf("%d", PROVA_PORT))

	// This is the prova command.
	cmdProva := exec.Command(
		fmt.Sprintf("%s/node_modules/.bin/prova", VANADIUM_BROWSER),
		options...,
	)
	fmt.Printf("\033[34m-Go to \033[32mhttp://0.0.0.0:%d\033[34m to see tests running.\033[0m\n", PROVA_PORT)
	fmt.Println(cmdProva)

	// Collect the prova stdout. This information needs to be sent to xunit.
	provaOut, err := cmdProva.StdoutPipe()
	exitOnError(err, "Failed to get prova stdout pipe")

	// Setup the tap to xunit command. It uses Prova's stdout as input.
	// The output will got the xunit output file.
	cmdTap := exec.Command(TAP_XUNIT, TAP_XUNIT_OPTIONS)
	cmdTap.Stdin = io.TeeReader(provaOut, os.Stdout) // Tee the prova output to see it on the console too.
	outfile, err := os.Create(XUNIT_OUTPUT_FILE)
	exitOnError(err, "Failed to create xunit output file")
	defer outfile.Close()
	bufferedWriter := bufio.NewWriter(outfile)
	cmdTap.Stdout = bufferedWriter
	defer bufferedWriter.Flush() // Ensure that the full xunit output is written.

	// We start the tap command...
	err = cmdTap.Start()
	exitOnError(err, "Failed to start tap to xunit command")

	// Meanwhile, run Prova to completion. If there was an error, print ERROR, otherwise PASS.
	err = cmdProva.Run()
	testsOk := true
	if err != nil {
		fmt.Println(err)
		fmt.Println("\033[31m\033[1mERROR\033[0m")
		testsOk = false
	} else {
		fmt.Println("\033[32m\033[1mPASS\033[0m")
	}

	// Wait for tap to xunit to finish itself off. This file will be ready for reading by Jenkins.
	fmt.Println("Converting Tap output to XUnit")
	err = cmdTap.Wait()
	exitOnError(err, "Failed tap to xunit conversion")

	return testsOk
}
<commit_msg>browser: Identityd should not be listening on proxy.<commit_after>// Copyright 2015 The Vanadium 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 main

import (
	"bufio"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"v.io/v23"
	"v.io/v23/context"
	"v.io/v23/options"
	"v.io/v23/rpc"
	"v.io/v23/security"
	"v.io/v23/security/access"
	"v.io/x/ref"
	"v.io/x/ref/runtime/factories/generic"
	"v.io/x/ref/services/identity/identitylib"
	"v.io/x/ref/services/mounttable/mounttablelib"
	"v.io/x/ref/services/xproxy/xproxy"
	"v.io/x/ref/test/expect"
	"v.io/x/ref/test/modules"
)

const (
	stdoutLog = "tmp/runner.stdout.log" // Used as stdout drain when shutting down.
	stderrLog = "tmp/runner.stderr.log" // Used as stderr drain when shutting down.
)

var (
	runTestsWatch bool
)

func init() {
	flag.BoolVar(&runTestsWatch, "runTestsWatch", false, "if true runs the tests in watch mode")
}

var runMT = modules.Register(func(env *modules.Env, args ...string) error {
	ctx, shutdown := v23.Init()
	defer shutdown()
	mp := args[0]

	mt, err := mounttablelib.NewMountTableDispatcher(ctx, "", "", "mounttable")
	if err != nil {
		return fmt.Errorf("mounttablelib.NewMountTableDispatcher failed: %s", err)
	}

	_, server, err := v23.WithNewDispatchingServer(ctx, mp, mt, options.ServesMountTable(true))
	if err != nil {
		return fmt.Errorf("root failed: %v", err)
	}
	fmt.Fprintf(env.Stdout, "PID=%d\n", os.Getpid())
	for _, ep := range server.Status().Endpoints {
		fmt.Fprintf(env.Stdout, "MT_NAME=%s\n", ep.Name())
	}
	modules.WaitForEOF(env.Stdin)
	return nil
}, "runMT")

// Helper function to simply print an error and then exit.
func exitOnError(err error, desc string) {
	if err != nil {
		fmt.Fprintln(os.Stderr, desc, err)
		os.Exit(1)
	}
}

// updateVars captures the vars from the given Handle's stdout and adds them to
// the given vars map, overwriting existing entries.
func updateVars(h modules.Handle, vars map[string]string, varNames ...string) error {
	varsToAdd := map[string]bool{}
	for _, v := range varNames {
		varsToAdd[v] = true
	}
	numLeft := len(varsToAdd)

	s := expect.NewSession(nil, h.Stdout(), 30*time.Second)
	for {
		l := s.ReadLine()
		if err := s.OriginalError(); err != nil {
			return err // EOF or otherwise
		}
		parts := strings.Split(l, "=")
		if len(parts) != 2 {
			return fmt.Errorf("Unexpected line: %s", l)
		}
		if _, ok := varsToAdd[parts[0]]; ok {
			numLeft--
			vars[parts[0]] = parts[1]
			if numLeft == 0 {
				break
			}
		}
	}
	return nil
}

func main() {
	if modules.IsChildProcess() {
		exitOnError(modules.Dispatch(), "Failed to dispatch module")
		return
	}

	// If we ever get a SIGHUP (terminal closes), then end the program.
	signalChannel := make(chan os.Signal)
	signal.Notify(signalChannel, syscall.SIGHUP)
	go func() {
		sig := <-signalChannel
		switch sig {
		case syscall.SIGHUP:
			os.Exit(1)
		}
	}()

	// Try running the program; on failure, exit with error status code.
	if !run() {
		os.Exit(1)
	}
}

// Runs the services and cleans up afterwards.
// Returns true if the run was successful.
func run() bool {
	ctx, shutdown := v23.Init()
	defer shutdown()

	// In order to prevent conflicts, tests and webapp use different mounttable ports.
	port := 8884
	cottagePort := 8885
	housePort := 8886
	host := "localhost"

	// Start a new shell module.
	vars := map[string]string{}
	sh, err := modules.NewShell(ctx, nil, false, nil)
	if err != nil {
		panic(fmt.Sprintf("modules.NewShell: %s", err))
	}

	// Collect the output of this shell on termination.
	err = os.MkdirAll("tmp", 0750)
	exitOnError(err, "Could not make temp directory")
	outFile, err := os.Create(stdoutLog)
	exitOnError(err, "Could not open stdout log file")
	defer outFile.Close()
	errFile, err := os.Create(stderrLog)
	exitOnError(err, "Could not open stderr log file")
	defer errFile.Close()
	defer sh.Cleanup(outFile, errFile)

	// Run a mounttable for tests
	hRoot, err := sh.Start(nil, runMT, "--v23.tcp.protocol=wsh", fmt.Sprintf("--v23.tcp.address=%s:%d", host, port), "root")
	exitOnError(err, "Failed to start root mount table")
	exitOnError(updateVars(hRoot, vars, "MT_NAME"), "Failed to get MT_NAME")
	defer hRoot.Shutdown(outFile, errFile)

	// Set ref.EnvNamespacePrefix env var, consumed downstream.
	sh.SetVar(ref.EnvNamespacePrefix, vars["MT_NAME"])
	v23.GetNamespace(ctx).SetRoots(vars["MT_NAME"])

	// Run the cottage mounttable at host/cottage.
	hCottage, err := sh.Start(nil, runMT, "--v23.tcp.protocol=wsh", fmt.Sprintf("--v23.tcp.address=%s:%d", host, cottagePort), "cottage")
	exitOnError(err, "Failed to start cottage mount table")
	expect.NewSession(nil, hCottage.Stdout(), 30*time.Second)
	defer hCottage.Shutdown(outFile, errFile)

	// run the house mounttable at host/house.
	hHouse, err := sh.Start(nil, runMT, "--v23.tcp.protocol=wsh", fmt.Sprintf("--v23.tcp.address=%s:%d", host, housePort), "house")
	exitOnError(err, "Failed to start house mount table")
	expect.NewSession(nil, hHouse.Stdout(), 30*time.Second)
	defer hHouse.Shutdown(outFile, errFile)

	// Just print out the collected variables. This is for debugging purposes.
	bytes, err := json.Marshal(vars)
	exitOnError(err, "Failed to marshal the collected variables")
	fmt.Println(string(bytes))

	// Also set HOUSE_MOUNTTABLE (used in the tests)
	os.Setenv("HOUSE_MOUNTTABLE", fmt.Sprintf("/%s:%d", host, housePort))

	lspec := v23.GetListenSpec(ctx)
	lspec.Addrs = rpc.ListenAddrs{{"wsh", ":0"}}
	// Allow all processes started by this runner to use the proxy.
	if ref.RPCTransitionState() >= ref.XServers {
		ctx, cancel := context.WithCancel(ctx)
		proxy, err := xproxy.New(ctx, "test/proxy", security.AllowEveryone())
		exitOnError(err, "Failed to start proxy")
		vars["PROXY_NAME"] = proxy.ListeningEndpoints()[0].Name()
		defer func() {
			cancel()
			<-proxy.Closed()
		}()
	} else {
		proxyACL := access.AccessList{In: security.DefaultBlessingPatterns(v23.GetPrincipal(ctx))}
		proxyShutdown, proxyEndpoint, err := generic.NewProxy(ctx, lspec, proxyACL, "test/proxy")
		exitOnError(err, "Failed to start proxy")
		defer proxyShutdown()
		vars["PROXY_NAME"] = proxyEndpoint.Name()
	}

	hIdentityd, err := sh.Start(nil, identitylib.TestIdentityd, "--v23.tcp.protocol=wsh", "--v23.tcp.address=:0", "--http-addr=localhost:0")
	exitOnError(err, "Failed to start identityd")
	exitOnError(updateVars(hIdentityd, vars, "TEST_IDENTITYD_NAME", "TEST_IDENTITYD_HTTP_ADDR"), "Failed to obtain identityd address")
	defer hIdentityd.Shutdown(outFile, errFile)

	// Setup a lot of environment variables; these are used for the tests and building the test extension.
	os.Setenv(ref.EnvNamespacePrefix, vars["MT_NAME"])
	os.Setenv("PROXY_ADDR", vars["PROXY_NAME"])
	os.Setenv("IDENTITYD", fmt.Sprintf("%s/google", vars["TEST_IDENTITYD_NAME"]))
	os.Setenv("IDENTITYD_BLESSING_URL", fmt.Sprintf("%s/auth/blessing-root", vars["TEST_IDENTITYD_HTTP_ADDR"]))
	os.Setenv("DEBUG", "false")

	testsOk := runProva()

	fmt.Println("Cleaning up launched services...")
	return testsOk
}

// Run the prova tests and convert its tap output to xunit.
func runProva() bool {
	// This is also useful information for routing the test output.
	JIRI_ROOT := os.Getenv("JIRI_ROOT")
	VANADIUM_JS := fmt.Sprintf("%s/release/javascript/core", JIRI_ROOT)
	VANADIUM_BROWSER := fmt.Sprintf("%s/release/projects/browser", JIRI_ROOT)

	TAP_XUNIT := fmt.Sprintf("%s/node_modules/.bin/tap-xunit", VANADIUM_BROWSER)
	XUNIT_OUTPUT_FILE := os.Getenv("XUNIT_OUTPUT_FILE")
	if XUNIT_OUTPUT_FILE == "" {
		XUNIT_OUTPUT_FILE = fmt.Sprintf("%s/test_output.xml", os.Getenv("TMPDIR"))
	}
	TAP_XUNIT_OPTIONS := " --package=namespace-browser"

	// Make sure we're in the right folder when we run make test-extension.
	vbroot, err := os.Open(VANADIUM_BROWSER)
	exitOnError(err, "Failed to open vanadium browser dir")
	err = vbroot.Chdir()
	exitOnError(err, "Failed to change to vanadium browser dir")

	// Make the test-extension, this should also remove the old one.
	fmt.Println("Rebuilding test extension...")
	cmdExtensionClean := exec.Command("rm", "-fr", fmt.Sprintf("%s/extension/build-test", VANADIUM_JS))
	err = cmdExtensionClean.Run()
	exitOnError(err, "Failed to clean test extension")
	cmdExtensionBuild := exec.Command("make", "-C", fmt.Sprintf("%s/extension", VANADIUM_JS), "build-test")
	out, err := cmdExtensionBuild.CombinedOutput()
	exitOnError(err, fmt.Sprintf("Failed to build test extension: make -C %s/extension build-test [[\n%s\n]]", VANADIUM_JS, out))

	// These are the basic prova options.
	options := []string{
		"test/**/*.js",
		"--browser",
		"--includeFilenameAsPackage",
		"--launch",
		"chrome",
		"--plugin",
		"proxyquireify/plugin",
		"--transform",
		"envify,./main-transform",
		"--log",
		"tmp/chrome.log",
		fmt.Sprintf("--options=--load-extension=%s/extension/build-test/,--ignore-certificate-errors,--enable-logging=stderr", VANADIUM_JS),
	}

	// Normal tests have a few more options and a different port from the watch tests.
	var PROVA_PORT int
	if !runTestsWatch {
		PROVA_PORT = 8893
		options = append(options, "--headless", "--quit", "--progress", "--tap")
		fmt.Printf("\033[34m-Executing tests. See %s for test xunit output.\033[0m\n", XUNIT_OUTPUT_FILE)
	} else {
		PROVA_PORT = 8894
		fmt.Println("\033[34m-Running tests in watch mode.\033[0m")
	}
	options = append(options, "--port", fmt.Sprintf("%d", PROVA_PORT))

	// This is the prova command.
	cmdProva := exec.Command(
		fmt.Sprintf("%s/node_modules/.bin/prova", VANADIUM_BROWSER),
		options...,
	)
	fmt.Printf("\033[34m-Go to \033[32mhttp://0.0.0.0:%d\033[34m to see tests running.\033[0m\n", PROVA_PORT)
	fmt.Println(cmdProva)

	// Collect the prova stdout. This information needs to be sent to xunit.
	provaOut, err := cmdProva.StdoutPipe()
	exitOnError(err, "Failed to get prova stdout pipe")

	// Setup the tap to xunit command. It uses Prova's stdout as input.
	// The output will got the xunit output file.
	cmdTap := exec.Command(TAP_XUNIT, TAP_XUNIT_OPTIONS)
	cmdTap.Stdin = io.TeeReader(provaOut, os.Stdout) // Tee the prova output to see it on the console too.
	outfile, err := os.Create(XUNIT_OUTPUT_FILE)
	exitOnError(err, "Failed to create xunit output file")
	defer outfile.Close()
	bufferedWriter := bufio.NewWriter(outfile)
	cmdTap.Stdout = bufferedWriter
	defer bufferedWriter.Flush() // Ensure that the full xunit output is written.

	// We start the tap command...
	err = cmdTap.Start()
	exitOnError(err, "Failed to start tap to xunit command")

	// Meanwhile, run Prova to completion. If there was an error, print ERROR, otherwise PASS.
	err = cmdProva.Run()
	testsOk := true
	if err != nil {
		fmt.Println(err)
		fmt.Println("\033[31m\033[1mERROR\033[0m")
		testsOk = false
	} else {
		fmt.Println("\033[32m\033[1mPASS\033[0m")
	}

	// Wait for tap to xunit to finish itself off. This file will be ready for reading by Jenkins.
	fmt.Println("Converting Tap output to XUnit")
	err = cmdTap.Wait()
	exitOnError(err, "Failed tap to xunit conversion")

	return testsOk
}
<|endoftext|>


<commit_before>package common

import (
	"image"
	"image/png"
	"io"
	"path"
	"strings"

	"github.com/boltdb/bolt"
	"go.skia.org/infra/go/fileutil"
	"go.skia.org/infra/golden/go/diff"
	"go.skia.org/infra/golden/go/types"
)

const (
	// IMG_EXTENSION is the default extension of images.
	IMG_EXTENSION = "png"

	// DiffImageSeparator is the character that separates two image ids in the
	// resulting diff image.
	DiffImageSeparator = "-"
)

// EncodeImg encodes the given image as a PNG and writes the result to the
// given writer.
func EncodeImg(w io.Writer, img *image.NRGBA) error {
	encoder := png.Encoder{CompressionLevel: png.BestSpeed}
	if err := encoder.Encode(w, img); err != nil {
		return err
	}
	return nil
}

// DecodeImg decodes an image from the given reader and returns it as a NRGBA image.
func DecodeImg(reader io.Reader) (*image.NRGBA, error) {
	im, err := png.Decode(reader)
	if err != nil {
		return nil, err
	}
	return diff.GetNRGBA(im), nil
}

// OpenBoltDB opens a boltDB in the given given directory with the given name.
func OpenBoltDB(baseDir, name string) (*bolt.DB, error) {
	baseDir, err := fileutil.EnsureDirExists(baseDir)
	if err != nil {
		return nil, err
	}

	return bolt.Open(path.Join(baseDir, name), 0600, nil)
}

func AsStrings(xd types.DigestSlice) []string {
	s := make([]string, 0, len(xd))
	for _, d := range xd {
		s = append(s, string(d))
	}
	return s
}

// Takes two image IDs and returns a unique diff ID.
// Note: DiffID(a,b) == DiffID(b, a) holds.
func DiffID(left, right types.Digest) string {
	_, _, diffID := getOrderedDiffID(left, right)
	return diffID
}

// Inverse function of DiffID.
// SplitDiffID(DiffID(a,b)) deterministically returns (a,b) or (b,a).
func SplitDiffID(diffID string) (types.Digest, types.Digest) {
	imageIDs := strings.Split(diffID, DiffImageSeparator)

	return types.Digest(imageIDs[0]), types.Digest(imageIDs[1])
}

func getOrderedDiffID(left, right types.Digest) (types.Digest, types.Digest, string) {
	if right < left {
		// Make sure the smaller digest is left imageID.
		left, right = right, left
	}
	return left, right, string(left) + DiffImageSeparator + string(right)
}
<commit_msg>[gold] Remove unused function common.OpenBoltDB().<commit_after>package common

import (
	"image"
	"image/png"
	"io"
	"strings"

	"go.skia.org/infra/golden/go/diff"
	"go.skia.org/infra/golden/go/types"
)

const (
	// IMG_EXTENSION is the default extension of images.
	IMG_EXTENSION = "png"

	// DiffImageSeparator is the character that separates two image ids in the
	// resulting diff image.
	DiffImageSeparator = "-"
)

// EncodeImg encodes the given image as a PNG and writes the result to the
// given writer.
func EncodeImg(w io.Writer, img *image.NRGBA) error {
	encoder := png.Encoder{CompressionLevel: png.BestSpeed}
	if err := encoder.Encode(w, img); err != nil {
		return err
	}
	return nil
}

// DecodeImg decodes an image from the given reader and returns it as a NRGBA image.
func DecodeImg(reader io.Reader) (*image.NRGBA, error) {
	im, err := png.Decode(reader)
	if err != nil {
		return nil, err
	}
	return diff.GetNRGBA(im), nil
}

func AsStrings(xd types.DigestSlice) []string {
	s := make([]string, 0, len(xd))
	for _, d := range xd {
		s = append(s, string(d))
	}
	return s
}

// Takes two image IDs and returns a unique diff ID.
// Note: DiffID(a,b) == DiffID(b, a) holds.
func DiffID(left, right types.Digest) string {
	_, _, diffID := getOrderedDiffID(left, right)
	return diffID
}

// Inverse function of DiffID.
// SplitDiffID(DiffID(a,b)) deterministically returns (a,b) or (b,a).
func SplitDiffID(diffID string) (types.Digest, types.Digest) {
	imageIDs := strings.Split(diffID, DiffImageSeparator)

	return types.Digest(imageIDs[0]), types.Digest(imageIDs[1])
}

func getOrderedDiffID(left, right types.Digest) (types.Digest, types.Digest, string) {
	if right < left {
		// Make sure the smaller digest is left imageID.
		left, right = right, left
	}
	return left, right, string(left) + DiffImageSeparator + string(right)
}
<|endoftext|>


<commit_before>package ginmon

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

const TestMode string = "test"

const checkMark = "\u2713"
const ballotX = "\u2717"

const testpath = "/foo/bar"

func internalGinCtx() *gin.Context {
	return &gin.Context{
		Request: &http.Request{
			URL: &url.URL{
				Path: testpath,
			},
		},
	}
}

func Test_Inc(t *testing.T) {
	ca := NewCounterAspect()
	expect := 1
	ca.Inc(internalGinCtx())
	ca.reset()
	if assert.Equal(t, ca.RequestsSum, expect, "Incrementation of counter does not work, expect %d but got %d %s",
		expect, ca.RequestsSum, ballotX) {
		t.Logf("Incrementation of counter works, expect %d and git %d %s",
			expect, ca.RequestsSum, checkMark)
	}
}

func Test_GetStats(t *testing.T) {
	ca := NewCounterAspect()
	if assert.NotNil(t, ca.GetStats(), "Return of Getstats() should not be nil") {
		t.Logf("Should be an interface %s", checkMark)
	}

	newCa := ca.GetStats().(CounterAspect)
	expect := 0
	if assert.Equal(t, newCa.RequestsSum, expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}

	ca.Inc(internalGinCtx())
	if assert.Equal(t, newCa.RequestsSum, expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, newCa.Requests[testpath], expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}

	ca.reset()
	newCa = ca.GetStats().(CounterAspect)
	expect = 1
	if assert.Equal(t, newCa.RequestsSum, expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, newCa.Requests[testpath], expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}
}

func Test_Name(t *testing.T) {
	ca := NewCounterAspect()
	expect := "Counter"
	if assert.Equal(t, ca.Name(), expect, "Return of counter name does not work, expect %s but got %s %s",
		expect, ca.Name(), ballotX) {
		t.Logf("Return of counter name works, expect %s and got %s %s",
			expect, ca.Name(), checkMark)
	}
}

func Test_InRoot(t *testing.T) {
	ca := NewCounterAspect()
	expect := false
	if assert.Equal(t, ca.InRoot(), expect, "Expect %v but got %v %s",
		expect, ca.InRoot(), ballotX) {
		t.Logf("Expect %v and got %v %s",
			expect, ca.InRoot(), checkMark)
	}
}

func Test_CounterHandler(t *testing.T) {
	gin.SetMode(TestMode)
	router := gin.New()
	ca := NewCounterAspect()
	expect := 1
	ca.Inc(internalGinCtx())
	ca.reset()

	router.Use(CounterHandler(ca))
	tryRequest(router, "GET", "/")
	if assert.Equal(t, ca.RequestsSum, expect, "Incrementation of counter does not work, expect %d but got %d %s", expect, ca.RequestsSum, ballotX) {
		t.Logf("CounterHandler works, expect %d and got %d %s", expect, ca.RequestsSum, checkMark)
	}
}

func tryRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}
<commit_msg>make counter test work with channel<commit_after>package ginmon

import (
	"net/http"
	"net/http/httptest"
	"net/url"
	"testing"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/assert"
)

const TestMode string = "test"

const checkMark = "\u2713"
const ballotX = "\u2717"

const testpath = "/foo/bar"

func internalGinCtx() *gin.Context {
	return &gin.Context{
		Request: &http.Request{
			URL: &url.URL{
				Path: testpath,
			},
		},
	}
}

func Test_Inc(t *testing.T) {
	ca := NewCounterAspect()
	ca.StartTimer(1 * time.Second)
	expect := 1
	ca.inc <- tuple{
		path: testpath,
		code: 404,
	}
	ca.reset()
	if assert.Equal(t, ca.RequestsSum, expect, "Incrementation of counter does not work, expect %d but got %d %s",
		expect, ca.RequestsSum, ballotX) {
		t.Logf("Incrementation of counter works, expect %d and git %d %s",
			expect, ca.RequestsSum, checkMark)
	}
}

func Test_GetStats(t *testing.T) {
	ca := NewCounterAspect()
	ca.StartTimer(1 * time.Second)
	if assert.NotNil(t, ca.GetStats(), "Return of Getstats() should not be nil") {
		t.Logf("Should be an interface %s", checkMark)
	}

	newCa := ca.GetStats().(CounterAspect)
	expect := 0
	if assert.Equal(t, newCa.RequestsSum, expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}

	ca.inc <- tuple{
		path: testpath,
		code: 404,
	}
	if assert.Equal(t, newCa.RequestsSum, expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, newCa.Requests[testpath], expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}

	ca.reset()
	newCa = ca.GetStats().(CounterAspect)
	expect = 1
	if assert.Equal(t, newCa.RequestsSum, expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.RequestsSum, ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.RequestsSum, checkMark)
	}
	if assert.Equal(t, newCa.Requests[testpath], expect, "Return of Getstats() does not work, expect %d but got %d %s",
		expect, newCa.Requests[testpath], ballotX) {
		t.Logf("Return of Getstats() works, expect %d and got %d %s",
			expect, newCa.Requests[testpath], checkMark)
	}
}

func Test_Name(t *testing.T) {
	ca := NewCounterAspect()
	expect := "Counter"
	if assert.Equal(t, ca.Name(), expect, "Return of counter name does not work, expect %s but got %s %s",
		expect, ca.Name(), ballotX) {
		t.Logf("Return of counter name works, expect %s and got %s %s",
			expect, ca.Name(), checkMark)
	}
}

func Test_InRoot(t *testing.T) {
	ca := NewCounterAspect()
	expect := false
	if assert.Equal(t, ca.InRoot(), expect, "Expect %v but got %v %s",
		expect, ca.InRoot(), ballotX) {
		t.Logf("Expect %v and got %v %s",
			expect, ca.InRoot(), checkMark)
	}
}

func Test_CounterHandler(t *testing.T) {
	gin.SetMode(TestMode)
	router := gin.New()
	ca := NewCounterAspect()
	ca.StartTimer(1 * time.Second)
	expect := 1
	ca.inc <- tuple{
		path: testpath,
		code: 404,
	}
	ca.reset()

	router.Use(CounterHandler(ca))
	tryRequest(router, "GET", "/")
	if assert.Equal(t, ca.RequestsSum, expect, "Incrementation of counter does not work, expect %d but got %d %s", expect, ca.RequestsSum, ballotX) {
		t.Logf("CounterHandler works, expect %d and got %d %s", expect, ca.RequestsSum, checkMark)
	}
}

func tryRequest(r http.Handler, method, path string) *httptest.ResponseRecorder {
	req, _ := http.NewRequest(method, path, nil)
	w := httptest.NewRecorder()
	r.ServeHTTP(w, req)
	return w
}
<|endoftext|>


<commit_before>// +build benchmark,!no-etcd,!integration

/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package master

import (
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/golang/glog"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/test/integration/framework"
)

// Command line flag globals, parsed in init and used by the benchmarks:
//	* pods && !tasks: Start -pods, scale number of parallel tasks with b.N
//  * !pods && tasks: Start -tasks, scale pods with b.N
//  * pods && tasks: Ignore b.N, benchmark behaves like a test constrained by -benchtime.
//  * !pods && !tasks: scale pods and workers with b.N.
// -workers specifies the number of workers to shard tasks across.
// Eg: go test bench . -bench-pods 3000 -bench-tasks 100 -bench-tasks 10:
// Create 100 tasks each listing 3000 pods, and run them 10 at a time.
var (
	Workers int
	Pods    int
	Tasks   int
)

const Glog_fatalf = 3

func init() {
	q := flag.Int("bench-quiet", 3, "Quietness, don't glog severities <= given level during the benchmark.")
	pods := flag.Int("bench-pods", -1, "Number of pods for the benchmark. If unspecified, uses b.N.")
	workers := flag.Int("bench-workers", -1, "Number workers for the benchmark. If unspecified, uses tasks.")
	tasks := flag.Int("bench-tasks", -1, "Number of tasks for the benchmark. These tasks are sharded across the workers. If unspecified, uses b.N.")
	flag.Parse()

	// Unfortunately this v level goes in the opposite direction as stderrthreshold.
	flag.Set("v", fmt.Sprintf("%d", *q))

	// We need quiet logs to parse benchmark results, which includes Errorf.
	flag.Set("logtostderr", "false")
	flag.Set("stderrthreshold", fmt.Sprintf("%d", Glog_fatalf-*q))
	Pods = *pods
	Workers = *workers
	Tasks = *tasks
}

// getPods returns the cmd line -pods or b.N if -pods wasn't specified.
// Benchmarks can call getPods to get the number of pods they need to
// create for a given benchmark.
func getPods(bN int) int {
	if Pods < 0 {
		return bN
	}
	return Pods
}

// getTasks returns the cmd line -workers or b.N if -workers wasn't specified.
// Benchmarks would call getTasks to get the number of workers required to
// perform the benchmark in parallel.
func getTasks(bN int) int {
	if Tasks < 0 {
		return bN
	}
	return Tasks
}

// getIterations returns the number of iterations required by each benchmark for
// go to produce reliable timing results.
func getIterations(bN int) int {
	// Anything with constant pods is only linear if we iterate b.N times.
	if Pods > 0 {
		return bN
	}
	return 1
}

// startPodsOnNodes creates numPods sharded across numNodes
func startPodsOnNodes(ns string, numPods, numNodes int, restClient *client.Client) {
	podsPerNode := numPods / numNodes
	if podsPerNode < 1 {
		podsPerNode = 1
	}
	framework.RunParallel(func(id int) error {
		return framework.StartPods(podsPerNode, fmt.Sprintf("host.%d", id), restClient)
	}, numNodes, -1)
}

// Benchmark pod listing by waiting on `Tasks` listers to list `Pods` pods via `Workers`.
func BenchmarkPodList(b *testing.B) {
	b.StopTimer()
	m := framework.NewMasterComponents(&framework.Config{nil, true, false, 250.0, 500})
	defer m.Stop(true, true)

	ns := framework.CreateTestingNamespace("benchmark-pod-list", s, t)
	defer framework.DeleteTestingNamespace(ns, s, t)

	numPods, numTasks, iter := getPods(b.N), getTasks(b.N), getIterations(b.N)
	podsPerNode := numPods / numTasks
	if podsPerNode < 1 {
		podsPerNode = 1
	}
	glog.Infof("Starting benchmark: b.N %d, pods %d, workers %d, podsPerNode %d",
		b.N, numPods, numTasks, podsPerNode)

	startPodsOnNodes(ns.Name, numPods, numTasks, m.RestClient)
	// Stop the rc manager so it doesn't steal resources
	m.Stop(false, true)

	b.StartTimer()
	for i := 0; i < iter; i++ {
		framework.RunParallel(func(id int) error {
			host := fmt.Sprintf("host.%d", id)
			now := time.Now()
			defer func() {
				glog.V(3).Infof("Worker %d: Node %v listing pods took %v", id, host, time.Since(now))
			}()
			if pods, err := m.RestClient.Pods(ns.Name).List(
				labels.Everything(),
				fields.OneTermEqualSelector(client.PodHost, host)); err != nil {
				return err
			} else if len(pods.Items) < podsPerNode {
				glog.Fatalf("List retrieved %d pods, which is less than %d", len(pods.Items), podsPerNode)
			}
			return nil
		}, numTasks, Workers)
	}
	b.StopTimer()
}

// Benchmark pod listing by waiting on `Tasks` listers to list `Pods` pods via `Workers`.
func BenchmarkPodListEtcd(b *testing.B) {
	b.StopTimer()
	m := framework.NewMasterComponents(&framework.Config{nil, true, false, 250.0, 500})
	defer m.Stop(true, true)

	ns := framework.CreateTestingNamespace("benchmark-pod-list-etcd", s, t)
	defer framework.DeleteTestingNamespace(ns, s, t)

	numPods, numTasks, iter := getPods(b.N), getTasks(b.N), getIterations(b.N)
	podsPerNode := numPods / numTasks
	if podsPerNode < 1 {
		podsPerNode = 1
	}

	startPodsOnNodes(ns.Name, numPods, numTasks, m.RestClient)
	// Stop the rc manager so it doesn't steal resources
	m.Stop(false, true)

	glog.Infof("Starting benchmark: b.N %d, pods %d, workers %d, podsPerNode %d",
		b.N, numPods, numTasks, podsPerNode)

	b.StartTimer()
	for i := 0; i < iter; i++ {
		framework.RunParallel(func(id int) error {
			now := time.Now()
			defer func() {
				glog.V(3).Infof("Worker %d: listing pods took %v", id, time.Since(now))
			}()
			pods, err := m.RestClient.Pods(ns.Name).List(labels.Everything(), fields.Everything())
			if err != nil {
				return err
			}
			if len(pods.Items) < numPods {
				glog.Fatalf("List retrieved %d pods, which is less than %d", len(pods.Items), numPods)
			}
			return nil
		}, numTasks, Workers)
	}
	b.StopTimer()
}

func TestMain(m *testing.M) {
	os.Exit(m.Run())
}
<commit_msg>Compile fixes for master_benchmark_test.go<commit_after>// +build benchmark,!no-etcd,!integration

/*
Copyright 2014 The Kubernetes Authors.

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

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

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

package master

import (
	"flag"
	"fmt"
	"os"
	"testing"
	"time"

	"github.com/golang/glog"
	"k8s.io/kubernetes/pkg/api"
	client "k8s.io/kubernetes/pkg/client/unversioned"
	"k8s.io/kubernetes/pkg/fields"
	"k8s.io/kubernetes/pkg/labels"
	"k8s.io/kubernetes/test/integration/framework"
)

// Command line flag globals, parsed in init and used by the benchmarks:
//	* pods && !tasks: Start -pods, scale number of parallel tasks with b.N
//  * !pods && tasks: Start -tasks, scale pods with b.N
//  * pods && tasks: Ignore b.N, benchmark behaves like a test constrained by -benchtime.
//  * !pods && !tasks: scale pods and workers with b.N.
// -workers specifies the number of workers to shard tasks across.
// Eg: go test bench . -bench-pods 3000 -bench-tasks 100 -bench-tasks 10:
// Create 100 tasks each listing 3000 pods, and run them 10 at a time.
var (
	Workers int
	Pods    int
	Tasks   int
)

const Glog_fatalf = 3

func init() {
	q := flag.Int("bench-quiet", 3, "Quietness, don't glog severities <= given level during the benchmark.")
	pods := flag.Int("bench-pods", -1, "Number of pods for the benchmark. If unspecified, uses b.N.")
	workers := flag.Int("bench-workers", -1, "Number workers for the benchmark. If unspecified, uses tasks.")
	tasks := flag.Int("bench-tasks", -1, "Number of tasks for the benchmark. These tasks are sharded across the workers. If unspecified, uses b.N.")
	flag.Parse()

	// Unfortunately this v level goes in the opposite direction as stderrthreshold.
	flag.Set("v", fmt.Sprintf("%d", *q))

	// We need quiet logs to parse benchmark results, which includes Errorf.
	flag.Set("logtostderr", "false")
	flag.Set("stderrthreshold", fmt.Sprintf("%d", Glog_fatalf-*q))
	Pods = *pods
	Workers = *workers
	Tasks = *tasks
}

// getPods returns the cmd line -pods or b.N if -pods wasn't specified.
// Benchmarks can call getPods to get the number of pods they need to
// create for a given benchmark.
func getPods(bN int) int {
	if Pods < 0 {
		return bN
	}
	return Pods
}

// getTasks returns the cmd line -workers or b.N if -workers wasn't specified.
// Benchmarks would call getTasks to get the number of workers required to
// perform the benchmark in parallel.
func getTasks(bN int) int {
	if Tasks < 0 {
		return bN
	}
	return Tasks
}

// getIterations returns the number of iterations required by each benchmark for
// go to produce reliable timing results.
func getIterations(bN int) int {
	// Anything with constant pods is only linear if we iterate b.N times.
	if Pods > 0 {
		return bN
	}
	return 1
}

// startPodsOnNodes creates numPods sharded across numNodes
func startPodsOnNodes(ns string, numPods, numNodes int, restClient *client.Client) {
	podsPerNode := numPods / numNodes
	if podsPerNode < 1 {
		podsPerNode = 1
	}
	framework.RunParallel(func(id int) error {
		return framework.StartPods(podsPerNode, fmt.Sprintf("host.%d", id), restClient)
	}, numNodes, -1)
}

// Benchmark pod listing by waiting on `Tasks` listers to list `Pods` pods via `Workers`.
func BenchmarkPodList(b *testing.B) {
	b.StopTimer()
	m := framework.NewMasterComponents(&framework.Config{nil, true, false, 250.0, 500})
	defer m.Stop(true, true)

	ns := framework.CreateTestingNamespace("benchmark-pod-list", s, t)
	defer framework.DeleteTestingNamespace(ns, s, t)

	numPods, numTasks, iter := getPods(b.N), getTasks(b.N), getIterations(b.N)
	podsPerNode := numPods / numTasks
	if podsPerNode < 1 {
		podsPerNode = 1
	}
	glog.Infof("Starting benchmark: b.N %d, pods %d, workers %d, podsPerNode %d",
		b.N, numPods, numTasks, podsPerNode)

	startPodsOnNodes(ns.Name, numPods, numTasks, m.RestClient)
	// Stop the rc manager so it doesn't steal resources
	m.Stop(false, true)

	b.StartTimer()
	for i := 0; i < iter; i++ {
		framework.RunParallel(func(id int) error {
			host := fmt.Sprintf("host.%d", id)
			now := time.Now()
			defer func() {
				glog.V(3).Infof("Worker %d: Node %v listing pods took %v", id, host, time.Since(now))
			}()
			if pods, err := m.RestClient.Pods(ns.Name).List(api.ListOptions{
				LabelSelector: labels.Everything(),
				FieldSelector: fields.OneTermEqualSelector(api.PodHostField, host),
			}); err != nil {
				return err
			} else if len(pods.Items) < podsPerNode {
				glog.Fatalf("List retrieved %d pods, which is less than %d", len(pods.Items), podsPerNode)
			}
			return nil
		}, numTasks, Workers)
	}
	b.StopTimer()
}

// Benchmark pod listing by waiting on `Tasks` listers to list `Pods` pods via `Workers`.
func BenchmarkPodListEtcd(b *testing.B) {
	b.StopTimer()
	m := framework.NewMasterComponents(&framework.Config{nil, true, false, 250.0, 500})
	defer m.Stop(true, true)

	ns := framework.CreateTestingNamespace("benchmark-pod-list-etcd", s, t)
	defer framework.DeleteTestingNamespace(ns, s, t)

	numPods, numTasks, iter := getPods(b.N), getTasks(b.N), getIterations(b.N)
	podsPerNode := numPods / numTasks
	if podsPerNode < 1 {
		podsPerNode = 1
	}

	startPodsOnNodes(ns.Name, numPods, numTasks, m.RestClient)
	// Stop the rc manager so it doesn't steal resources
	m.Stop(false, true)

	glog.Infof("Starting benchmark: b.N %d, pods %d, workers %d, podsPerNode %d",
		b.N, numPods, numTasks, podsPerNode)

	b.StartTimer()
	for i := 0; i < iter; i++ {
		framework.RunParallel(func(id int) error {
			now := time.Now()
			defer func() {
				glog.V(3).Infof("Worker %d: listing pods took %v", id, time.Since(now))
			}()
			pods, err := m.RestClient.Pods(ns.Name).List(api.ListOptions{
				LabelSelector: labels.Everything(),
				FieldSelector: fields.Everything(),
			})
			if err != nil {
				return err
			}
			if len(pods.Items) < numPods {
				glog.Fatalf("List retrieved %d pods, which is less than %d", len(pods.Items), numPods)
			}
			return nil
		}, numTasks, Workers)
	}
	b.StopTimer()
}

func TestMain(m *testing.M) {
	os.Exit(m.Run())
}
<|endoftext|>


<commit_before>package checkdisk

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
	gpud "github.com/shirou/gopsutil/disk"
)

var opts struct {
	Warning       *string `short:"w" long:"warning" value-name:"N, N%" description:"Exit with WARNING status if less than N units or N% of disk are free"`
	Critical      *string `short:"c" long:"critical" value-name:"N, N%" description:"Exit with CRITICAL status if less than N units or N% of disk are free"`
	InodeWarning  *string `short:"W" long:"iwarning" value-name:"N%" description:"Exit with WARNING status if less than PERCENT of inode space is free"`
	InodeCritical *string `short:"K" long:"icritical" value-name:"N%" description:"Exit with CRITICAL status if less than PERCENT of inode space is free"`
	Path          *string `short:"p" long:"path" value-name:"PATH" description:"Mount point or block device as emitted by the mount(8) command"`
	Exclude       *string `short:"x" long:"exclude-device" value-name:"EXCLUDE PATH" description:"Ignore device (only works if -p unspecified)"`
	All           bool    `short:"A" long:"all" description:"Explicitly select all paths."`
	Units         *string `short:"u" long:"units" value-name:"STRING" description:"Choose bytes, kB, MB, GB, TB (default: MB)"`
}

const (
	b  = float64(1)
	kb = float64(1024) * b
	mb = float64(1024) * kb
	gb = float64(1024) * mb
	tb = float64(1024) * gb
)

type unit struct {
	Name string
	Size float64
}

func checkStatus(current checkers.Status, threshold string, units float64, disk *gpud.UsageStat, chkInode bool, status checkers.Status) (checkers.Status, error) {
	if strings.HasSuffix(threshold, "%") {
		v, err := strconv.ParseFloat(strings.TrimRight(threshold, "%"), 64)
		if err != nil {
			return checkers.UNKNOWN, err
		}

		freePct := float64(100) - disk.UsedPercent
		inodesFreePct := float64(100) - disk.InodesUsedPercent

		if chkInode && v > inodesFreePct {
			current = status
		}

		if !chkInode && (v > freePct || v > inodesFreePct) {
			current = status
		}
	} else {
		if chkInode {
			return checkers.UNKNOWN, errors.New("-W, -K value should be N%")
		}

		v, err := strconv.ParseFloat(threshold, 64)
		if err != nil {
			return checkers.UNKNOWN, err
		}

		free := float64(disk.Free) / units
		if v > free {
			current = status
		}
	}

	return current, nil
}

func genMessage(disk *gpud.UsageStat, u unit) string {
	all := float64(disk.Total) / u.Size
	used := float64(disk.Used) / u.Size
	free := float64(disk.Free) / u.Size
	freePct := float64(100) - disk.UsedPercent
	inodesFreePct := float64(100) - disk.InodesUsedPercent

	return fmt.Sprintf("Path: %v, All: %.2f %v, Used: %.2f %v, Free: %.2f %v, Free percentage: %.2f (inodes: %.2f)", disk.Path, all, u.Name, used, u.Name, free, u.Name, freePct, inodesFreePct)
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "Disk"
	ckr.Exit()
}

func run(args []string) *checkers.Checker {
	_, err := flags.ParseArgs(&opts, args)
	if err != nil {
		os.Exit(1)
	}

	partitions, err := listPartitions()
	if err != nil {
		return checkers.Unknown(fmt.Sprintf("Faild to fetch partitions: %s", err))
	}

	if !opts.All {
		if opts.Path != nil {
			if opts.Exclude != nil {
				return checkers.Unknown(fmt.Sprintf("Invalid arguments: %s", errors.New("-x does not work with -p")))
			}

			exist := false
			for _, partition := range partitions {
				if *opts.Path == partition.Mountpoint {
					partitions = []gpud.PartitionStat{partition}
					exist = true
					break
				}
			}

			if !exist {
				return checkers.Unknown(fmt.Sprintf("Faild to fetch mountpoint: %s", errors.New("Invalid argument flag '-p, --path'")))
			}
		}

		if opts.Path == nil && opts.Exclude != nil {
			var tmp []gpud.PartitionStat
			for _, partition := range partitions {
				if *opts.Exclude != partition.Mountpoint {
					tmp = append(tmp, partition)
				}
			}
			partitions = tmp
		}
	}

	var disks []*gpud.UsageStat

	for _, partition := range partitions {
		disk, err := gpud.Usage(partition.Mountpoint)
		if err != nil {
			return checkers.Unknown(fmt.Sprintf("Faild to fetch disk usage: %s", err))
		}

		disks = append(disks, disk)
	}

	u := unit{"MB", mb}
	if opts.Units != nil {
		us := strings.ToLower(*opts.Units)
		if us == "bytes" {
			u = unit{us, b}
		} else if us == "kb" {
			u = unit{us, mb}
		} else if us == "gb" {
			u = unit{us, gb}
		} else if us == "tb" {
			u = unit{us, tb}
		} else {
			return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", errors.New("Invalid argument flag '-u, --units'")))
		}
	}

	checkSt := checkers.OK
	if opts.InodeCritical != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.InodeCritical, u.Size, disk, true, checkers.CRITICAL)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.CRITICAL {
				break
			}
		}
	}

	if checkSt != checkers.CRITICAL && opts.Critical != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.Critical, u.Size, disk, false, checkers.CRITICAL)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.CRITICAL {
				break
			}
		}
	}

	if checkSt != checkers.CRITICAL && opts.InodeWarning != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.InodeWarning, u.Size, disk, true, checkers.WARNING)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.WARNING {
				break
			}
		}
	}

	if checkSt == checkers.OK && opts.Warning != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.Warning, u.Size, disk, false, checkers.WARNING)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.WARNING {
				break
			}
		}
	}

	var msgs []string
	for _, disk := range disks {
		msg := genMessage(disk, u)
		msgs = append(msgs, msg)
	}
	msgss := strings.Join(msgs, ";\n")

	return checkers.NewChecker(checkSt, msgss)
}

// ref: mountlist.c in gnulib
// https://github.com/coreutils/gnulib/blob/a742bdb3/lib/mountlist.c#L168
func listPartitions() ([]gpud.PartitionStat, error) {
	allPartitions, err := gpud.Partitions(true)
	if err != nil {
		return nil, err
	}
	partitions := make([]gpud.PartitionStat, 0, len(allPartitions))
	for _, p := range allPartitions {
		switch p.Fstype {
		case "autofs",
			"proc",
			"subfs",
			"debugfs",
			"devpts",
			"fusectl",
			"mqueue",
			"rpc_pipefs",
			"sysfs",
			"devfs",
			"kernfs",
			"ignore":
			continue
		case "none":
			if !strings.Contains(p.Opts, "bind") {
				partitions = append(partitions, p)
			}
		default:
			partitions = append(partitions, p)
		}
	}
	return partitions, nil
}
<commit_msg>Add -N, -X opptions<commit_after>package checkdisk

import (
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"

	"github.com/jessevdk/go-flags"
	"github.com/mackerelio/checkers"
	gpud "github.com/shirou/gopsutil/disk"
)

var opts struct {
	Warning       *string `short:"w" long:"warning" value-name:"N, N%" description:"Exit with WARNING status if less than N units or N% of disk are free"`
	Critical      *string `short:"c" long:"critical" value-name:"N, N%" description:"Exit with CRITICAL status if less than N units or N% of disk are free"`
	InodeWarning  *string `short:"W" long:"iwarning" value-name:"N%" description:"Exit with WARNING status if less than PERCENT of inode space is free"`
	InodeCritical *string `short:"K" long:"icritical" value-name:"N%" description:"Exit with CRITICAL status if less than PERCENT of inode space is free"`
	Path          *string `short:"p" long:"path" value-name:"PATH" description:"Mount point or block device as emitted by the mount(8) command"`
	Exclude       *string `short:"x" long:"exclude-device" value-name:"EXCLUDE PATH" description:"Ignore device (only works if -p unspecified)"`
	All           bool    `short:"A" long:"all" description:"Explicitly select all paths."`
	ExcludeType   *string `short:"X" long:"exclude-type" value-name:"TYPE" description:"Ignore all filesystems of indicated type"`
	IncludeType   *string `short:"N" long:"include-type" value-name:"TYPE" description:"Check only filesystems of indicated type"`
	Units         *string `short:"u" long:"units" value-name:"STRING" description:"Choose bytes, kB, MB, GB, TB (default: MB)"`
}

const (
	b  = float64(1)
	kb = float64(1024) * b
	mb = float64(1024) * kb
	gb = float64(1024) * mb
	tb = float64(1024) * gb
)

type unit struct {
	Name string
	Size float64
}

func checkStatus(current checkers.Status, threshold string, units float64, disk *gpud.UsageStat, chkInode bool, status checkers.Status) (checkers.Status, error) {
	if strings.HasSuffix(threshold, "%") {
		v, err := strconv.ParseFloat(strings.TrimRight(threshold, "%"), 64)
		if err != nil {
			return checkers.UNKNOWN, err
		}

		freePct := float64(100) - disk.UsedPercent
		inodesFreePct := float64(100) - disk.InodesUsedPercent

		if chkInode && v > inodesFreePct {
			current = status
		}

		if !chkInode && (v > freePct || v > inodesFreePct) {
			current = status
		}
	} else {
		if chkInode {
			return checkers.UNKNOWN, errors.New("-W, -K value should be N%")
		}

		v, err := strconv.ParseFloat(threshold, 64)
		if err != nil {
			return checkers.UNKNOWN, err
		}

		free := float64(disk.Free) / units
		if v > free {
			current = status
		}
	}

	return current, nil
}

func genMessage(disk *gpud.UsageStat, u unit) string {
	all := float64(disk.Total) / u.Size
	used := float64(disk.Used) / u.Size
	free := float64(disk.Free) / u.Size
	freePct := float64(100) - disk.UsedPercent
	inodesFreePct := float64(100) - disk.InodesUsedPercent

	return fmt.Sprintf("Path: %v, All: %.2f %v, Used: %.2f %v, Free: %.2f %v, Free percentage: %.2f (inodes: %.2f)", disk.Path, all, u.Name, used, u.Name, free, u.Name, freePct, inodesFreePct)
}

// Do the plugin
func Do() {
	ckr := run(os.Args[1:])
	ckr.Name = "Disk"
	ckr.Exit()
}

func run(args []string) *checkers.Checker {
	_, err := flags.ParseArgs(&opts, args)
	if err != nil {
		os.Exit(1)
	}

	partitions, err := listPartitions(opts.IncludeType, opts.ExcludeType)
	if err != nil {
		return checkers.Unknown(fmt.Sprintf("Faild to fetch partitions: %s", err))
	}

	if !opts.All {
		if opts.Path != nil {
			if opts.Exclude != nil {
				return checkers.Unknown(fmt.Sprintf("Invalid arguments: %s", errors.New("-x does not work with -p")))
			}

			exist := false
			for _, partition := range partitions {
				if *opts.Path == partition.Mountpoint {
					partitions = []gpud.PartitionStat{partition}
					exist = true
					break
				}
			}

			if !exist {
				return checkers.Unknown(fmt.Sprintf("Faild to fetch mountpoint: %s", errors.New("Invalid argument flag '-p, --path'")))
			}
		}

		if opts.Path == nil && opts.Exclude != nil {
			var tmp []gpud.PartitionStat
			for _, partition := range partitions {
				if *opts.Exclude != partition.Mountpoint {
					tmp = append(tmp, partition)
				}
			}
			partitions = tmp
		}
	}

	var disks []*gpud.UsageStat

	for _, partition := range partitions {
		disk, err := gpud.Usage(partition.Mountpoint)
		if err != nil {
			return checkers.Unknown(fmt.Sprintf("Faild to fetch disk usage: %s", err))
		}

		disks = append(disks, disk)
	}

	u := unit{"MB", mb}
	if opts.Units != nil {
		us := strings.ToLower(*opts.Units)
		if us == "bytes" {
			u = unit{us, b}
		} else if us == "kb" {
			u = unit{us, mb}
		} else if us == "gb" {
			u = unit{us, gb}
		} else if us == "tb" {
			u = unit{us, tb}
		} else {
			return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", errors.New("Invalid argument flag '-u, --units'")))
		}
	}

	checkSt := checkers.OK
	if opts.InodeCritical != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.InodeCritical, u.Size, disk, true, checkers.CRITICAL)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.CRITICAL {
				break
			}
		}
	}

	if checkSt != checkers.CRITICAL && opts.Critical != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.Critical, u.Size, disk, false, checkers.CRITICAL)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.CRITICAL {
				break
			}
		}
	}

	if checkSt != checkers.CRITICAL && opts.InodeWarning != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.InodeWarning, u.Size, disk, true, checkers.WARNING)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.WARNING {
				break
			}
		}
	}

	if checkSt == checkers.OK && opts.Warning != nil {
		for _, disk := range disks {
			checkSt, err = checkStatus(checkSt, *opts.Warning, u.Size, disk, false, checkers.WARNING)
			if err != nil {
				return checkers.Unknown(fmt.Sprintf("Faild to check disk status: %s", err))
			}

			if checkSt == checkers.WARNING {
				break
			}
		}
	}

	var msgs []string
	for _, disk := range disks {
		msg := genMessage(disk, u)
		msgs = append(msgs, msg)
	}
	msgss := strings.Join(msgs, ";\n")

	return checkers.NewChecker(checkSt, msgss)
}

// ref: mountlist.c in gnulib
// https://github.com/coreutils/gnulib/blob/a742bdb3/lib/mountlist.c#L168
func listPartitions(includeType *string, excludeType *string) ([]gpud.PartitionStat, error) {
	allPartitions, err := gpud.Partitions(true)
	if err != nil {
		return nil, err
	}
	partitions := make([]gpud.PartitionStat, 0, len(allPartitions))
	for _, p := range allPartitions {
		if includeType != nil {
			if p.Fstype == *includeType {
				partitions = append(partitions, p)
				continue
			}
		}

		if excludeType != nil {
			if p.Fstype == *excludeType {
				continue
			}
		}

		switch p.Fstype {
		case "autofs",
			"proc",
			"subfs",
			"debugfs",
			"devpts",
			"fusectl",
			"mqueue",
			"rpc_pipefs",
			"sysfs",
			"devfs",
			"kernfs",
			"ignore":
			continue
		case "none":
			if !strings.Contains(p.Opts, "bind") {
				partitions = append(partitions, p)
			}
		default:
			partitions = append(partitions, p)
		}
	}
	return partitions, nil
}
<|endoftext|>


<commit_before>package grim

import (
	"sync"

	"golang.org/x/net/context"
)

// GrimReaper is a task runner which will wait until all the tasks
// are complete until comtinuing. Reapers can also kill tasks prematurely if
// they are listening to the given context.
type GrimReaper interface {
	// New creates a sub-reaper which is attached to the parent context. If
	// the parent context is killed, so are the children. However, Wait must
	// still be called for child reapers.
	New() GrimReaper

	// SpawnFunc starts a new goroutine for the given function. The task should
	// return as soon as possible after the context completes.
	SpawnFunc(TaskFunc)

	// Spawn starts a new goroutine for the given Task. The task should
	// return as soon as possible after the context completes.
	Spawn(Task)

	// Kill sends a message to all running tasks to stop. If child reapers
	// have also been created, they will be triggered to stop as well.
	Kill()

	// Wait will block until all tasks have completed. This will NOT block
	// until chlid reapers are finished. Each reaper must call wait
	// independently.
	Wait()
}

// Reaper returns an implementation of the GrimReaper interface.
func Reaper() GrimReaper {
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(context.Background())
	return &reaper{wg, ctx, cancel}
}

// ReaperWithContext creates a new GrimReaper implementation and uses the given
// context as the parent context.
func ReaperWithContext(c context.Context) GrimReaper {
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(c)
	return &reaper{wg, ctx, cancel}
}

// TaskFunc is a killable function which runs in a separate go routine. In order to fulfill the contract the function MUST listen to the context and exit if it fires.
type TaskFunc func(context.Context)

type Task interface {
	Execute(context.Context)
}

// reaper implements the GrimReaper interface.
type reaper struct {
	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc
}

// SpawnFunc runs a task.
func (r *reaper) SpawnFunc(t TaskFunc) {
	r.wg.Add(1)
	c, _ := context.WithCancel(r.ctx)
	go func(ctx context.Context) {
		defer r.wg.Done()
		t(ctx)
	}(c)
}

// Spawn runs a task.
func (r *reaper) Spawn(t Task) {
	r.wg.Add(1)
	c, _ := context.WithCancel(r.ctx)
	go func(ctx context.Context) {
		defer r.wg.Done()
		t.Execute(ctx)
	}(c)
}

// Wait blocks until all tasks have completed.
func (r *reaper) Wait() {
	r.wg.Wait()
}

// Kill cancels the context and waits for all tasks to exit.
func (r *reaper) Kill() {
	r.cancel()
	r.Wait()
}

// New creates a new reaper with the current context as the parent context of
// the new reaper.
func (r *reaper) New() GrimReaper {
	return ReaperWithContext(r.ctx)
}
<commit_msg>Fix typo<commit_after>package grim

import (
	"sync"

	"golang.org/x/net/context"
)

// GrimReaper is a task runner which will wait until all the tasks
// are complete until continuing. Reapers can also kill tasks prematurely if
// they are listening to the given context.
type GrimReaper interface {
	// New creates a sub-reaper which is attached to the parent context. If
	// the parent context is killed, so are the children. However, Wait must
	// still be called for child reapers.
	New() GrimReaper

	// SpawnFunc starts a new goroutine for the given function. The task should
	// return as soon as possible after the context completes.
	SpawnFunc(TaskFunc)

	// Spawn starts a new goroutine for the given Task. The task should
	// return as soon as possible after the context completes.
	Spawn(Task)

	// Kill sends a message to all running tasks to stop. If child reapers
	// have also been created, they will be triggered to stop as well.
	Kill()

	// Wait will block until all tasks have completed. This will NOT block
	// until chlid reapers are finished. Each reaper must call wait
	// independently.
	Wait()
}

// Reaper returns an implementation of the GrimReaper interface.
func Reaper() GrimReaper {
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(context.Background())
	return &reaper{wg, ctx, cancel}
}

// ReaperWithContext creates a new GrimReaper implementation and uses the given
// context as the parent context.
func ReaperWithContext(c context.Context) GrimReaper {
	var wg sync.WaitGroup
	ctx, cancel := context.WithCancel(c)
	return &reaper{wg, ctx, cancel}
}

// TaskFunc is a killable function which runs in a separate go routine. In order to fulfill the contract the function MUST listen to the context and exit if it fires.
type TaskFunc func(context.Context)

type Task interface {
	Execute(context.Context)
}

// reaper implements the GrimReaper interface.
type reaper struct {
	wg     sync.WaitGroup
	ctx    context.Context
	cancel context.CancelFunc
}

// SpawnFunc runs a task.
func (r *reaper) SpawnFunc(t TaskFunc) {
	r.wg.Add(1)
	c, _ := context.WithCancel(r.ctx)
	go func(ctx context.Context) {
		defer r.wg.Done()
		t(ctx)
	}(c)
}

// Spawn runs a task.
func (r *reaper) Spawn(t Task) {
	r.wg.Add(1)
	c, _ := context.WithCancel(r.ctx)
	go func(ctx context.Context) {
		defer r.wg.Done()
		t.Execute(ctx)
	}(c)
}

// Wait blocks until all tasks have completed.
func (r *reaper) Wait() {
	r.wg.Wait()
}

// Kill cancels the context and waits for all tasks to exit.
func (r *reaper) Kill() {
	r.cancel()
	r.Wait()
}

// New creates a new reaper with the current context as the parent context of
// the new reaper.
func (r *reaper) New() GrimReaper {
	return ReaperWithContext(r.ctx)
}
<|endoftext|>


<commit_before>package main

import (
	"log"
	"time"

	"github.com/99designs/ecs-cli"
	"github.com/99designs/ecs-cli/cli"
	"github.com/99designs/ecs-cli/templates"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/aws/aws-sdk-go/service/ecs"
)

type EcsClusterParameters struct {
	KeyName       string
	InstanceType  string
	InstanceCount int
}

type CreateClusterCommandInput struct {
	ClusterName string
	Parameters  EcsClusterParameters
}

func CreateClusterCommand(ui *cli.Ui, input CreateClusterCommandInput) {
	ui.Printf("Creating cluster %s", input.ClusterName)

	_, err := ecsSvc.CreateCluster(&ecs.CreateClusterInput{
		ClusterName: aws.String(input.ClusterName),
	})

	timer := time.Now()
	stackName := input.ClusterName + "-ecs-" + time.Now().Format("20060102-150405")
	ui.Printf("Creating cloudformation stack %s", stackName)

	err = ecscli.CreateStack(cfnSvc, stackName, templates.EcsStack(), map[string]string{
		"KeyName":         input.Parameters.KeyName,
		"ECSCluster":      input.ClusterName,
		"InstanceType":    input.Parameters.InstanceType,
		"DesiredCapacity": string(input.Parameters.InstanceCount),
	})
	if err != nil {
		ui.Fatal(err)
	}

	err = ecscli.PollStackEvents(cfnSvc, stackName, func(event *cloudformation.StackEvent) {
		ui.Printf("%s\n", ecscli.FormatStackEvent(event))
	})
	if err != nil {
		ui.Fatal(err)
	}

	log.Printf("Cluster %s created in %s", input.ClusterName, time.Now().Sub(timer).String())
}
<commit_msg>Fix string encoding of DesiredCapacity<commit_after>package main

import (
	"log"
	"strconv"
	"time"

	"github.com/99designs/ecs-cli"
	"github.com/99designs/ecs-cli/cli"
	"github.com/99designs/ecs-cli/templates"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/cloudformation"
	"github.com/aws/aws-sdk-go/service/ecs"
)

type EcsClusterParameters struct {
	KeyName       string
	InstanceType  string
	InstanceCount int
}

type CreateClusterCommandInput struct {
	ClusterName string
	Parameters  EcsClusterParameters
}

func CreateClusterCommand(ui *cli.Ui, input CreateClusterCommandInput) {
	ui.Printf("Creating cluster %s", input.ClusterName)

	_, err := ecsSvc.CreateCluster(&ecs.CreateClusterInput{
		ClusterName: aws.String(input.ClusterName),
	})

	timer := time.Now()
	stackName := input.ClusterName + "-ecs-" + time.Now().Format("20060102-150405")
	ui.Printf("Creating cloudformation stack %s", stackName)

	err = ecscli.CreateStack(cfnSvc, stackName, templates.EcsStack(), map[string]string{
		"KeyName":         input.Parameters.KeyName,
		"ECSCluster":      input.ClusterName,
		"InstanceType":    input.Parameters.InstanceType,
		"DesiredCapacity": strconv.Itoa(input.Parameters.InstanceCount),
	})
	if err != nil {
		ui.Fatal(err)
	}

	err = ecscli.PollStackEvents(cfnSvc, stackName, func(event *cloudformation.StackEvent) {
		ui.Printf("%s\n", ecscli.FormatStackEvent(event))
	})
	if err != nil {
		ui.Fatal(err)
	}

	log.Printf("Cluster %s created in %s", input.ClusterName, time.Now().Sub(timer).String())
}
<|endoftext|>


<commit_before>// Package cnveval provides a way to evalute CNVs based on a truth-set.
package cnveval

import (
	"fmt"
	"math"
	"sort"
)

// CNV indicates the region, sample and copy-number of a region.
type CNV struct {
	Chrom   string
	Start   int
	End     int
	Sample  string
	CN      int
	counted bool
}

// Truth indicates the region, samples and copy-number for a truth-set.
type Truth struct {
	Chrom   string
	Start   int
	End     int
	Samples []string
	CN      int
	used    map[string]bool
}

func (t Truth) String() string {
	return fmt.Sprintf("%s:%d-%d[%d]", t.Chrom, t.Start, t.End, t.CN)
}

func (c CNV) String() string {
	return fmt.Sprintf("%s:%d-%d[%d] in %s", c.Chrom, c.Start, c.End, c.CN, c.Sample)
}

func (t Truth) equals(o Truth) bool {
	return t.Start == o.Start && t.End == o.End
}

// SC is a size-class
type SC int

var (
	// Any is used to avoid any size-class designation when reporting results.
	Any    SC
	Small  SC = 20000
	Medium SC = 100000
	Large  SC = math.MaxInt64
)

func (s SC) Order() int {
	switch s {
	case Small:
		return 0
	case Medium:
		return 1
	case Large:
		return 2
	case Any:
		return 3
	default:
		panic(fmt.Sprintf("unknown size class: %d", s))
	}
}

func (s SC) String() string {
	switch s {
	case Any:
		return "all"
	case Small:
		return fmt.Sprintf("0-%d", Small)
	case Medium:
		return fmt.Sprintf("%d-%d", Small, Medium)
	case Large:
		return fmt.Sprintf(">=%d", Medium)
	default:
		panic(fmt.Sprintf("unknown size class: %d", s))
	}
}

// TS indicates a sample and SizeClass
type TS struct {
	SizeClass SC
	Sample    string
}

type Stat struct {
	FP, FN, TP, TN int
}

func (s Stat) String() string {
	return fmt.Sprintf("precision: %.4f (%-04d / (%-04d + %-04d)) recall: %.4f (%-04d / (%-04d + %-04d))", s.Precision(), s.TP, s.TP, s.FP,
		s.Recall(), s.TP, s.TP, s.FN)
}

func (s Stat) Recall() float64 {
	return float64(s.TP) / float64(s.TP+s.FN)
}
func (s Stat) Precision() float64 {
	return float64(s.TP) / float64(s.TP+s.FP)
}

type CohortStats map[TS]*Stat

func (c CohortStats) TP(class SC) int {
	var TP int
	for cls, st := range c {
		if class != Any && cls.SizeClass != class {
			continue
		}
		TP += st.TP
	}
	return TP
}

func (c CohortStats) Tabulate() [4]Stat {
	var stats [4]Stat
	for cls, st := range c {
		stats[cls.SizeClass.Order()].TP += st.TP
		stats[cls.SizeClass.Order()].FP += st.FP
		stats[cls.SizeClass.Order()].TN += st.TN
		stats[cls.SizeClass.Order()].FN += st.FN
	}
	for _, cl := range []SC{Small, Medium, Large} {
		stats[Any.Order()].TP += stats[cl.Order()].TP
		stats[Any.Order()].FP += stats[cl.Order()].FP
		stats[Any.Order()].TN += stats[cl.Order()].TN
		stats[Any.Order()].FN += stats[cl.Order()].FN
	}
	return stats
}

func (c CohortStats) Precision(class SC) float64 {
	var TP, FP float64
	for cls, st := range c {
		if class != Any && cls.SizeClass != class {
			continue
		}
		TP += float64(st.TP)
		FP += float64(st.FP)
	}
	//log.Println("TP, FP:", TP, FP)
	return TP / (TP + FP)
}

func (c CohortStats) Recall(class SC) float64 {
	var TP, FN, FP, TN float64
	for cls, st := range c {
		if class != Any && cls.SizeClass != class {
			continue
		}
		TP += float64(st.TP)
		FN += float64(st.FN)
		FP += float64(st.FP)
		TN += float64(st.TN)
	}
	//log.Println("TP, FN, FP, TN:", TP, FN, FP, TN)
	return TP / (TP + FN)
}

func Evaluate(cnvs []CNV, truths []Truth, po float64) CohortStats {
	stat := make(map[TS]*Stat, 200)
	allSamples := make(map[string]bool)
	for _, t := range truths {
		for _, s := range t.Samples {
			allSamples[s] = true
		}
	}
	for _, c := range cnvs {
		allSamples[c.Sample] = true
	}
	sampleList := fromMap(allSamples)
	truthBySample := make(map[string][]Truth)
	truthWithoutSample := make(map[string][]Truth)
	for _, t := range truths {
		for _, s := range t.Samples {
			truthBySample[s] = append(truthBySample[s], t)
		}
		for _, s := range sampleList {
			if notin(s, t.Samples) {
				truthWithoutSample[s] = append(truthWithoutSample[s], t)
			}
		}
	}

	cnvBySample := make(map[string][]CNV)
	for _, c := range cnvs {
		cnvBySample[c.Sample] = append(cnvBySample[c.Sample], c)
	}

	for _, sample := range sampleList {
		truths, _ := truthBySample[sample]
		sort.Slice(truths, func(i, j int) bool {
			return truths[i].Chrom < truths[j].Chrom || (truths[i].Chrom == truths[j].Chrom && truths[i].Start < truths[j].Start)
		})
		cnvs, _ := cnvBySample[sample]
		sort.Slice(cnvs, func(i, j int) bool {
			return cnvs[i].Chrom < cnvs[j].Chrom || (cnvs[i].Chrom == cnvs[j].Chrom && cnvs[i].Start < cnvs[j].Start)
		})

		updatePositive(stat, truths, cnvs, po)
		others, _ := truthWithoutSample[sample]
		sort.Slice(others, func(i, j int) bool {
			return others[i].Chrom < others[j].Chrom || (others[i].Chrom == others[j].Chrom && others[i].Start < others[j].Start)
		})
		updateFP(stat, others, cnvs, truths, po)
	}

	return stat
}

func notin(a string, bs []string) bool {
	for _, b := range bs {
		if a == b {
			return false
		}
	}
	return true
}

func fromMap(samples map[string]bool) []string {
	m := make([]string, 0, len(samples))
	for k := range samples {
		m = append(m, k)
	}
	return m
}

func updateFP(stat map[TS]*Stat, others []Truth, cnvs []CNV, truths []Truth, po float64) {
	if len(cnvs) == 0 || len(others) == 0 {
		return
	}

	var i int
	// others should not include CNVs from this sample.
	for _, o := range others {
		ts := TS{Sample: cnvs[0].Sample, SizeClass: sizeClass(o)}
		val := stat[ts]
		if val == nil {
			val = &Stat{}
			stat[ts] = val
		}
		// don't need to reset i because others is sorted.
		for ; i < len(cnvs) && (cnvs[i].Chrom < o.Chrom || (cnvs[i].Chrom == o.Chrom && cnvs[i].End < o.Start)); i++ {
		}
		if i > 0 {
			i--
		}
		tpfound := false
		fpfound := false
		found := false
		// TODO: flip this cnvs loop with the others loop above because we need to have a value for every cnv.
		for k, cnv := range cnvs[i:] {
			if cnv.Chrom > o.Chrom || (o.Chrom == cnv.Chrom && cnv.Start > o.End) {
				break
			}
			if poverlap(cnv, o) >= po && sameCN(cnv.CN, o.CN) {
				fpfound = true
				// we have a putative FP, but need to check if there's something else in the truth something
				// that also fits this criteria to avoid incorrectly calling an FP.
				for _, t := range truths {
					if t.Chrom != cnv.Chrom {
						continue
					}
					if poverlap(cnv, t) >= po && sameCN(cnv.CN, t.CN) {
						tpfound = true
						break
					}

				}
			}
			// if here, then cnv didn't overlap something in the truth set.
			if fpfound && !tpfound {
				val.FP++
				found = true
				cnvs[i+k].counted = true
			}
		}
		if !(found || tpfound) {
			val.TN++
		}
	}
}

// given a set of truths and cnvs from a sample, update the FPs.
// truths and cnvs are sorted by chrom, then by start
func updatePositive(stat map[TS]*Stat, truths []Truth, cnvs []CNV, po float64) {
	if len(cnvs) == 0 {
		return
	}
	var i int
	for _, t := range truths {
		ts := TS{Sample: cnvs[0].Sample, SizeClass: sizeClass(t)}
		found := false
		val := stat[ts]
		if val == nil {
			val = &Stat{}
			stat[ts] = val
		}
		// since truths are sorted as well, we don't need to reset i.
		for ; i < len(cnvs) && (cnvs[i].Chrom < t.Chrom || (cnvs[i].Chrom == t.Chrom && cnvs[i].End < t.Start)); i++ {
		}
		if i > 0 {
			i--
		}
		for k, cnv := range cnvs[i:] {
			if cnv.Chrom > t.Chrom || (t.Chrom == cnv.Chrom && cnv.Start > t.End) {
				break
			}
			if poverlap(cnv, t) >= po && sameCN(cnv.CN, t.CN) {
				// used map make sure we don't double count cnvs from the sameCN
				// sample that are subsets of the full truth interval.
				if _, ok := t.used[cnv.Sample]; !ok {
					val.TP++
					cnvs[i+k].counted = true
					found = true
					if t.used == nil {
						t.used = make(map[string]bool)
					}
					t.used[cnv.Sample] = true
				}
			}
		}
		if !found {
			val.FN++
		}
	}
	for _, cnv := range cnvs {
		if !cnv.counted {
			ts := TS{Sample: cnv.Sample, SizeClass: sizeClass(Truth{Start: cnv.Start, End: cnv.End})}
			val := stat[ts]
			if val == nil {
				val = &Stat{}
				stat[ts] = val
			}
			val.FP++
		}
	}
}

func sizeClass(t Truth) SC {
	l := t.End - t.Start
	if l < int(Small) {
		return Small
	}
	if l < int(Medium) {
		return Medium
	}
	return Large
}

func sameCN(a, b int) bool {
	if a > 2 {
		a = 3
	}
	if b > 2 {
		b = 3
	}
	return math.Abs(float64(a)-float64(b)) < 1
}

func imin(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func imax(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func poverlap(a CNV, b Truth) float64 {
	if a.Chrom != b.Chrom {
		return 0
	}
	total := math.Min(float64(a.End-a.Start), float64(b.End-b.Start))
	ovl := imin(a.End, b.End) - imax(a.Start, b.Start)
	if ovl < 0 {
		return 0
	}
	return float64(ovl) / total
}
<commit_msg>simplify comparison<commit_after>// Package cnveval provides a way to evalute CNVs based on a truth-set.
package cnveval

import (
	"fmt"
	"math"
	"sort"
)

// CNV indicates the region, sample and copy-number of a region.
type CNV struct {
	Chrom   string
	Start   int
	End     int
	Sample  string
	CN      int
	counted bool
}

// Truth indicates the region, samples and copy-number for a truth-set.
type Truth struct {
	Chrom   string
	Start   int
	End     int
	Samples []string
	CN      int
	used    map[string]bool
}

func (t Truth) String() string {
	return fmt.Sprintf("%s:%d-%d[%d]", t.Chrom, t.Start, t.End, t.CN)
}

func (c CNV) String() string {
	return fmt.Sprintf("%s:%d-%d[%d] in %s", c.Chrom, c.Start, c.End, c.CN, c.Sample)
}

func (t Truth) equals(o Truth) bool {
	return t.Start == o.Start && t.End == o.End
}

// SC is a size-class
type SC int

var (
	// Any is used to avoid any size-class designation when reporting results.
	Any    SC
	Small  SC = 20000
	Medium SC = 100000
	Large  SC = math.MaxInt64
)

func (s SC) Order() int {
	switch s {
	case Small:
		return 0
	case Medium:
		return 1
	case Large:
		return 2
	case Any:
		return 3
	default:
		panic(fmt.Sprintf("unknown size class: %d", s))
	}
}

func (s SC) String() string {
	switch s {
	case Any:
		return "all"
	case Small:
		return fmt.Sprintf("0-%d", Small)
	case Medium:
		return fmt.Sprintf("%d-%d", Small, Medium)
	case Large:
		return fmt.Sprintf(">=%d", Medium)
	default:
		panic(fmt.Sprintf("unknown size class: %d", s))
	}
}

// TS indicates a sample and SizeClass
type TS struct {
	SizeClass SC
	Sample    string
}

type Stat struct {
	FP, FN, TP, TN int
}

func (s Stat) String() string {
	return fmt.Sprintf("precision: %.4f (%-04d / (%-04d + %-04d)) recall: %.4f (%-04d / (%-04d + %-04d))", s.Precision(), s.TP, s.TP, s.FP,
		s.Recall(), s.TP, s.TP, s.FN)
}

func (s Stat) Recall() float64 {
	return float64(s.TP) / float64(s.TP+s.FN)
}
func (s Stat) Precision() float64 {
	return float64(s.TP) / float64(s.TP+s.FP)
}

type CohortStats map[TS]*Stat

func (c CohortStats) TP(class SC) int {
	var TP int
	for cls, st := range c {
		if class != Any && cls.SizeClass != class {
			continue
		}
		TP += st.TP
	}
	return TP
}

func (c CohortStats) Tabulate() [4]Stat {
	var stats [4]Stat
	for cls, st := range c {
		stats[cls.SizeClass.Order()].TP += st.TP
		stats[cls.SizeClass.Order()].FP += st.FP
		stats[cls.SizeClass.Order()].TN += st.TN
		stats[cls.SizeClass.Order()].FN += st.FN
	}
	for _, cl := range []SC{Small, Medium, Large} {
		stats[Any.Order()].TP += stats[cl.Order()].TP
		stats[Any.Order()].FP += stats[cl.Order()].FP
		stats[Any.Order()].TN += stats[cl.Order()].TN
		stats[Any.Order()].FN += stats[cl.Order()].FN
	}
	return stats
}

func (c CohortStats) Precision(class SC) float64 {
	var TP, FP float64
	for cls, st := range c {
		if class != Any && cls.SizeClass != class {
			continue
		}
		TP += float64(st.TP)
		FP += float64(st.FP)
	}
	//log.Println("TP, FP:", TP, FP)
	return TP / (TP + FP)
}

func (c CohortStats) Recall(class SC) float64 {
	var TP, FN, FP, TN float64
	for cls, st := range c {
		if class != Any && cls.SizeClass != class {
			continue
		}
		TP += float64(st.TP)
		FN += float64(st.FN)
		FP += float64(st.FP)
		TN += float64(st.TN)
	}
	//log.Println("TP, FN, FP, TN:", TP, FN, FP, TN)
	return TP / (TP + FN)
}

func Evaluate(cnvs []CNV, truths []Truth, po float64) CohortStats {
	stat := make(map[TS]*Stat, 200)
	allSamples := make(map[string]bool)
	for _, t := range truths {
		for _, s := range t.Samples {
			allSamples[s] = true
		}
	}
	for _, c := range cnvs {
		allSamples[c.Sample] = true
	}
	sampleList := fromMap(allSamples)
	truthBySample := make(map[string][]Truth)
	truthWithoutSample := make(map[string][]Truth)
	for _, t := range truths {
		for _, s := range t.Samples {
			truthBySample[s] = append(truthBySample[s], t)
		}
		for _, s := range sampleList {
			if notin(s, t.Samples) {
				truthWithoutSample[s] = append(truthWithoutSample[s], t)
			}
		}
	}

	cnvBySample := make(map[string][]CNV)
	for _, c := range cnvs {
		cnvBySample[c.Sample] = append(cnvBySample[c.Sample], c)
	}

	for _, sample := range sampleList {
		truths, _ := truthBySample[sample]
		sort.Slice(truths, func(i, j int) bool {
			return truths[i].Chrom < truths[j].Chrom || (truths[i].Chrom == truths[j].Chrom && truths[i].Start < truths[j].Start)
		})
		cnvs, _ := cnvBySample[sample]
		sort.Slice(cnvs, func(i, j int) bool {
			return cnvs[i].Chrom < cnvs[j].Chrom || (cnvs[i].Chrom == cnvs[j].Chrom && cnvs[i].Start < cnvs[j].Start)
		})

		updatePositive(stat, truths, cnvs, po)
		others, _ := truthWithoutSample[sample]
		sort.Slice(others, func(i, j int) bool {
			return others[i].Chrom < others[j].Chrom || (others[i].Chrom == others[j].Chrom && others[i].Start < others[j].Start)
		})
		updateFP(stat, others, cnvs, truths, po)
	}

	return stat
}

func notin(a string, bs []string) bool {
	for _, b := range bs {
		if a == b {
			return false
		}
	}
	return true
}

func fromMap(samples map[string]bool) []string {
	m := make([]string, 0, len(samples))
	for k := range samples {
		m = append(m, k)
	}
	return m
}

func updateFP(stat map[TS]*Stat, others []Truth, cnvs []CNV, truths []Truth, po float64) {
	if len(cnvs) == 0 || len(others) == 0 {
		return
	}

	var i int
	// others should not include CNVs from this sample.
	for _, o := range others {
		ts := TS{Sample: cnvs[0].Sample, SizeClass: sizeClass(o)}
		val := stat[ts]
		if val == nil {
			val = &Stat{}
			stat[ts] = val
		}
		// don't need to reset i because others is sorted.
		for ; i < len(cnvs) && (cnvs[i].Chrom < o.Chrom || (cnvs[i].Chrom == o.Chrom && cnvs[i].End < o.Start)); i++ {
		}
		if i > 0 {
			i--
		}
		tpfound := false
		fpfound := false
		found := false
		// TODO: flip this cnvs loop with the others loop above because we need to have a value for every cnv.
		for k, cnv := range cnvs[i:] {
			if cnv.Chrom > o.Chrom || (o.Chrom == cnv.Chrom && cnv.Start > o.End) {
				break
			}
			if poverlap(cnv, o) >= po && sameCN(cnv.CN, o.CN) {
				fpfound = true
				// we have a putative FP, but need to check if there's something else in the truth something
				// that also fits this criteria to avoid incorrectly calling an FP.
				for _, t := range truths {
					if t.Chrom != cnv.Chrom {
						continue
					}
					if poverlap(cnv, t) >= po && sameCN(cnv.CN, t.CN) {
						tpfound = true
						break
					}

				}
			}
			// if here, then cnv didn't overlap something in the truth set.
			if fpfound && !tpfound {
				val.FP++
				found = true
				cnvs[i+k].counted = true
			}
		}
		if !(found || tpfound) {
			val.TN++
		}
	}
}

// given a set of truths and cnvs from a sample, update the FPs.
// truths and cnvs are sorted by chrom, then by start
func updatePositive(stat map[TS]*Stat, truths []Truth, cnvs []CNV, po float64) {
	if len(cnvs) == 0 {
		return
	}
	var i int
	for _, t := range truths {
		ts := TS{Sample: cnvs[0].Sample, SizeClass: sizeClass(t)}
		found := false
		val := stat[ts]
		if val == nil {
			val = &Stat{}
			stat[ts] = val
		}
		// since truths are sorted as well, we don't need to reset i.
		for ; i < len(cnvs) && (cnvs[i].Chrom < t.Chrom || (cnvs[i].Chrom == t.Chrom && cnvs[i].End < t.Start)); i++ {
		}
		if i > 0 {
			i--
		}
		for k, cnv := range cnvs[i:] {
			if cnv.Chrom > t.Chrom || (t.Chrom == cnv.Chrom && cnv.Start > t.End) {
				break
			}
			if poverlap(cnv, t) >= po && sameCN(cnv.CN, t.CN) {
				// used map make sure we don't double count cnvs from the sameCN
				// sample that are subsets of the full truth interval.
				if _, ok := t.used[cnv.Sample]; !ok {
					val.TP++
					cnvs[i+k].counted = true
					found = true
					if t.used == nil {
						t.used = make(map[string]bool)
					}
					t.used[cnv.Sample] = true
				}
			}
		}
		if !found {
			val.FN++
		}
	}
	for _, cnv := range cnvs {
		if !cnv.counted {
			ts := TS{Sample: cnv.Sample, SizeClass: sizeClass(Truth{Start: cnv.Start, End: cnv.End})}
			val := stat[ts]
			if val == nil {
				val = &Stat{}
				stat[ts] = val
			}
			val.FP++
		}
	}
}

func sizeClass(t Truth) SC {
	l := t.End - t.Start
	if l < int(Small) {
		return Small
	}
	if l < int(Medium) {
		return Medium
	}
	return Large
}

func sameCN(a, b int) bool {
	if a > 2 {
		a = 3
	}
	if b > 2 {
		b = 3
	}
	return a == b
}

func imin(a, b int) int {
	if a < b {
		return a
	}
	return b
}

func imax(a, b int) int {
	if a > b {
		return a
	}
	return b
}

func poverlap(a CNV, b Truth) float64 {
	if a.Chrom != b.Chrom {
		return 0
	}
	total := math.Min(float64(a.End-a.Start), float64(b.End-b.Start))
	ovl := imin(a.End, b.End) - imax(a.Start, b.Start)
	if ovl < 0 {
		return 0
	}
	return float64(ovl) / total
}
<|endoftext|>