<commit_before>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/opts"
	flag "github.com/docker/docker/pkg/mflag"

	dockerClient "github.com/fsouza/go-dockerclient"
)

var (
	SYSFS       string        = "/sys/fs/cgroup"
	PROCS       string        = "cgroup.procs"
	CGROUP_PROC string        = "/proc/%d/cgroup"
	INTERVAL    time.Duration = 1000
)

type Context struct {
	Args         []string
	Cgroups      []string
	AllCgroups   bool
	Logs         bool
	Notify       bool
	Name         string
	Env          bool
	Rm           bool
	Id           string
	NotifySocket string
	Cmd          *exec.Cmd
	Pid          int
	PidFile      string
	Client       *dockerClient.Client
}

func setupEnvironment(c *Context) {
	newArgs := []string{}
	if c.Notify && len(c.NotifySocket) > 0 {
		newArgs = append(newArgs, "-e", fmt.Sprintf("NOTIFY_SOCKET=%s", c.NotifySocket))
		newArgs = append(newArgs, "-v", fmt.Sprintf("%s:%s", c.NotifySocket, c.NotifySocket))
	} else {
		c.Notify = false
	}

	if c.Env {
		for _, val := range os.Environ() {
			if !strings.HasPrefix(val, "HOME=") && !strings.HasPrefix(val, "PATH=") {
				newArgs = append(newArgs, "-e", val)
			}
		}
	}

	if len(newArgs) > 0 {
		c.Args = append(newArgs, c.Args...)
	}
}

func parseContext(args []string) (*Context, error) {
	c := &Context{
		Logs:       true,
		AllCgroups: false,
	}

	flags := flag.NewFlagSet("systemd-docker", flag.ContinueOnError)

	flCgroups := opts.NewListOpts(nil)

	flags.StringVar(&c.PidFile, []string{"p", "-pid-file"}, "", "pipe file")
	flags.BoolVar(&c.Logs, []string{"l", "-logs"}, true, "pipe logs")
	flags.BoolVar(&c.Notify, []string{"n", "-notify"}, false, "setup systemd notify for container")
	flags.BoolVar(&c.Env, []string{"e", "-env"}, false, "inherit environment variable")
	flags.Var(&flCgroups, []string{"c", "-cgroups"}, "cgroups to take ownership of or 'all' for all cgroups available")

	err := flags.Parse(args)
	if err != nil {
		return nil, err
	}

	foundD := false
	var name string

	runArgs := flags.Args()
	if len(runArgs) == 0 || runArgs[0] != "run" {
		log.Println("Args:", runArgs)
		return nil, errors.New("run not found in arguments")
	}

	runArgs = runArgs[1:]
	newArgs := make([]string, 0, len(runArgs))

	for i, arg := range runArgs {
		/* This is tedious, but flag can't ignore unknown flags and I don't want to define them all */
		add := true

		switch {
		case arg == "-rm" || arg == "--rm":
			c.Rm = true
			add = false
		case arg == "-d" || arg == "-detach" || arg == "--detach":
			foundD = true
		case strings.HasPrefix(arg, "-name") || strings.HasPrefix(arg, "--name"):
			if strings.Contains(arg, "=") {
				name = strings.SplitN(arg, "=", 2)[1]
			} else if len(runArgs) > i+1 {
				name = runArgs[i+1]
			}
		}

		if add {
			newArgs = append(newArgs, arg)
		}
	}

	if !foundD {
		newArgs = append([]string{"-d"}, newArgs...)
	}

	c.Name = name
	c.NotifySocket = os.Getenv("NOTIFY_SOCKET")
	c.Args = newArgs
	c.Cgroups = flCgroups.GetAll()

	for _, val := range c.Cgroups {
		if val == "all" {
			c.Cgroups = nil
			c.AllCgroups = true
			break
		}
	}

	setupEnvironment(c)

	return c, nil
}

func lookupNamedContainer(c *Context) error {
	client, err := getClient(c)
	if err != nil {
		return err
	}

	container, err := client.InspectContainer(c.Name)
	if _, ok := err.(*dockerClient.NoSuchContainer); ok {
		return nil
	}
	if err != nil || container == nil {
		return err
	}

	if container.State.Running {
		c.Id = container.ID
		c.Pid = container.State.Pid
		return nil
	} else if c.Rm {
		return client.RemoveContainer(dockerClient.RemoveContainerOptions{
			ID:    container.ID,
			Force: true,
		})
	} else {
		client, err := getClient(c)
		err = client.StartContainer(container.ID, container.HostConfig)
		if err != nil {
			return err
		}

		container, err = client.InspectContainer(c.Name)
		if err != nil {
			return err
		}

		c.Id = container.ID
		c.Pid = container.State.Pid

		return nil
	}
}

func launchContainer(c *Context) error {
	args := append([]string{"run"}, c.Args...)
	c.Cmd = exec.Command("docker", args...)

	errorPipe, err := c.Cmd.StderrPipe()
	if err != nil {
		return err
	}

	outputPipe, err := c.Cmd.StdoutPipe()
	if err != nil {
		return err
	}

	err = c.Cmd.Start()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, errorPipe)

	bytes, err := ioutil.ReadAll(outputPipe)
	if err != nil {
		return err
	}

	c.Id = strings.TrimSpace(string(bytes))

	err = c.Cmd.Wait()
	if err != nil {
		return err
	}

	if !c.Cmd.ProcessState.Success() {
		return err
	}

	c.Pid, err = getContainerPid(c)

	return err
}

func runContainer(c *Context) error {
	if len(c.Name) > 0 {
		err := lookupNamedContainer(c)
		if err != nil {
			return err
		}

	}

	if len(c.Id) == 0 {
		err := launchContainer(c)
		if err != nil {
			return err
		}
	}

	if c.Pid == 0 {
		return errors.New("Failed to launch container, pid is 0")
	}

	return nil
}

func getClient(c *Context) (*dockerClient.Client, error) {
	if c.Client != nil {
		return c.Client, nil
	}

	endpoint := os.Getenv("DOCKER_HOST")
	if len(endpoint) == 0 {
		endpoint = "unix:///var/run/docker.sock"
	}

	return dockerClient.NewVersionedClient(endpoint, "1.11")
}

func getContainerPid(c *Context) (int, error) {
	client, err := getClient(c)
	if err != nil {
		return 0, err
	}

	container, err := client.InspectContainer(c.Id)
	if err != nil {
		return 0, err
	}

	if container == nil {
		return 0, errors.New(fmt.Sprintf("Failed to find container %s", c.Id))
	}

	if container.State.Pid <= 0 {
		return 0, errors.New(fmt.Sprintf("Pid is %d for container %s", container.State.Pid, c.Id))
	}

	return container.State.Pid, nil
}

func getCgroupsForPid(pid int) (map[string]string, error) {
	file, err := os.Open(fmt.Sprintf(CGROUP_PROC, pid))
	if err != nil {
		return nil, err
	}

	ret := map[string]string{}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.SplitN(scanner.Text(), ":", 3)
		if len(line) != 3 {
			continue
		}

		ret[line[1]] = line[2]
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func constructCgroupPath(cgroupName string, cgroupPath string) string {
	return path.Join(SYSFS, strings.TrimPrefix(cgroupName, "name="), cgroupPath, PROCS)
}

func getCgroupPids(cgroupName string, cgroupPath string) ([]string, error) {
	ret := []string{}

	file, err := os.Open(constructCgroupPath(cgroupName, cgroupPath))
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		ret = append(ret, strings.TrimSpace(scanner.Text()))
	}

	if err = scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func writePid(pid string, path string) error {
	return ioutil.WriteFile(path, []byte(pid), 0644)
}

func moveCgroups(c *Context) (bool, error) {
	moved := false
	currentCgroups, err := getCgroupsForPid(os.Getpid())
	if err != nil {
		return false, err
	}

	containerCgroups, err := getCgroupsForPid(c.Pid)
	if err != nil {
		return false, err
	}

	var ns []string

	if c.AllCgroups || c.Cgroups == nil || len(c.Cgroups) == 0 {
		ns = make([]string, 0, len(containerCgroups))
		for value, _ := range containerCgroups {
			ns = append(ns, value)
		}
	} else {
		ns = c.Cgroups
	}

	for _, nsName := range ns {
		currentPath, ok := currentCgroups[nsName]
		if !ok {
			continue
		}

		containerPath, ok := containerCgroups[nsName]
		if !ok {
			continue
		}

		if currentPath == containerPath || containerPath == "/" {
			continue
		}

		pids, err := getCgroupPids(nsName, containerPath)
		if err != nil {
			return false, err
		}

		for _, pid := range pids {
			pidInt, err := strconv.Atoi(pid)
			if err != nil {
				continue
			}

			if pidDied(pidInt) {
				continue
			}

			currentFullPath := constructCgroupPath(nsName, currentPath)
			log.Printf("Moving pid %s to %s\n", pid, currentFullPath)
			err = writePid(pid, currentFullPath)
			if err != nil {
				return false, err
			}

			moved = true
		}
	}

	return moved, nil
}

func pidDied(pid int) bool {
	_, err := os.Stat(fmt.Sprintf("/proc/%d", pid))
	return os.IsNotExist(err)
}

func notify(c *Context) error {
	if pidDied(c.Pid) {
		return errors.New("Container exited before we could notify systemd")
	}

	if len(c.NotifySocket) == 0 {
		return nil
	}

	conn, err := net.Dial("unixgram", c.NotifySocket)
	if err != nil {
		return err
	}

	defer conn.Close()

	_, err = conn.Write([]byte(fmt.Sprintf("MAINPID=%d", c.Pid)))
	if err != nil {
		return err
	}

	if pidDied(c.Pid) {
		conn.Write([]byte(fmt.Sprintf("MAINPID=%d", os.Getpid())))
		return errors.New("Container exited before we could notify systemd")
	}

	if !c.Notify {
		_, err = conn.Write([]byte("READY=1"))
		if err != nil {
			return err
		}
	}

	return nil
}

func pidFile(c *Context) error {
	if len(c.PidFile) == 0 || c.Pid <= 0 {
		return nil
	}

	err := ioutil.WriteFile(c.PidFile, []byte(strconv.Itoa(c.Pid)), 0644)
	if err != nil {
		return err
	}

	return nil
}

func pipeLogs(c *Context) error {
	if !c.Logs {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	err = client.Logs(dockerClient.LogsOptions{
		Container:    c.Id,
		Follow:       true,
		Stdout:       true,
		Stderr:       true,
		OutputStream: os.Stdout,
		ErrorStream:  os.Stderr,
	})

	return err
}

func keepAlive(c *Context) error {
	if c.Logs || c.Rm {
		client, err := getClient(c)
		if err != nil {
			return err
		}

		/* Good old polling... */
		for true {
			container, err := client.InspectContainer(c.Id)
			if err != nil {
				return err
			}

			if container.State.Running {
				client.WaitContainer(c.Id)
			} else {
				return nil
			}
		}
	}

	return nil
}

func rmContainer(c *Context) error {
	if !c.Rm {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	return client.RemoveContainer(dockerClient.RemoveContainerOptions{
		ID:    c.Id,
		Force: true,
	})
}

func mainWithArgs(args []string) (*Context, error) {
	c, err := parseContext(args)
	if err != nil {
		return c, err
	}

	err = runContainer(c)
	if err != nil {
		return c, err
	}

	_, err = moveCgroups(c)
	if err != nil {
		return c, err
	}

	err = notify(c)
	if err != nil {
		return c, err
	}

	err = pidFile(c)
	if err != nil {
		return c, err
	}

	go pipeLogs(c)

	err = keepAlive(c)
	if err != nil {
		return c, err
	}

	err = rmContainer(c)
	if err != nil {
		return c, err
	}

	return c, nil
}

func main() {
	_, err := mainWithArgs(os.Args[1:])
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>Update version<commit_after>package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/exec"
	"path"
	"strconv"
	"strings"
	"time"

	"github.com/docker/docker/opts"
	flag "github.com/docker/docker/pkg/mflag"

	dockerClient "github.com/fsouza/go-dockerclient"
)

var (
	SYSFS       string        = "/sys/fs/cgroup"
	PROCS       string        = "cgroup.procs"
	CGROUP_PROC string        = "/proc/%d/cgroup"
	INTERVAL    time.Duration = 1000
)

type Context struct {
	Args         []string
	Cgroups      []string
	AllCgroups   bool
	Logs         bool
	Notify       bool
	Name         string
	Env          bool
	Rm           bool
	Id           string
	NotifySocket string
	Cmd          *exec.Cmd
	Pid          int
	PidFile      string
	Client       *dockerClient.Client
}

func setupEnvironment(c *Context) {
	newArgs := []string{}
	if c.Notify && len(c.NotifySocket) > 0 {
		newArgs = append(newArgs, "-e", fmt.Sprintf("NOTIFY_SOCKET=%s", c.NotifySocket))
		newArgs = append(newArgs, "-v", fmt.Sprintf("%s:%s", c.NotifySocket, c.NotifySocket))
	} else {
		c.Notify = false
	}

	if c.Env {
		for _, val := range os.Environ() {
			if !strings.HasPrefix(val, "HOME=") && !strings.HasPrefix(val, "PATH=") {
				newArgs = append(newArgs, "-e", val)
			}
		}
	}

	if len(newArgs) > 0 {
		c.Args = append(newArgs, c.Args...)
	}
}

func parseContext(args []string) (*Context, error) {
	c := &Context{
		Logs:       true,
		AllCgroups: false,
	}

	flags := flag.NewFlagSet("systemd-docker", flag.ContinueOnError)

	flCgroups := opts.NewListOpts(nil)

	flags.StringVar(&c.PidFile, []string{"p", "-pid-file"}, "", "pipe file")
	flags.BoolVar(&c.Logs, []string{"l", "-logs"}, true, "pipe logs")
	flags.BoolVar(&c.Notify, []string{"n", "-notify"}, false, "setup systemd notify for container")
	flags.BoolVar(&c.Env, []string{"e", "-env"}, false, "inherit environment variable")
	flags.Var(&flCgroups, []string{"c", "-cgroups"}, "cgroups to take ownership of or 'all' for all cgroups available")

	err := flags.Parse(args)
	if err != nil {
		return nil, err
	}

	foundD := false
	var name string

	runArgs := flags.Args()
	if len(runArgs) == 0 || runArgs[0] != "run" {
		log.Println("Args:", runArgs)
		return nil, errors.New("run not found in arguments")
	}

	runArgs = runArgs[1:]
	newArgs := make([]string, 0, len(runArgs))

	for i, arg := range runArgs {
		/* This is tedious, but flag can't ignore unknown flags and I don't want to define them all */
		add := true

		switch {
		case arg == "-rm" || arg == "--rm":
			c.Rm = true
			add = false
		case arg == "-d" || arg == "-detach" || arg == "--detach":
			foundD = true
		case strings.HasPrefix(arg, "-name") || strings.HasPrefix(arg, "--name"):
			if strings.Contains(arg, "=") {
				name = strings.SplitN(arg, "=", 2)[1]
			} else if len(runArgs) > i+1 {
				name = runArgs[i+1]
			}
		}

		if add {
			newArgs = append(newArgs, arg)
		}
	}

	if !foundD {
		newArgs = append([]string{"-d"}, newArgs...)
	}

	c.Name = name
	c.NotifySocket = os.Getenv("NOTIFY_SOCKET")
	c.Args = newArgs
	c.Cgroups = flCgroups.GetAll()

	for _, val := range c.Cgroups {
		if val == "all" {
			c.Cgroups = nil
			c.AllCgroups = true
			break
		}
	}

	setupEnvironment(c)

	return c, nil
}

func lookupNamedContainer(c *Context) error {
	client, err := getClient(c)
	if err != nil {
		return err
	}

	container, err := client.InspectContainer(c.Name)
	if _, ok := err.(*dockerClient.NoSuchContainer); ok {
		return nil
	}
	if err != nil || container == nil {
		return err
	}

	if container.State.Running {
		c.Id = container.ID
		c.Pid = container.State.Pid
		return nil
	} else if c.Rm {
		return client.RemoveContainer(dockerClient.RemoveContainerOptions{
			ID:    container.ID,
			Force: true,
		})
	} else {
		client, err := getClient(c)
		err = client.StartContainer(container.ID, container.HostConfig)
		if err != nil {
			return err
		}

		container, err = client.InspectContainer(c.Name)
		if err != nil {
			return err
		}

		c.Id = container.ID
		c.Pid = container.State.Pid

		return nil
	}
}

func launchContainer(c *Context) error {
	args := append([]string{"run"}, c.Args...)
	c.Cmd = exec.Command("docker", args...)

	errorPipe, err := c.Cmd.StderrPipe()
	if err != nil {
		return err
	}

	outputPipe, err := c.Cmd.StdoutPipe()
	if err != nil {
		return err
	}

	err = c.Cmd.Start()
	if err != nil {
		return err
	}

	go io.Copy(os.Stderr, errorPipe)

	bytes, err := ioutil.ReadAll(outputPipe)
	if err != nil {
		return err
	}

	c.Id = strings.TrimSpace(string(bytes))

	err = c.Cmd.Wait()
	if err != nil {
		return err
	}

	if !c.Cmd.ProcessState.Success() {
		return err
	}

	c.Pid, err = getContainerPid(c)

	return err
}

func runContainer(c *Context) error {
	if len(c.Name) > 0 {
		err := lookupNamedContainer(c)
		if err != nil {
			return err
		}

	}

	if len(c.Id) == 0 {
		err := launchContainer(c)
		if err != nil {
			return err
		}
	}

	if c.Pid == 0 {
		return errors.New("Failed to launch container, pid is 0")
	}

	return nil
}

func getClient(c *Context) (*dockerClient.Client, error) {
	if c.Client != nil {
		return c.Client, nil
	}

	endpoint := os.Getenv("DOCKER_HOST")
	if len(endpoint) == 0 {
		endpoint = "unix:///var/run/docker.sock"
	}

	return dockerClient.NewVersionedClient(endpoint, "1.12")
}

func getContainerPid(c *Context) (int, error) {
	client, err := getClient(c)
	if err != nil {
		return 0, err
	}

	container, err := client.InspectContainer(c.Id)
	if err != nil {
		return 0, err
	}

	if container == nil {
		return 0, errors.New(fmt.Sprintf("Failed to find container %s", c.Id))
	}

	if container.State.Pid <= 0 {
		return 0, errors.New(fmt.Sprintf("Pid is %d for container %s", container.State.Pid, c.Id))
	}

	return container.State.Pid, nil
}

func getCgroupsForPid(pid int) (map[string]string, error) {
	file, err := os.Open(fmt.Sprintf(CGROUP_PROC, pid))
	if err != nil {
		return nil, err
	}

	ret := map[string]string{}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		line := strings.SplitN(scanner.Text(), ":", 3)
		if len(line) != 3 {
			continue
		}

		ret[line[1]] = line[2]
	}

	if err := scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func constructCgroupPath(cgroupName string, cgroupPath string) string {
	return path.Join(SYSFS, strings.TrimPrefix(cgroupName, "name="), cgroupPath, PROCS)
}

func getCgroupPids(cgroupName string, cgroupPath string) ([]string, error) {
	ret := []string{}

	file, err := os.Open(constructCgroupPath(cgroupName, cgroupPath))
	if err != nil {
		return nil, err
	}

	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		ret = append(ret, strings.TrimSpace(scanner.Text()))
	}

	if err = scanner.Err(); err != nil {
		return nil, err
	}

	return ret, nil
}

func writePid(pid string, path string) error {
	return ioutil.WriteFile(path, []byte(pid), 0644)
}

func moveCgroups(c *Context) (bool, error) {
	moved := false
	currentCgroups, err := getCgroupsForPid(os.Getpid())
	if err != nil {
		return false, err
	}

	containerCgroups, err := getCgroupsForPid(c.Pid)
	if err != nil {
		return false, err
	}

	var ns []string

	if c.AllCgroups || c.Cgroups == nil || len(c.Cgroups) == 0 {
		ns = make([]string, 0, len(containerCgroups))
		for value, _ := range containerCgroups {
			ns = append(ns, value)
		}
	} else {
		ns = c.Cgroups
	}

	for _, nsName := range ns {
		currentPath, ok := currentCgroups[nsName]
		if !ok {
			continue
		}

		containerPath, ok := containerCgroups[nsName]
		if !ok {
			continue
		}

		if currentPath == containerPath || containerPath == "/" {
			continue
		}

		pids, err := getCgroupPids(nsName, containerPath)
		if err != nil {
			return false, err
		}

		for _, pid := range pids {
			pidInt, err := strconv.Atoi(pid)
			if err != nil {
				continue
			}

			if pidDied(pidInt) {
				continue
			}

			currentFullPath := constructCgroupPath(nsName, currentPath)
			log.Printf("Moving pid %s to %s\n", pid, currentFullPath)
			err = writePid(pid, currentFullPath)
			if err != nil {
				return false, err
			}

			moved = true
		}
	}

	return moved, nil
}

func pidDied(pid int) bool {
	_, err := os.Stat(fmt.Sprintf("/proc/%d", pid))
	return os.IsNotExist(err)
}

func notify(c *Context) error {
	if pidDied(c.Pid) {
		return errors.New("Container exited before we could notify systemd")
	}

	if len(c.NotifySocket) == 0 {
		return nil
	}

	conn, err := net.Dial("unixgram", c.NotifySocket)
	if err != nil {
		return err
	}

	defer conn.Close()

	_, err = conn.Write([]byte(fmt.Sprintf("MAINPID=%d", c.Pid)))
	if err != nil {
		return err
	}

	if pidDied(c.Pid) {
		conn.Write([]byte(fmt.Sprintf("MAINPID=%d", os.Getpid())))
		return errors.New("Container exited before we could notify systemd")
	}

	if !c.Notify {
		_, err = conn.Write([]byte("READY=1"))
		if err != nil {
			return err
		}
	}

	return nil
}

func pidFile(c *Context) error {
	if len(c.PidFile) == 0 || c.Pid <= 0 {
		return nil
	}

	err := ioutil.WriteFile(c.PidFile, []byte(strconv.Itoa(c.Pid)), 0644)
	if err != nil {
		return err
	}

	return nil
}

func pipeLogs(c *Context) error {
	if !c.Logs {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	err = client.Logs(dockerClient.LogsOptions{
		Container:    c.Id,
		Follow:       true,
		Stdout:       true,
		Stderr:       true,
		OutputStream: os.Stdout,
		ErrorStream:  os.Stderr,
	})

	return err
}

func keepAlive(c *Context) error {
	if c.Logs || c.Rm {
		client, err := getClient(c)
		if err != nil {
			return err
		}

		/* Good old polling... */
		for true {
			container, err := client.InspectContainer(c.Id)
			if err != nil {
				return err
			}

			if container.State.Running {
				client.WaitContainer(c.Id)
			} else {
				return nil
			}
		}
	}

	return nil
}

func rmContainer(c *Context) error {
	if !c.Rm {
		return nil
	}

	client, err := getClient(c)
	if err != nil {
		return err
	}

	return client.RemoveContainer(dockerClient.RemoveContainerOptions{
		ID:    c.Id,
		Force: true,
	})
}

func mainWithArgs(args []string) (*Context, error) {
	c, err := parseContext(args)
	if err != nil {
		return c, err
	}

	err = runContainer(c)
	if err != nil {
		return c, err
	}

	_, err = moveCgroups(c)
	if err != nil {
		return c, err
	}

	err = notify(c)
	if err != nil {
		return c, err
	}

	err = pidFile(c)
	if err != nil {
		return c, err
	}

	go pipeLogs(c)

	err = keepAlive(c)
	if err != nil {
		return c, err
	}

	err = rmContainer(c)
	if err != nil {
		return c, err
	}

	return c, nil
}

func main() {
	_, err := mainWithArgs(os.Args[1:])
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>


<commit_before>// Virtua SmartOS Prometheus exporter
//
// Worflow :
//  - detect if launched in GZ or into a Zone
//  - retrieve useful metrics
//  - expose them to http://xxx:9100/metrics (same as node_exporter)

package main

import (
	"log"
	"net/http"
	"os/exec"
	"strings"
	//  "fmt"

	"github.com/virtua-network/smartos_exporter/collector"

	// Prometheus Go toolset
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	// Global variables
	exporterPort = ":9100"
)

// Global Helpers

// try to determine if its executed inside the GZ or not.
// return 1 if in GZ
//        0 if in zone
func isGlobalZone() int {
	out, eerr := exec.Command("zonename").Output()
	if eerr != nil {
		log.Fatal(eerr)
	}
	if (strings.Contains(string(out), "global")) == false {
		return 0
	}
	return 1
}

// XXX add a OS check in func init
// func init()

// program starter
func main() {
	// check if it is a GZ or a zone
	gz := isGlobalZone()

	// common metrics
	loadAvg, _ := collector.NewLoadAverageExporter()
	prometheus.MustRegister(loadAvg)

	if gz == 0 {
		// Zone metrics
		zoneDf, _ := collector.NewZoneDfExporter()
		prometheus.MustRegister(zoneDf)
		zoneKstat, _ := collector.NewZoneKstatExporter()
		prometheus.MustRegister(zoneKstat)
	}

	if gz == 1 {
		// Global Zone metrics
		gzFreeMem, _ := collector.NewGZFreeMemExporter()
		prometheus.MustRegister(gzFreeMem)

		gzMLAGUsage, _ := collector.NewGZMLAGUsageExporter()
		prometheus.MustRegister(gzMLAGUsage)

		loadAvg, _ := collector.NewLoadAverageExporter()
		prometheus.MustRegister(loadAvg)

		cpuUsage, _ := collector.NewGZCPUUsageExporter()
		prometheus.MustRegister(cpuUsage)

		gzDiskErrors, _ := collector.NewGZDiskErrorsExporter()
		prometheus.MustRegister(gzDiskErrors)

		gzZpoolList, _ := collector.NewGZZpoolListExporter()
		prometheus.MustRegister(gzZpoolList)
	}

	// The Handler function provides a default handler to expose metrics
	// via an HTTP server. "/metrics" is the usual endpoint for that.
	http.Handle("/metrics", promhttp.Handler())
	log.Fatal(http.ListenAndServe(exporterPort, nil))
}
<commit_msg>Fix LoadAverage collector which was duplicated on GZ<commit_after>// Virtua SmartOS Prometheus exporter
//
// Worflow :
//  - detect if launched in GZ or into a Zone
//  - retrieve useful metrics
//  - expose them to http://xxx:9100/metrics (same as node_exporter)

package main

import (
	"log"
	"net/http"
	"os/exec"
	"strings"
	//  "fmt"

	"github.com/virtua-network/smartos_exporter/collector"

	// Prometheus Go toolset
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

var (
	// Global variables
	exporterPort = ":9100"
)

// Global Helpers

// try to determine if its executed inside the GZ or not.
// return 1 if in GZ
//        0 if in zone
func isGlobalZone() int {
	out, eerr := exec.Command("zonename").Output()
	if eerr != nil {
		log.Fatal(eerr)
	}
	if (strings.Contains(string(out), "global")) == false {
		return 0
	}
	return 1
}

// XXX add a OS check in func init
// func init()

// program starter
func main() {
	// check if it is a GZ or a zone
	gz := isGlobalZone()

	// common metrics
	loadAvg, _ := collector.NewLoadAverageExporter()
	prometheus.MustRegister(loadAvg)

	if gz == 0 {
		// Zone metrics
		zoneDf, _ := collector.NewZoneDfExporter()
		prometheus.MustRegister(zoneDf)

		zoneKstat, _ := collector.NewZoneKstatExporter()
		prometheus.MustRegister(zoneKstat)
	}

	if gz == 1 {
		// Global Zone metrics
		gzFreeMem, _ := collector.NewGZFreeMemExporter()
		prometheus.MustRegister(gzFreeMem)

		gzMLAGUsage, _ := collector.NewGZMLAGUsageExporter()
		prometheus.MustRegister(gzMLAGUsage)

		cpuUsage, _ := collector.NewGZCPUUsageExporter()
		prometheus.MustRegister(cpuUsage)

		gzDiskErrors, _ := collector.NewGZDiskErrorsExporter()
		prometheus.MustRegister(gzDiskErrors)

		gzZpoolList, _ := collector.NewGZZpoolListExporter()
		prometheus.MustRegister(gzZpoolList)
	}

	// The Handler function provides a default handler to expose metrics
	// via an HTTP server. "/metrics" is the usual endpoint for that.
	http.Handle("/metrics", promhttp.Handler())
	log.Fatal(http.ListenAndServe(exporterPort, nil))
}
<|endoftext|>


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

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/lancetw/hcfd-forecast/db"
	"github.com/lancetw/hcfd-forecast/rain"
	"github.com/line/line-bot-sdk-go/linebot"
)

const timeZone = "Asia/Taipei"

var bot *linebot.Client

func main() {
	strID := os.Getenv("ChannelID")
	numID, err := strconv.ParseInt(strID, 10, 64)
	if err != nil {
		log.Fatal("Wrong environment setting about ChannelID")
	}
	bot, err = linebot.NewClient(numID, os.Getenv("ChannelSecret"), os.Getenv("MID"))
	if err != nil {
		log.Println("Bot:", bot, " err:", err)
	}

	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)
}

func callbackHandler(w http.ResponseWriter, r *http.Request) {
	received, err := bot.ParseRequest(r)
	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}
	for _, result := range received.Results {
		content := result.Content()
		if content != nil && content.IsOperation && content.OpType == linebot.OpTypeAddedAsFriend {
			op, err := content.OperationContent()
			if err != nil {
				log.Println(err)
				return
			}
			from := op.Params[0]

			user, err := bot.GetUserProfile([]string{from})
			if err != nil {
				log.Println(err)
				return
			}
			_, err = bot.SendText([]string{from}, user.Contacts[0].DisplayName+" 您好，目前可用指令為：「加入」「退出」「雨量」「警報」「貓圖」「狀態」「時間」")
			if err != nil {
				log.Println(err)
			}
		}

		if content != nil && content.IsMessage && content.ContentType == linebot.ContentTypeText {
			user, err := bot.GetUserProfile([]string{content.From})
			if err != nil {
				return
			}
			text, err := content.TextContent()
			if err != nil {
				log.Println(err)
			}
			cmd := strings.Fields(text.Text)

			switch cmd[0] {
			case "加入":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status, addErr := c.Do("SADD", "user", content.From)
				if addErr != nil {
					log.Println("SADD to redis error", addErr, status)
				} else {
					_, err = bot.SendText([]string{content.From}, user.Contacts[0].DisplayName+" 您好，已將您加入傳送對象，未來將會傳送天氣警報資訊給您 ^＿^ ")
					if err != nil {
						log.Println(err)
					}
				}
				defer c.Close()

			case "退出":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status, setErr := c.Do("SREM", "user", content.From)
				if setErr != nil {
					log.Println("DEL to redis error", setErr, status)
				} else {
					_, err = bot.SendText([]string{content.From}, user.Contacts[0].DisplayName+" 掰掰 Q＿Q")
					if err != nil {
						log.Println(err)
					}
				}
				defer c.Close()

			case "服務":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				count, countErr := redis.Int(c.Do("SCARD", "user"))
				if countErr != nil {
					log.Println(countErr)
				} else {
					_, err = bot.SendText([]string{content.From}, fmt.Sprintf("目前有 %d 人加入自動警訊服務。", count))
					if err != nil {
						log.Println(err)
					}
				}

				defer c.Close()

			case "狀態":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status, getErr := redis.Int(c.Do("SISMEMBER", "user", content.From))
				if getErr != nil || status == 0 {
					_, err = bot.SendText([]string{content.From}, "目前沒有登記您的編號喔！")
					if err != nil {
						log.Println(err)
					}
				} else {
					_, err = bot.SendText([]string{content.From}, "您已經是傳送對象 :D")
					if err != nil {
						log.Println(err)
					}
				}
				defer c.Close()

			case "時間":
				local := time.Now()
				location, timezoneErr := time.LoadLocation(timeZone)
				if timezoneErr == nil {
					local = local.In(location)
				}
				_, err = bot.SendText([]string{content.From}, local.Format("2006/01/02 15:04:05"))
				if err != nil {
					log.Println(err)
				}

			case "雨量":
				target := []string{"新竹市"}
				if len(cmd) > 1 {
					target[0] = cmd[1]
				}

				msgs, _ := rain.GetRainingInfo(target, true)
				if len(msgs) == 0 {
					_, err = bot.SendText([]string{content.From}, "目前沒有雨量資訊！")
					if err != nil {
						log.Println(err)
					}
				} else {
					for _, msg := range msgs {
						_, err = bot.SendText([]string{content.From}, msg)
						if err != nil {
							log.Println(err)
						}
					}
				}

			case "警報":
				msgs, _ := rain.GetWarningInfo(nil)
				for _, msg := range msgs {
					_, err = bot.SendText([]string{content.From}, msg)
					if err != nil {
						log.Println(err)
					}
				}

			case "重開":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status0, clearErr0 := c.Do("DEL", "token0")
				if clearErr0 != nil {
					log.Println("DEL to redis error", clearErr0, status0)
				}
				status1, clearErr1 := c.Do("DEL", "token1")
				if clearErr1 != nil {
					log.Println("DEL to redis error", clearErr1, status1)
				}

				_, err = bot.SendText([]string{content.From}, "已重開")
				if err != nil {
					log.Println(err)
				}

			case "貓圖":
				type Cat struct {
					File string `json:"file"`
				}
				cat := new(Cat)
				getJSON("http://random.cat/meow", &cat)
				if cat.File != "" {
					_, err = bot.SendImage([]string{content.From}, cat.File, cat.File)
					if err != nil {
						log.Println(err)
					}
				}

			case "妹子":
				type MeisData struct {
					ID         string `json:"id"`
					Name       string `json:"name"`
					ImgNo      string `json:"img_no"`
					Fanpage    string `json:"fanpage"`
					Creator    string `json:"creator"`
					UpdateTime string `json:"update_time"`
				}

				type Beauty struct {
					Meis map[string]MeisData `json:"meis"`
				}

				beauty := new(Beauty)
				getJSONErr := getJSON("http://beauty.zones.gamebase.com.tw/wall?json", &beauty)
				if len(beauty.Meis) > 0 {
					for fbid, data := range beauty.Meis {
						image := fmt.Sprintf("http://graph.facebook.com/%s/picture?type=large", fbid)
						_, err = bot.SendImage([]string{content.From}, image, image)
						if err != nil {
							log.Println(err)
						}

						link := fmt.Sprintf("http://www.facebook.com/%s", data.ID)
						description := fmt.Sprintf("%s %s", data.Name, link)
						_, err = bot.SendText([]string{content.From}, description)
						if err != nil {
							log.Println(err)
						}
						break
					}

				}

				if getJSONErr != nil {
					log.Println(getJSONErr)
				}

			default:
				_, err = bot.SendText([]string{content.From}, "指令錯誤，請重試")
				if err != nil {
					log.Println(err)
				}

			}
		}
	}
}

// getJSON func
func getJSON(url string, target interface{}) error {
	r, err := http.Get(url)
	if err != nil {
		return err
	}
	defer r.Body.Close()

	return json.NewDecoder(r.Body).Decode(target)
}
<commit_msg>fixed typo<commit_after>// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/garyburd/redigo/redis"
	"github.com/lancetw/hcfd-forecast/db"
	"github.com/lancetw/hcfd-forecast/rain"
	"github.com/line/line-bot-sdk-go/linebot"
)

const timeZone = "Asia/Taipei"

var bot *linebot.Client

func main() {
	strID := os.Getenv("ChannelID")
	numID, err := strconv.ParseInt(strID, 10, 64)
	if err != nil {
		log.Fatal("Wrong environment setting about ChannelID")
	}
	bot, err = linebot.NewClient(numID, os.Getenv("ChannelSecret"), os.Getenv("MID"))
	if err != nil {
		log.Println("Bot:", bot, " err:", err)
	}

	http.HandleFunc("/callback", callbackHandler)
	port := os.Getenv("PORT")
	addr := fmt.Sprintf(":%s", port)
	http.ListenAndServe(addr, nil)
}

func callbackHandler(w http.ResponseWriter, r *http.Request) {
	received, err := bot.ParseRequest(r)
	if err != nil {
		if err == linebot.ErrInvalidSignature {
			w.WriteHeader(400)
		} else {
			w.WriteHeader(500)
		}
		return
	}
	for _, result := range received.Results {
		content := result.Content()
		if content != nil && content.IsOperation && content.OpType == linebot.OpTypeAddedAsFriend {
			op, err := content.OperationContent()
			if err != nil {
				log.Println(err)
				return
			}
			from := op.Params[0]

			user, err := bot.GetUserProfile([]string{from})
			if err != nil {
				log.Println(err)
				return
			}
			_, err = bot.SendText([]string{from}, user.Contacts[0].DisplayName+" 您好，目前可用指令為：「加入」「退出」「雨量」「警報」「貓圖」「狀態」「時間」")
			if err != nil {
				log.Println(err)
			}
		}

		if content != nil && content.IsMessage && content.ContentType == linebot.ContentTypeText {
			user, err := bot.GetUserProfile([]string{content.From})
			if err != nil {
				return
			}
			text, err := content.TextContent()
			if err != nil {
				log.Println(err)
			}
			cmd := strings.Fields(text.Text)

			switch cmd[0] {
			case "加入":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status, addErr := c.Do("SADD", "user", content.From)
				if addErr != nil {
					log.Println("SADD to redis error", addErr, status)
				} else {
					_, err = bot.SendText([]string{content.From}, user.Contacts[0].DisplayName+" 您好，已將您加入傳送對象，未來將會傳送天氣警報資訊給您 ^＿^ ")
					if err != nil {
						log.Println(err)
					}
				}
				defer c.Close()

			case "退出":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status, setErr := c.Do("SREM", "user", content.From)
				if setErr != nil {
					log.Println("DEL to redis error", setErr, status)
				} else {
					_, err = bot.SendText([]string{content.From}, user.Contacts[0].DisplayName+" 掰掰 Q＿Q")
					if err != nil {
						log.Println(err)
					}
				}
				defer c.Close()

			case "服務":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				count, countErr := redis.Int(c.Do("SCARD", "user"))
				if countErr != nil {
					log.Println(countErr)
				} else {
					_, err = bot.SendText([]string{content.From}, fmt.Sprintf("目前有 %d 人加入自動警訊服務。", count))
					if err != nil {
						log.Println(err)
					}
				}

				defer c.Close()

			case "狀態":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status, getErr := redis.Int(c.Do("SISMEMBER", "user", content.From))
				if getErr != nil || status == 0 {
					_, err = bot.SendText([]string{content.From}, "目前沒有登記您的編號喔！")
					if err != nil {
						log.Println(err)
					}
				} else {
					_, err = bot.SendText([]string{content.From}, "您已經是傳送對象 :D")
					if err != nil {
						log.Println(err)
					}
				}
				defer c.Close()

			case "時間":
				local := time.Now()
				location, timezoneErr := time.LoadLocation(timeZone)
				if timezoneErr == nil {
					local = local.In(location)
				}
				_, err = bot.SendText([]string{content.From}, local.Format("2006/01/02 15:04:05"))
				if err != nil {
					log.Println(err)
				}

			case "雨量":
				target := []string{"新竹市"}
				if len(cmd) > 1 {
					target[0] = cmd[1]
				}

				msgs, _ := rain.GetRainingInfo(target, true)
				if len(msgs) == 0 {
					_, err = bot.SendText([]string{content.From}, "目前沒有雨量資訊！")
					if err != nil {
						log.Println(err)
					}
				} else {
					for _, msg := range msgs {
						_, err = bot.SendText([]string{content.From}, msg)
						if err != nil {
							log.Println(err)
						}
					}
				}

			case "警報":
				msgs, _ := rain.GetWarningInfo(nil)
				for _, msg := range msgs {
					_, err = bot.SendText([]string{content.From}, msg)
					if err != nil {
						log.Println(err)
					}
				}

			case "重開":
				c := db.Connect(os.Getenv("REDISTOGO_URL"))
				status0, clearErr0 := c.Do("DEL", "token0")
				if clearErr0 != nil {
					log.Println("DEL to redis error", clearErr0, status0)
				}
				status1, clearErr1 := c.Do("DEL", "token1")
				if clearErr1 != nil {
					log.Println("DEL to redis error", clearErr1, status1)
				}

				_, err = bot.SendText([]string{content.From}, "已重開")
				if err != nil {
					log.Println(err)
				}

			case "貓圖":
				type Cat struct {
					File string `json:"file"`
				}
				cat := new(Cat)
				getJSON("http://random.cat/meow", &cat)
				if cat.File != "" {
					_, err = bot.SendImage([]string{content.From}, cat.File, cat.File)
					if err != nil {
						log.Println(err)
					}
				}

			case "妹子":
				type MeisData struct {
					ID         string `json:"id"`
					Name       string `json:"name"`
					ImgNo      string `json:"img_no"`
					Fanpage    string `json:"fanpage"`
					Creator    string `json:"creator"`
					UpdateTime string `json:"update_time"`
				}

				type Beauty struct {
					Meis map[string]MeisData `json:"meis"`
				}

				beauty := new(Beauty)
				getJSONErr := getJSON("http://beauty.zones.gamebase.com.tw/wall?json", &beauty)
				if len(beauty.Meis) > 0 {
					for fbid, data := range beauty.Meis {
						image := fmt.Sprintf("http://graph.facebook.com/%s/picture?type=large", fbid)
						_, err = bot.SendImage([]string{content.From}, image, image)
						if err != nil {
							log.Println(err)
						}

						link := fmt.Sprintf("http://www.facebook.com/%s", fbid)
						description := fmt.Sprintf("%s %s", data.Name, link)
						_, err = bot.SendText([]string{content.From}, description)
						if err != nil {
							log.Println(err)
						}
						break
					}

				}

				if getJSONErr != nil {
					log.Println(getJSONErr)
				}

			default:
				_, err = bot.SendText([]string{content.From}, "指令錯誤，請重試")
				if err != nil {
					log.Println(err)
				}

			}
		}
	}
}

// getJSON func
func getJSON(url string, target interface{}) error {
	r, err := http.Get(url)
	if err != nil {
		return err
	}
	defer r.Body.Close()

	return json.NewDecoder(r.Body).Decode(target)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
)

func main() {
	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

	cwe := cloudwatchevents.New(sess)
	result, err := cwe.ListRules(nil)

	if err != nil {
		fmt.Println("Error", err)

	} else {
		fmt.Println("Success")
		fmt.Println(result)
	}
}
<commit_msg>Parse basic flags apply and dry-run<commit_after>package main

import (
	"flag"
	"fmt"

	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatchevents"
)

func main() {
	var apply bool
	var dryrun bool

	flag.BoolVar(&apply, "apply", false, "apply to CloudWatch Events")
	flag.BoolVar(&dryrun, "dry-run", false, "dry-run")
	flag.Parse()

	fmt.Println(apply)
	fmt.Println(dryrun)

	sess := session.Must(session.NewSessionWithOptions(session.Options{
		SharedConfigState: session.SharedConfigEnable,
	}))

	cwe := cloudwatchevents.New(sess)
	result, err := cwe.ListRules(nil)

	if err != nil {
		fmt.Println("Error", err)

	} else {
		fmt.Println("Success")
		fmt.Println(result)
	}
}
<|endoftext|>


<commit_before>package main

import (
    "log"
    "math/rand"
    "time"
)

func main() {
    log.Print("started.")

    // チャネル
    sleep1_finished := make(chan bool)

    go func() {
        // 0.2秒かかるコマンド
        log.Print("sleep1 started.")
        time.Sleep(200 * 1000 * 1000 * time.Nanosecond)
        log.Print("sleep1 finished.")
        sleep1_finished <- true
    }()

    // 終わるまで待つ
    <- sleep1_finished

    dice()
}

func dice() {
    log.Printf("%v", random())
}

// http://qiita.com/cubicdaiya/items/819886c57e9d17e4b019
// 3までの数字をランダムに返す
func random()(int) {
    rand.Seed(time.Now().UnixNano())
    return rand.Intn(4)
}
<commit_msg>Update<commit_after>package main

import (
    "log"
    "math/rand"
    "time"
)

func main() {
    log.Print("started.")

    // チャネル
    sleep1_finished := make(chan bool)

    go func() {
        // 0.2秒かかるコマンド
        log.Print("sleep1 started.")
        time.Sleep(200 * 1000 * 1000 * time.Nanosecond)
        log.Print("sleep1 finished.")
        sleep1_finished <- true
    }()

    // 終わるまで待つ
    <- sleep1_finished

    dice()
}

func dice() {
    log.Printf("%v", random())
}

// http://qiita.com/cubicdaiya/items/819886c57e9d17e4b019
// 東西南北をランダムに返す
func random()(string) {
    a := []string{"ton", "nan", "sha", "pe"}
    rand.Seed(time.Now().UnixNano())
    return a[rand.Intn(4)]
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/nacl/secretbox"
	"io"
	"log"
	"net"
	"os"
)

type header struct {
	Length uint64
	Nonce  [24]byte
}

type secureReader struct {
	src io.Reader
	key *[32]byte
	buf []byte
}

type secureWriter struct {
	dst io.Writer
	key *[32]byte
}

type secureConn struct {
	io.Reader
	io.Writer
	io.Closer
}

func generateNonce(nonce *[24]byte) error {
	_, err := io.ReadFull(rand.Reader, nonce[:])
	return err
}

func (p secureReader) Read(b []byte) (int, error) {
	if p.buf == nil {
		var h header
		if err := binary.Read(p.src, binary.LittleEndian, &h); err != nil {
			return 0, err
		}

		p.buf = make([]byte, h.Length-secretbox.Overhead)
		tmp := make([]byte, h.Length)

		if _, err := io.ReadFull(p.src, tmp); err != nil {
			return 0, err
		}

		_, auth := secretbox.Open(p.buf[:0], tmp[:], &h.Nonce, p.key)
		if !auth {
			return 0, fmt.Errorf("Message failed authentication")
		}
	}
	var size int
	if len(b) >= len(p.buf) {
		size = len(p.buf)
	} else {
		size = len(b)
	}

	for i := 0; i < size; i++ {
		b[i] = p.buf[i]
	}

	if size == len(p.buf) {
		p.buf = nil
	} else {
		p.buf = p.buf[size:]
	}
	return size, nil
}

func (p secureWriter) Write(b []byte) (int, error) {
	var h header
	h.Length = uint64(len(b) + secretbox.Overhead)
	if err := generateNonce(&h.Nonce); err != nil {
		return 0, err
	}

	hbuf := new(bytes.Buffer)
	binary.Write(hbuf, binary.LittleEndian, &h)
	if _, err := p.dst.Write(hbuf.Bytes()); err != nil {
		return 0, err
	}

	tmp := make([]byte, h.Length)
	secretbox.Seal(tmp[:0], b, &h.Nonce, p.key)
	return p.dst.Write(tmp)
}

// NewSecureReader instantiates a new SecureReader
func NewSecureReader(r io.Reader, priv, pub *[32]byte) io.Reader {
	sr := secureReader{src: r, key: new([32]byte), buf: nil}
	box.Precompute(sr.key, pub, priv)
	return sr
}

// NewSecureWriter instantiates a new SecureWriter
func NewSecureWriter(w io.Writer, priv, pub *[32]byte) io.Writer {
	sw := secureWriter{dst: w, key: new([32]byte)}
	box.Precompute(sw.key, pub, priv)
	return sw
}

func generateKeys() (priv, pub *[32]byte, err error) {
	pub, priv, err = box.GenerateKey(rand.Reader)
	return
}

// Dial generates a private/public key pair,
// connects to the server, perform the handshake
// and return a reader/writer.
func Dial(addr string) (io.ReadWriteCloser, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	priv, pub, err := generateKeys()
	if err != nil {
		conn.Close()
		return nil, err
	}

	if _, err := conn.Write(pub[:]); err != nil {
		conn.Close()
		return nil, err
	}

	servpub := new([32]byte)
	if _, err := io.ReadFull(conn, servpub[:]); err != nil {
		conn.Close()
		return nil, err
	}

	secconn := secureConn{
		NewSecureReader(conn, priv, servpub),
		NewSecureWriter(conn, priv, servpub),
		conn}

	return secconn, nil
}

func connect(c net.Conn) {
	defer c.Close()

	clientpub := new([32]byte)
	_, err := io.ReadFull(c, clientpub[:])
	if err != nil {
		log.Fatal(err)
	}

	priv, pub, err := generateKeys()
	if err != nil {
		log.Fatal(err)
	}
	if _, err := c.Write(pub[:]); err != nil {
		log.Fatal(err)
	}

	r := NewSecureReader(c, priv, clientpub)
	w := NewSecureWriter(c, priv, clientpub)

	buf := make([]byte, 2048)
	for {
		n, err := r.Read(buf)
		if err != nil && err != io.EOF {
			log.Fatal(err)
		}
		if n > 0 {
			if _, err := w.Write(buf[:n]); err != nil {
				log.Fatal(err)
			}
		}
		if err == io.EOF {
			return
		}
	}
}

// Serve starts a secure echo server on the given listener.
func Serve(l net.Listener) error {
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go connect(conn)
	}
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	port := flag.Int("l", 0, "Listen mode. Specify port")
	flag.Parse()

	// Server mode
	if *port != 0 {
		l, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer l.Close()
		log.Fatal(Serve(l))
	}

	// Client mode
	if len(os.Args) != 3 {
		log.Fatalf("Usage: %s <port> <message>", os.Args[0])
	}
	conn, err := Dial("localhost:" + os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	if _, err := conn.Write([]byte(os.Args[2])); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, len(os.Args[2]))
	n, err := conn.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
}
<commit_msg>All tests working<commit_after>package main

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"golang.org/x/crypto/nacl/box"
	"golang.org/x/crypto/nacl/secretbox"
	"io"
	"log"
	"net"
	"os"
)

type header struct {
	Length uint64
	Nonce  [24]byte
}

type secureReader struct {
	src io.Reader
	key *[32]byte
	buf []byte
}

type secureWriter struct {
	dst io.Writer
	key *[32]byte
}

type secureConn struct {
	io.Reader
	io.Writer
	io.Closer
}

func generateNonce(nonce *[24]byte) error {
	_, err := io.ReadFull(rand.Reader, nonce[:])
	return err
}

func (p secureReader) Read(b []byte) (int, error) {
	if p.buf == nil {
		var h header
		if err := binary.Read(p.src, binary.LittleEndian, &h); err != nil {
			return 0, err
		}

		p.buf = make([]byte, h.Length-secretbox.Overhead)
		tmp := make([]byte, h.Length)

		if _, err := io.ReadFull(p.src, tmp); err != nil {
			return 0, err
		}

		_, auth := secretbox.Open(p.buf[:0], tmp[:], &h.Nonce, p.key)
		if !auth {
			return 0, fmt.Errorf("Message failed authentication")
		}
	}
	var size int
	if len(b) >= len(p.buf) {
		size = len(p.buf)
	} else {
		size = len(b)
	}

	for i := 0; i < size; i++ {
		b[i] = p.buf[i]
	}

	if size == len(p.buf) {
		p.buf = nil
	} else {
		p.buf = p.buf[size:]
	}
	return size, nil
}

func (p secureWriter) Write(b []byte) (int, error) {
	var h header
	h.Length = uint64(len(b) + secretbox.Overhead)
	if err := generateNonce(&h.Nonce); err != nil {
		return 0, err
	}

	hbuf := new(bytes.Buffer)
	binary.Write(hbuf, binary.LittleEndian, &h)
	if _, err := p.dst.Write(hbuf.Bytes()); err != nil {
		return 0, err
	}

	tmp := make([]byte, h.Length)
	secretbox.Seal(tmp[:0], b, &h.Nonce, p.key)
	return p.dst.Write(tmp)
}

// NewSecureReader instantiates a new SecureReader
func NewSecureReader(r io.Reader, priv, pub *[32]byte) io.Reader {
	sr := secureReader{src: r, key: new([32]byte), buf: nil}
	box.Precompute(sr.key, pub, priv)
	return sr
}

// NewSecureWriter instantiates a new SecureWriter
func NewSecureWriter(w io.Writer, priv, pub *[32]byte) io.Writer {
	sw := secureWriter{dst: w, key: new([32]byte)}
	box.Precompute(sw.key, pub, priv)
	return sw
}

func swapKeys(r io.Reader, w io.Writer) (priv, pub, peer *[32]byte, err error) {
	pub, priv, err = box.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, nil, err
	}

	sent := make(chan error)
	defer close(sent)

	go func() {
		_, err := w.Write(pub[:])
		sent <- err
	}()

	peer = new([32]byte)
	_, err = io.ReadFull(r, peer[:])
	if err != nil {
		return nil, nil, nil, err
	}

	if err = <-sent; err != nil {
		return nil, nil, nil, err
	}

	return
}

// Dial generates a private/public key pair,
// connects to the server, perform the handshake
// and return a reader/writer.
func Dial(addr string) (io.ReadWriteCloser, error) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		return nil, err
	}

	priv, _, servpub, err := swapKeys(conn, conn)
	if err != nil {
		return nil, err
	}

	secconn := secureConn{
		NewSecureReader(conn, priv, servpub),
		NewSecureWriter(conn, priv, servpub),
		conn}

	return secconn, nil
}

func connect(c net.Conn) {
	defer c.Close()

	priv, _, clientpub, err := swapKeys(c, c)
	if err != nil {
		log.Fatal(err)
	}

	r := NewSecureReader(c, priv, clientpub)
	w := NewSecureWriter(c, priv, clientpub)

	buf := make([]byte, 2048)
	for {
		n, err := r.Read(buf)
		if err != nil && err != io.EOF {
			log.Fatal(err)
		}
		if n > 0 {
			if _, err := w.Write(buf[:n]); err != nil {
				log.Fatal(err)
			}
		}
		if err == io.EOF {
			return
		}
	}
}

// Serve starts a secure echo server on the given listener.
func Serve(l net.Listener) error {
	for {
		conn, err := l.Accept()
		if err != nil {
			return err
		}
		go connect(conn)
	}
}

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	port := flag.Int("l", 0, "Listen mode. Specify port")
	flag.Parse()

	// Server mode
	if *port != 0 {
		l, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
		if err != nil {
			log.Fatal(err)
		}
		defer l.Close()
		log.Fatal(Serve(l))
	}

	// Client mode
	if len(os.Args) != 3 {
		log.Fatalf("Usage: %s <port> <message>", os.Args[0])
	}
	conn, err := Dial("localhost:" + os.Args[1])
	if err != nil {
		log.Fatal(err)
	}
	if _, err := conn.Write([]byte(os.Args[2])); err != nil {
		log.Fatal(err)
	}
	buf := make([]byte, len(os.Args[2]))
	n, err := conn.Read(buf)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("%s\n", buf[:n])
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"io"

	"path/filepath"

	"os/exec"

	"github.com/cep21/gobuild/internal/golang.org/x/net/context"
)

type gobuildMain struct {
	args []string

	flags struct {
		verbose     bool
		verboseFile string
		chunkSize   int
		forceAbs    bool
	}

	tc                templateCache
	storageDir        string
	testrunStorageDir string

	verboseLog logger
	errLog     logger

	stderr io.Writer

	onClose []func() error
}

var mainInstance = gobuildMain{
	tc: templateCache{
		cache: make(map[string]*buildTemplate),
	},
	stderr: os.Stdout,
}

func init() {
	flag.BoolVar(&mainInstance.flags.verbose, "verbose", false, "Add verbose log to stderr")
	flag.StringVar(&mainInstance.flags.verboseFile, "verbosefile", "", "Will verbose log to a filename rather than stderr")
	flag.IntVar(&mainInstance.flags.chunkSize, "chunksize", 250, "size to chunk xargs into")
	flag.BoolVar(&mainInstance.flags.forceAbs, "abs", false, "will force abs paths for ... dirs")
}

func main() {
	flag.Parse()
	mainInstance.args = flag.Args()
	if err := mainInstance.main(); err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err.Error())
		os.Exit(1)
	}
}

func (g *gobuildMain) parseFlags() error {
	vlog := ioutil.Discard
	if g.flags.verbose {
		vlog = os.Stderr
		if g.flags.verboseFile != "" {
			verboseFile, err := os.Create(g.flags.verboseFile)
			if err != nil {
				return wraperr(err, "cannot create verbose file %s", g.flags.verboseFile)
			}
			vlog = verboseFile
			g.onClose = append(g.onClose, verboseFile.Close)
		}
	}
	g.verboseLog = log.New(vlog, "[gobuild-verbose]", log.LstdFlags|log.Lshortfile)
	g.errLog = log.New(os.Stderr, "[gobuild-err]", log.LstdFlags|log.Lshortfile)
	g.tc.verboseLog = g.verboseLog

	var err error
	g.storageDir, err = g.storageDirectory()
	if err != nil {
		return wraperr(err, "cannot create test storage directory")
	}
	g.testrunStorageDir, err = g.testReportDirectory()
	if err != nil {
		return wraperr(err, "cannot create test run storage directory")
	}
	g.verboseLog.Printf("Storing results to %s", g.storageDir)

	return nil
}

func (g *gobuildMain) getArgs() (string, []string) {
	if len(g.args) == 0 {
		return "check", []string{"./..."}
	}
	if len(g.args) == 1 {
		return g.args[0], []string{"./..."}
	}
	return g.args[0], g.args[1:]
}

func (g *gobuildMain) fix(ctx context.Context, dirs []string) error {
	if err := g.install(ctx, dirs); err != nil {
		return wraperr(err, "cannot install subcommands")
	}
	c := fixCmd{
		dirs:       dirs,
		chunkSize:  g.flags.chunkSize,
		verboseOut: g.verboseLog,
		errOut:     g.errLog,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) lint(ctx context.Context, dirs []string) error {
	if err := g.install(ctx, dirs); err != nil {
		return wraperr(err, "cannot install subcommands")
	}
	testDirs, err := dirsWithFileGob(dirs, "*.go")
	if err != nil {
		return wraperr(err, "cannot find *.go files in dirs")
	}
	c := gometalinterCmd{
		verboseLog: g.verboseLog,
		errLog:     g.errLog,
		metaOutput: &myselfOutput{&nopCloseWriter{os.Stderr}},
		dirsToLint: testDirs,
		cache:      &g.tc,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) build(ctx context.Context, dirs []string) error {
	buildableDirs, err := dirsWithFileGob(dirs, "*.go")
	if err != nil {
		return wraperr(err, "cannot find *.go files in dirs")
	}
	c := cmdBuild{
		verboseLog: g.verboseLog,
		errorLog:   g.errLog,
		cmdStdout:  &myselfOutput{&nopCloseWriter{os.Stdout}},
		cmdStderr:  &myselfOutput{&nopCloseWriter{os.Stderr}},
		dirs:       buildableDirs,
		cache:      &g.tc,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) dupl(ctx context.Context, dirs []string) error {
	if err := g.install(ctx, dirs); err != nil {
		return wraperr(err, "cannot install subcommands")
	}
	tmpl, err := g.tc.loadInDir(".")
	if err != nil {
		return wraperr(err, "cannot load root dir template")
	}
	htmlOut, err := os.Create(filepath.Join(g.storageDir, "dupl.html"))
	if err != nil {
		return wraperr(err, "cannot create coverage html file")
	}
	c := duplCmd{
		verboseLog: g.verboseLog,
		dirs:       dirs,
		consoleOut: os.Stdout,
		htmlOut:    htmlOut,
		tmpl:       tmpl,
	}
	return multiErr([]error{c.Run(ctx), htmlOut.Close()})
}

func (g *gobuildMain) install(ctx context.Context, dirs []string) error {
	tmpl, err := g.tc.loadInDir(".")
	if err != nil {
		return wraperr(err, "cannot load root dir template")
	}
	c := installCmd{
		forceReinstall: false,
		verboseLog:     g.verboseLog,
		errLog:         g.errLog,
		stdoutOutput:   os.Stdout,
		stderrOutput:   os.Stderr,
		tmpl:           tmpl,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) testReportDirectory() (string, error) {
	return g.makeEnvDirectory("testReportEnv", "gobuild-tests")
}

func (g *gobuildMain) storageDirectory() (string, error) {
	return g.makeEnvDirectory("artifactsEnv", "gobuild")
}

func (g *gobuildMain) makeEnvDirectory(envName string, subdirName string) (string, error) {
	tmpl, err := g.tc.loadInDir(".")
	if err != nil {
		return "", wraperr(err, "cannot load root template directory")
	}
	fromEnv := os.Getenv(tmpl.varStr(envName))
	if fromEnv != "" {
		return fromEnv, nil
	}
	artifactDir := filepath.Join(os.TempDir(), subdirName)
	if err := os.RemoveAll(artifactDir); err != nil {
		return "", wraperr(err, "Cannot clean directory %s", artifactDir)
	}
	if err := os.MkdirAll(artifactDir, 0777); err != nil {
		return "", wraperr(err, "Cannot create directory %s", artifactDir)
	}

	return artifactDir, nil
}

func (g *gobuildMain) test(ctx context.Context, dirs []string) error {
	testDirs, err := dirsWithFileGob(dirs, "*.go")
	if err != nil {
		return wraperr(err, "cannot find *.go files in dirs")
	}

	fullCoverageFilename := filepath.Join(g.storageDir, "full_coverage_output.cover.txt")
	fullOut, err := os.Create(fullCoverageFilename)
	if err != nil {
		return wraperr(err, "cannot create full coverage profile file")
	}
	fullTestOutputFilename := filepath.Join(g.storageDir, "full_test_output.txt")
	fullTestStdout, err := os.Create(fullTestOutputFilename)
	if err != nil {
		return wraperr(err, "cannot create full test output file")
	}
	c := goCoverageCheck{
		dirs:                testDirs,
		cache:               &g.tc,
		coverProfileOutTo:   inDirStreamer(g.storageDir, ".cover.txt"),
		testStdoutOutputTo:  &myselfOutput{&nopCloseWriter{os.Stdout}},
		testStderrOutputTo:  &myselfOutput{&nopCloseWriter{os.Stderr}},
		requiredCoverage:    0,
		verboseLog:          g.verboseLog,
		errLog:              g.errLog,
		fullCoverageOutput:  fullOut,
		aggregateTestStdout: fullTestStdout,
	}
	e1 := c.Run(ctx)
	e2 := fullOut.Close()
	var e3 error
	if e2 == nil {
		htmlFilename := filepath.Join(g.storageDir, "full_coverage_output.cover.html")
		e3 = g.genCoverageHTML(ctx, fullCoverageFilename, htmlFilename)
	}
	e4 := fullTestStdout.Close()
	var e5 error
	if e4 == nil {
		e5 = g.genJunitXML(ctx, fullTestOutputFilename)
	}
	return multiErr([]error{e1, e2, e3, e4, e5})
}

func (g *gobuildMain) genJunitXML(ctx context.Context, fullTestOutputFilename string) error {
	junitXMLOutputFileDir := filepath.Join(g.testrunStorageDir, "gotest")
	if err := os.Mkdir(junitXMLOutputFileDir, 0777); err != nil {
		return wraperr(err, "cannot make temp dir %s", junitXMLOutputFileDir)
	}
	xmlOutFilename := filepath.Join(junitXMLOutputFileDir, "junit-gotest.xml")
	xmlOutFile, err := os.Create(xmlOutFilename)
	if err != nil {
		return wraperr(err, "cannot open xml run output file")
	}
	fullTestOutput, err := os.Open(fullTestOutputFilename)
	if err != nil {
		return wraperr(err, "cannot open test file for output")
	}
	return multiErr([]error{g.genJunitXMLFromBuffers(ctx, fullTestOutput, xmlOutFile), xmlOutFile.Close(), fullTestOutput.Close()})
}

func (g *gobuildMain) genJunitXMLFromBuffers(ctx context.Context, testInput io.Reader, testOutput io.Writer) error {
	cmd := exec.Command("go-junit-report")
	cmd.Stdin = testInput
	cmd.Stderr = os.Stderr
	cmd.Stdout = testOutput
	g.verboseLog.Printf("Generating junit XML with %v", cmd)
	if err := cmd.Run(); err != nil {
		return wraperr(err, "test junit XML generation failed")
	}
	return nil
}

func (g *gobuildMain) genCoverageHTML(ctx context.Context, coverFilename string, htmlFilename string) error {
	cmd := exec.Command("go", "tool", "cover", "-o", htmlFilename, "-html", coverFilename)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	g.verboseLog.Printf("Generating coverage html %s => %s with %v", coverFilename, htmlFilename, cmd)
	if err := cmd.Run(); err != nil {
		return wraperr(err, "coverage HTML generation failed")
	}
	return nil
}

func (g *gobuildMain) check(ctx context.Context, dirs []string) error {
	buildErr := g.build(ctx, dirs)
	lintErr := g.lint(ctx, dirs)
	duplErr := g.dupl(ctx, dirs)
	testErr := g.test(ctx, dirs)
	return multiErr([]error{buildErr, lintErr, duplErr, testErr})
}

func (g *gobuildMain) list(ctx context.Context, dirs []string) error {
	g.verboseLog.Printf("len(dirs) = %d", len(dirs))
	fmt.Printf("%s\n", strings.Join(dirs, "\n"))
	return nil
}

func (g *gobuildMain) Close() error {
	errs := make([]error, 0, len(g.onClose))
	for _, f := range g.onClose {
		if err := f(); err != nil {
			errs = append(errs, err)
		}
	}
	return multiErr(errs)
}

func (g *gobuildMain) main() error {
	defer func() {
		if err := g.Close(); err != nil {
			fmt.Fprintf(os.Stderr, "Cannot close mainInstance: %s\n", err.Error())
		}
	}()

	if err := g.parseFlags(); err != nil {
		return wraperr(err, "cannot parse flags")
	}
	ctx := context.Background()

	pe := pathExpansion{
		forceAbs: g.flags.forceAbs,
		log:      g.verboseLog,
		template: &g.tc,
	}

	cmdMap := map[string]func(context.Context, []string) error{
		"fix":     g.fix,
		"lint":    g.lint,
		"list":    g.list,
		"build":   g.build,
		"test":    g.test,
		"dupl":    g.dupl,
		"install": g.install,
		"check":   g.check,
	}

	cmd, args := g.getArgs()
	f, exists := cmdMap[cmd]
	if !exists {
		fmt.Fprintf(g.stderr, "Unknown command %s\nValid commands:\n", cmd)
		for k := range cmdMap {
			fmt.Fprintf(g.stderr, "  %s\n", k)
		}
		return fmt.Errorf("unknown command %s", cmd)
	}
	dirs, err := pe.expandPaths(args)
	if err != nil {
		return wraperr(err, "cannot expand paths %s", strings.Join(args, ","))
	}
	if err := f(ctx, dirs); err != nil {
		return wraperr(err, "Failure in command %s", cmd)
	}
	return nil
}
<commit_msg>Allow filename prefixes<commit_after>package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strings"

	"io"

	"path/filepath"

	"os/exec"

	"github.com/cep21/gobuild/internal/golang.org/x/net/context"
)

type gobuildMain struct {
	args []string

	flags struct {
		verbose        bool
		verboseFile    string
		chunkSize      int
		forceAbs       bool
		filenamePrefix string
	}

	tc                templateCache
	storageDir        string
	testrunStorageDir string

	verboseLog logger
	errLog     logger

	stderr io.Writer

	onClose []func() error
}

var mainInstance = gobuildMain{
	tc: templateCache{
		cache: make(map[string]*buildTemplate),
	},
	stderr: os.Stdout,
}

func init() {
	flag.BoolVar(&mainInstance.flags.verbose, "verbose", false, "Add verbose log to stderr")
	flag.StringVar(&mainInstance.flags.verboseFile, "verbosefile", "", "Will verbose log to a filename rather than stderr")
	flag.IntVar(&mainInstance.flags.chunkSize, "chunksize", 250, "size to chunk xargs into")
	flag.StringVar(&mainInstance.flags.filenamePrefix, "filename_prefix", "", "Prefix to append to all generated files")
	flag.BoolVar(&mainInstance.flags.forceAbs, "abs", false, "will force abs paths for ... dirs")
}

func main() {
	flag.Parse()
	mainInstance.args = flag.Args()
	if err := mainInstance.main(); err != nil {
		fmt.Fprintf(os.Stderr, "%s\n", err.Error())
		os.Exit(1)
	}
}

func (g *gobuildMain) parseFlags() error {
	vlog := ioutil.Discard
	if g.flags.verbose {
		vlog = os.Stderr
		if g.flags.verboseFile != "" {
			verboseFile, err := os.Create(g.flags.verboseFile)
			if err != nil {
				return wraperr(err, "cannot create verbose file %s", g.flags.verboseFile)
			}
			vlog = verboseFile
			g.onClose = append(g.onClose, verboseFile.Close)
		}
	}
	g.verboseLog = log.New(vlog, "[gobuild-verbose]", log.LstdFlags|log.Lshortfile)
	g.errLog = log.New(os.Stderr, "[gobuild-err]", log.LstdFlags|log.Lshortfile)
	g.tc.verboseLog = g.verboseLog

	var err error
	g.storageDir, err = g.storageDirectory()
	if err != nil {
		return wraperr(err, "cannot create test storage directory")
	}
	g.testrunStorageDir, err = g.testReportDirectory()
	if err != nil {
		return wraperr(err, "cannot create test run storage directory")
	}
	g.verboseLog.Printf("Storing results to %s", g.storageDir)

	return nil
}

func (g *gobuildMain) getArgs() (string, []string) {
	if len(g.args) == 0 {
		return "check", []string{"./..."}
	}
	if len(g.args) == 1 {
		return g.args[0], []string{"./..."}
	}
	return g.args[0], g.args[1:]
}

func (g *gobuildMain) fix(ctx context.Context, dirs []string) error {
	if err := g.install(ctx, dirs); err != nil {
		return wraperr(err, "cannot install subcommands")
	}
	c := fixCmd{
		dirs:       dirs,
		chunkSize:  g.flags.chunkSize,
		verboseOut: g.verboseLog,
		errOut:     g.errLog,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) lint(ctx context.Context, dirs []string) error {
	if err := g.install(ctx, dirs); err != nil {
		return wraperr(err, "cannot install subcommands")
	}
	testDirs, err := dirsWithFileGob(dirs, "*.go")
	if err != nil {
		return wraperr(err, "cannot find *.go files in dirs")
	}
	c := gometalinterCmd{
		verboseLog: g.verboseLog,
		errLog:     g.errLog,
		metaOutput: &myselfOutput{&nopCloseWriter{os.Stderr}},
		dirsToLint: testDirs,
		cache:      &g.tc,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) build(ctx context.Context, dirs []string) error {
	buildableDirs, err := dirsWithFileGob(dirs, "*.go")
	if err != nil {
		return wraperr(err, "cannot find *.go files in dirs")
	}
	c := cmdBuild{
		verboseLog: g.verboseLog,
		errorLog:   g.errLog,
		cmdStdout:  &myselfOutput{&nopCloseWriter{os.Stdout}},
		cmdStderr:  &myselfOutput{&nopCloseWriter{os.Stderr}},
		dirs:       buildableDirs,
		cache:      &g.tc,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) dupl(ctx context.Context, dirs []string) error {
	if err := g.install(ctx, dirs); err != nil {
		return wraperr(err, "cannot install subcommands")
	}
	tmpl, err := g.tc.loadInDir(".")
	if err != nil {
		return wraperr(err, "cannot load root dir template")
	}
	htmlOut, err := os.Create(filepath.Join(g.storageDir, g.flags.filenamePrefix+"dupl.html"))
	if err != nil {
		return wraperr(err, "cannot create coverage html file")
	}
	c := duplCmd{
		verboseLog: g.verboseLog,
		dirs:       dirs,
		consoleOut: os.Stdout,
		htmlOut:    htmlOut,
		tmpl:       tmpl,
	}
	return multiErr([]error{c.Run(ctx), htmlOut.Close()})
}

func (g *gobuildMain) install(ctx context.Context, dirs []string) error {
	tmpl, err := g.tc.loadInDir(".")
	if err != nil {
		return wraperr(err, "cannot load root dir template")
	}
	c := installCmd{
		forceReinstall: false,
		verboseLog:     g.verboseLog,
		errLog:         g.errLog,
		stdoutOutput:   os.Stdout,
		stderrOutput:   os.Stderr,
		tmpl:           tmpl,
	}
	return c.Run(ctx)
}

func (g *gobuildMain) testReportDirectory() (string, error) {
	return g.makeEnvDirectory("testReportEnv", "gobuild-tests")
}

func (g *gobuildMain) storageDirectory() (string, error) {
	return g.makeEnvDirectory("artifactsEnv", "gobuild")
}

func (g *gobuildMain) makeEnvDirectory(envName string, subdirName string) (string, error) {
	tmpl, err := g.tc.loadInDir(".")
	if err != nil {
		return "", wraperr(err, "cannot load root template directory")
	}
	fromEnv := os.Getenv(tmpl.varStr(envName))
	if fromEnv != "" {
		return fromEnv, nil
	}
	artifactDir := filepath.Join(os.TempDir(), subdirName)
	if err := os.RemoveAll(artifactDir); err != nil {
		return "", wraperr(err, "Cannot clean directory %s", artifactDir)
	}
	if err := os.MkdirAll(artifactDir, 0777); err != nil {
		return "", wraperr(err, "Cannot create directory %s", artifactDir)
	}

	return artifactDir, nil
}

func (g *gobuildMain) test(ctx context.Context, dirs []string) error {
	testDirs, err := dirsWithFileGob(dirs, "*.go")
	if err != nil {
		return wraperr(err, "cannot find *.go files in dirs")
	}

	fullCoverageFilename := filepath.Join(g.storageDir, g.flags.filenamePrefix+"full_coverage_output.cover.txt")
	fullOut, err := os.Create(fullCoverageFilename)
	if err != nil {
		return wraperr(err, "cannot create full coverage profile file")
	}
	fullTestOutputFilename := filepath.Join(g.storageDir, g.flags.filenamePrefix+"full_test_output.txt")
	fullTestStdout, err := os.Create(fullTestOutputFilename)
	if err != nil {
		return wraperr(err, "cannot create full test output file")
	}
	c := goCoverageCheck{
		dirs:                testDirs,
		cache:               &g.tc,
		coverProfileOutTo:   inDirStreamer(g.storageDir, ".cover.txt"),
		testStdoutOutputTo:  &myselfOutput{&nopCloseWriter{os.Stdout}},
		testStderrOutputTo:  &myselfOutput{&nopCloseWriter{os.Stderr}},
		requiredCoverage:    0,
		verboseLog:          g.verboseLog,
		errLog:              g.errLog,
		fullCoverageOutput:  fullOut,
		aggregateTestStdout: fullTestStdout,
	}
	e1 := c.Run(ctx)
	e2 := fullOut.Close()
	var e3 error
	if e2 == nil {
		htmlFilename := filepath.Join(g.storageDir, g.flags.filenamePrefix+"full_coverage_output.cover.html")
		e3 = g.genCoverageHTML(ctx, fullCoverageFilename, htmlFilename)
	}
	e4 := fullTestStdout.Close()
	var e5 error
	if e4 == nil {
		e5 = g.genJunitXML(ctx, fullTestOutputFilename)
	}
	return multiErr([]error{e1, e2, e3, e4, e5})
}

func (g *gobuildMain) genJunitXML(ctx context.Context, fullTestOutputFilename string) error {
	junitXMLOutputFileDir := filepath.Join(g.testrunStorageDir, "gotest")
	if _, err := os.Stat(junitXMLOutputFileDir); err != nil {
		if err := os.Mkdir(junitXMLOutputFileDir, 0777); err != nil {
			return wraperr(err, "cannot make temp dir %s", junitXMLOutputFileDir)
		}
	}
	xmlOutFilename := filepath.Join(junitXMLOutputFileDir, g.flags.filenamePrefix+"junit-gotest.xml")
	xmlOutFile, err := os.Create(xmlOutFilename)
	if err != nil {
		return wraperr(err, "cannot open xml run output file")
	}
	fullTestOutput, err := os.Open(fullTestOutputFilename)
	if err != nil {
		return wraperr(err, "cannot open test file for output")
	}
	return multiErr([]error{g.genJunitXMLFromBuffers(ctx, fullTestOutput, xmlOutFile), xmlOutFile.Close(), fullTestOutput.Close()})
}

func (g *gobuildMain) genJunitXMLFromBuffers(ctx context.Context, testInput io.Reader, testOutput io.Writer) error {
	cmd := exec.Command("go-junit-report")
	cmd.Stdin = testInput
	cmd.Stderr = os.Stderr
	cmd.Stdout = testOutput
	g.verboseLog.Printf("Generating junit XML with %v", cmd)
	if err := cmd.Run(); err != nil {
		return wraperr(err, "test junit XML generation failed")
	}
	return nil
}

func (g *gobuildMain) genCoverageHTML(ctx context.Context, coverFilename string, htmlFilename string) error {
	cmd := exec.Command("go", "tool", "cover", "-o", htmlFilename, "-html", coverFilename)
	cmd.Stderr = os.Stderr
	cmd.Stdout = os.Stdout
	g.verboseLog.Printf("Generating coverage html %s => %s with %v", coverFilename, htmlFilename, cmd)
	if err := cmd.Run(); err != nil {
		return wraperr(err, "coverage HTML generation failed")
	}
	return nil
}

func (g *gobuildMain) check(ctx context.Context, dirs []string) error {
	buildErr := g.build(ctx, dirs)
	lintErr := g.lint(ctx, dirs)
	duplErr := g.dupl(ctx, dirs)
	testErr := g.test(ctx, dirs)
	return multiErr([]error{buildErr, lintErr, duplErr, testErr})
}

func (g *gobuildMain) list(ctx context.Context, dirs []string) error {
	g.verboseLog.Printf("len(dirs) = %d", len(dirs))
	fmt.Printf("%s\n", strings.Join(dirs, "\n"))
	return nil
}

func (g *gobuildMain) Close() error {
	errs := make([]error, 0, len(g.onClose))
	for _, f := range g.onClose {
		if err := f(); err != nil {
			errs = append(errs, err)
		}
	}
	return multiErr(errs)
}

func (g *gobuildMain) main() error {
	defer func() {
		if err := g.Close(); err != nil {
			fmt.Fprintf(os.Stderr, "Cannot close mainInstance: %s\n", err.Error())
		}
	}()

	if err := g.parseFlags(); err != nil {
		return wraperr(err, "cannot parse flags")
	}
	ctx := context.Background()

	pe := pathExpansion{
		forceAbs: g.flags.forceAbs,
		log:      g.verboseLog,
		template: &g.tc,
	}

	cmdMap := map[string]func(context.Context, []string) error{
		"fix":     g.fix,
		"lint":    g.lint,
		"list":    g.list,
		"build":   g.build,
		"test":    g.test,
		"dupl":    g.dupl,
		"install": g.install,
		"check":   g.check,
	}

	cmd, args := g.getArgs()
	f, exists := cmdMap[cmd]
	if !exists {
		fmt.Fprintf(g.stderr, "Unknown command %s\nValid commands:\n", cmd)
		for k := range cmdMap {
			fmt.Fprintf(g.stderr, "  %s\n", k)
		}
		return fmt.Errorf("unknown command %s", cmd)
	}
	dirs, err := pe.expandPaths(args)
	if err != nil {
		return wraperr(err, "cannot expand paths %s", strings.Join(args, ","))
	}
	if err := f(ctx, dirs); err != nil {
		return wraperr(err, "Failure in command %s", cmd)
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 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.

// Benchstat computes and compares statistics about benchmarks.
//
// Usage:
//
//	benchstat old.txt [new.txt]
//
// Each input file should contain the concatenated output of a number
// of runs of ``go test -bench.'' For each different benchmark listed in an input file,
// benchstat computes the mean, minimum, and maximum run time,
// after removing outliers using the interquartile range rule.
//
// If invoked on a single input file, benchstat prints the per-benchmark statistics
// for that file.
//
// If invoked on a pair of input files, benchstat adds to the output a column
// showing the statistics from the second file and a column showing the
// percent change in mean from the first to the second file.
// Next to the percent change, benchstat shows the p value from the
// two-sample Welch t-test, which measures the statistical significance
// of the difference. If the t-test indicates that the measured difference is
// not statistically significant (defined as p > 0.05), benchstat replaces
// the percent change with a single ~.
//
// Example
//
// Suppose we collect benchmark results from running ``go test -bench=Encode''
// five times before and after a particular change.
//
// The file old.txt contains:
//
//	BenchmarkGobEncode   	100	  13552735 ns/op	  56.63 MB/s
//	BenchmarkJSONEncode  	 50	  32395067 ns/op	  59.90 MB/s
//	BenchmarkGobEncode   	100	  13553943 ns/op	  56.63 MB/s
//	BenchmarkJSONEncode  	 50	  32334214 ns/op	  60.01 MB/s
//	BenchmarkGobEncode   	100	  13606356 ns/op	  56.41 MB/s
//	BenchmarkJSONEncode  	 50	  31992891 ns/op	  60.65 MB/s
//	BenchmarkGobEncode   	100	  13683198 ns/op	  56.09 MB/s
//	BenchmarkJSONEncode  	 50	  31735022 ns/op	  61.15 MB/s
//
// The file new.txt contains:
//
//	BenchmarkGobEncode   	 100	  11773189 ns/op	  65.19 MB/s
//	BenchmarkJSONEncode  	  50	  32036529 ns/op	  60.57 MB/s
//	BenchmarkGobEncode   	 100	  11942588 ns/op	  64.27 MB/s
//	BenchmarkJSONEncode  	  50	  32156552 ns/op	  60.34 MB/s
//	BenchmarkGobEncode   	 100	  11786159 ns/op	  65.12 MB/s
//	BenchmarkJSONEncode  	  50	  31288355 ns/op	  62.02 MB/s
//	BenchmarkGobEncode   	 100	  11628583 ns/op	  66.00 MB/s
//	BenchmarkJSONEncode  	  50	  31559706 ns/op	  61.49 MB/s
//	BenchmarkGobEncode   	 100	  11815924 ns/op	  64.96 MB/s
//	BenchmarkJSONEncode  	  50	  31765634 ns/op	  61.09 MB/s
//
// The order of the lines in the file does not matter, except that the
// output lists benchmarks in order of appearance.
//
// If run with just one input file, benchstat summarizes that file:
//
//	$ benchstat old.txt
//	name        mean
//	GobEncode   13.6ms × (1.00,1.01)
//	JSONEncode  32.1ms × (0.99,1.01)
//	$
//
// If run with two input files, benchstat summarizes and compares:
//
//	$ benchstat old.txt new.txt
//	name        old mean              new mean              delta
//	GobEncode   13.6ms × (1.00,1.01)  11.8ms × (0.99,1.01)  -13.31% (p=0.000)
//	JSONEncode  32.1ms × (0.99,1.01)  31.8ms × (0.99,1.01)     ~    (p=0.154)
//	$
//
// Note that the JSONEncode result is reported as
// statistically insignificant instead of a -1% delta.
//
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"unicode/utf8"

	"rsc.io/benchstat/internal/stats"
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: benchstat old.txt [new.txt]\n")
	os.Exit(2)
}

func main() {
	log.SetPrefix("benchstats: ")
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() < 1 || flag.NArg() > 2 {
		flag.Usage()
	}

	before := readFile(flag.Arg(0))
	var out [][]string
	if flag.NArg() == 1 {
		out = append(out, []string{"name", "mean"})
		for _, old := range before.Benchmarks {
			out = append(out, []string{old.Name, old.Time(old.Scaler())})
		}
	} else {
		after := readFile(flag.Arg(1))
		out = append(out, []string{"name", "old mean            ", "new mean            ", "delta"})
		for _, old := range before.Benchmarks {
			new := after.ByName[old.Name]
			if new == nil {
				continue
			}

			ttest, err := stats.TwoSampleWelchTTest(stats.Sample{Xs: old.RTimes}, stats.Sample{Xs: new.RTimes})
			significant := false
			if err == nil {
				significant = ttest.P <= 0.05
			}

			scaler := old.Scaler()
			row := []string{old.Name, old.Time(scaler), new.Time(scaler), "~   "}
			if err != nil {
				row[3] = fmt.Sprintf("(%s)", err)
			} else if significant {
				row[3] = fmt.Sprintf("%+.2f%%", ((new.Mean/old.Mean)-1.0)*100.0)
			}
			if ttest != nil {
				row[3] += fmt.Sprintf(" (p=%0.3f)", ttest.P)
			}
			out = append(out, row)
		}
	}

	max := []int{0, 0, 0, 0}
	for _, row := range out {
		for i, s := range row {
			n := utf8.RuneCountInString(s)
			if max[i] < n {
				max[i] = n
			}
		}
	}

	if flag.NArg() == 1 {
		row := out[0]
		fmt.Printf("%-*s  %s\n", max[0], row[0], row[1])
		for _, row := range out[1:] {
			fmt.Printf("%-*s  %*s\n", max[0], row[0], max[1], row[1])
		}
	} else {
		row := out[0]
		fmt.Printf("%-*s  %*s  %*s  %s\n", max[0], row[0], max[1], row[1], max[2], row[2], row[3])
		for _, row := range out[1:] {
			fmt.Printf("%-*s  %*s  %*s  %*s\n", max[0], row[0], max[1], row[1], max[2], row[2], max[3], row[3])
		}
	}
}

func (b *Benchmark) Scaler() func(*Benchmark) string {
	var format string
	var scale float64
	switch x := b.Mean / 1e9; {
	case x >= 99.5:
		format, scale = "%.0fs", 1
	case x >= 9.95:
		format, scale = "%.1fs", 1
	case x >= 0.995:
		format, scale = "%.2fs", 1
	case x >= 0.0995:
		format, scale = "%.0fms", 1000
	case x >= 0.00995:
		format, scale = "%.1fms", 1000
	case x >= 0.000995:
		format, scale = "%.2fms", 1000
	case x >= 0.0000995:
		format, scale = "%.0fµs", 1000*1000
	case x >= 0.00000995:
		format, scale = "%.1fµs", 1000*1000
	case x >= 0.000000995:
		format, scale = "%.2fµs", 1000*1000
	case x >= 0.0000000995:
		format, scale = "%.0fns", 1000*1000*1000
	case x >= 0.00000000995:
		format, scale = "%.1fns", 1000*1000*1000
	default:
		format, scale = "%.2fns", 1000*1000*1000
	}
	return func(b *Benchmark) string {
		return fmt.Sprintf(format, b.Mean/1e9*scale)
	}
}

func (b *Benchmark) Time(scaler func(*Benchmark) string) string {
	return fmt.Sprintf("%s × (%.2f,%.2f)", scaler(b), b.Min/b.Mean, b.Max/b.Mean)
}

// A Group is a collection of benchmark results for a specific version of the code.
type Group struct {
	Benchmarks []*Benchmark
	ByName     map[string]*Benchmark
}

// A Benchmark is a collection of runs of a specific benchmark.
type Benchmark struct {
	Name   string
	Runs   []Run
	Times  []float64 // Times of runs in nanoseconds
	RTimes []float64 // Times with outliers removed
	Min    float64   // minimum run
	Mean   float64
	Max    float64 // maximum run
}

// A Run records the result of a single benchmark run.
type Run struct {
	N  int
	Ns float64
}

// readFile reads a benchmark group from a file.
func readFile(file string) *Group {
	g := &Group{ByName: make(map[string]*Benchmark)}

	text, err := ioutil.ReadFile(file)
	if err != nil {
		log.Fatal(err)
	}
	for _, line := range strings.Split(string(text), "\n") {
		f := strings.Fields(line)
		if len(f) < 4 {
			continue
		}
		name := f[0]
		if !strings.HasPrefix(name, "Benchmark") {
			continue
		}
		name = strings.TrimPrefix(name, "Benchmark")
		n, _ := strconv.Atoi(f[1])
		var ns float64
		for i := 2; i+2 <= len(f); i += 2 {
			if f[i+1] == "ns/op" {
				ns, _ = strconv.ParseFloat(f[i], 64)
				break
			}
		}
		if n == 0 || ns == 0 {
			continue
		}
		b := g.ByName[name]
		if b == nil {
			b = &Benchmark{Name: name}
			g.ByName[name] = b
			g.Benchmarks = append(g.Benchmarks, b)
		}
		b.Runs = append(b.Runs, Run{n, ns})
		b.Times = append(b.Times, ns)
	}

	for _, b := range g.Benchmarks {
		// Discard outliers.
		times := stats.Sample{Xs: b.Times}
		q1, q3 := times.Percentile(0.25), times.Percentile(0.75)
		lo, hi := q1-1.5*(q3-q1), q3+1.5*(q3-q1)
		for _, time := range b.Times {
			if lo <= time && time <= hi {
				b.RTimes = append(b.RTimes, time)
			}
		}

		// Compute statistics of remaining data.
		b.Min, b.Max = stats.Bounds(b.RTimes)
		b.Mean = stats.Mean(b.RTimes)
	}

	return g
}
<commit_msg>benchstat: shorter messages for common stats errors<commit_after>// Copyright 2015 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.

// Benchstat computes and compares statistics about benchmarks.
//
// Usage:
//
//	benchstat old.txt [new.txt]
//
// Each input file should contain the concatenated output of a number
// of runs of ``go test -bench.'' For each different benchmark listed in an input file,
// benchstat computes the mean, minimum, and maximum run time,
// after removing outliers using the interquartile range rule.
//
// If invoked on a single input file, benchstat prints the per-benchmark statistics
// for that file.
//
// If invoked on a pair of input files, benchstat adds to the output a column
// showing the statistics from the second file and a column showing the
// percent change in mean from the first to the second file.
// Next to the percent change, benchstat shows the p value from the
// two-sample Welch t-test, which measures the statistical significance
// of the difference. If the t-test indicates that the measured difference is
// not statistically significant (defined as p > 0.05), benchstat replaces
// the percent change with a single ~.
//
// Example
//
// Suppose we collect benchmark results from running ``go test -bench=Encode''
// five times before and after a particular change.
//
// The file old.txt contains:
//
//	BenchmarkGobEncode   	100	  13552735 ns/op	  56.63 MB/s
//	BenchmarkJSONEncode  	 50	  32395067 ns/op	  59.90 MB/s
//	BenchmarkGobEncode   	100	  13553943 ns/op	  56.63 MB/s
//	BenchmarkJSONEncode  	 50	  32334214 ns/op	  60.01 MB/s
//	BenchmarkGobEncode   	100	  13606356 ns/op	  56.41 MB/s
//	BenchmarkJSONEncode  	 50	  31992891 ns/op	  60.65 MB/s
//	BenchmarkGobEncode   	100	  13683198 ns/op	  56.09 MB/s
//	BenchmarkJSONEncode  	 50	  31735022 ns/op	  61.15 MB/s
//
// The file new.txt contains:
//
//	BenchmarkGobEncode   	 100	  11773189 ns/op	  65.19 MB/s
//	BenchmarkJSONEncode  	  50	  32036529 ns/op	  60.57 MB/s
//	BenchmarkGobEncode   	 100	  11942588 ns/op	  64.27 MB/s
//	BenchmarkJSONEncode  	  50	  32156552 ns/op	  60.34 MB/s
//	BenchmarkGobEncode   	 100	  11786159 ns/op	  65.12 MB/s
//	BenchmarkJSONEncode  	  50	  31288355 ns/op	  62.02 MB/s
//	BenchmarkGobEncode   	 100	  11628583 ns/op	  66.00 MB/s
//	BenchmarkJSONEncode  	  50	  31559706 ns/op	  61.49 MB/s
//	BenchmarkGobEncode   	 100	  11815924 ns/op	  64.96 MB/s
//	BenchmarkJSONEncode  	  50	  31765634 ns/op	  61.09 MB/s
//
// The order of the lines in the file does not matter, except that the
// output lists benchmarks in order of appearance.
//
// If run with just one input file, benchstat summarizes that file:
//
//	$ benchstat old.txt
//	name        mean
//	GobEncode   13.6ms × (1.00,1.01)
//	JSONEncode  32.1ms × (0.99,1.01)
//	$
//
// If run with two input files, benchstat summarizes and compares:
//
//	$ benchstat old.txt new.txt
//	name        old mean              new mean              delta
//	GobEncode   13.6ms × (1.00,1.01)  11.8ms × (0.99,1.01)  -13.31% (p=0.000)
//	JSONEncode  32.1ms × (0.99,1.01)  31.8ms × (0.99,1.01)     ~    (p=0.154)
//	$
//
// Note that the JSONEncode result is reported as
// statistically insignificant instead of a -1% delta.
//
package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"strconv"
	"strings"
	"unicode/utf8"

	"rsc.io/benchstat/internal/stats"
)

func usage() {
	fmt.Fprintf(os.Stderr, "usage: benchstat old.txt [new.txt]\n")
	os.Exit(2)
}

func main() {
	log.SetPrefix("benchstats: ")
	log.SetFlags(0)
	flag.Usage = usage
	flag.Parse()
	if flag.NArg() < 1 || flag.NArg() > 2 {
		flag.Usage()
	}

	before := readFile(flag.Arg(0))
	var out [][]string
	if flag.NArg() == 1 {
		out = append(out, []string{"name", "mean"})
		for _, old := range before.Benchmarks {
			out = append(out, []string{old.Name, old.Time(old.Scaler())})
		}
	} else {
		after := readFile(flag.Arg(1))
		out = append(out, []string{"name", "old mean            ", "new mean            ", "delta"})
		for _, old := range before.Benchmarks {
			new := after.ByName[old.Name]
			if new == nil {
				continue
			}

			ttest, err := stats.TwoSampleWelchTTest(stats.Sample{Xs: old.RTimes}, stats.Sample{Xs: new.RTimes})
			significant := false
			if err == nil {
				significant = ttest.P <= 0.05
			}

			scaler := old.Scaler()
			row := []string{old.Name, old.Time(scaler), new.Time(scaler), "~   "}
			if err == stats.ErrZeroVariance {
				row[3] = "zero variance"
			} else if err == stats.ErrSampleSize {
				row[3] = "too few samples"
			} else if err != nil {
				row[3] = fmt.Sprintf("(%s)", err)
			} else if significant {
				row[3] = fmt.Sprintf("%+.2f%%", ((new.Mean/old.Mean)-1.0)*100.0)
			}
			if ttest != nil {
				row[3] += fmt.Sprintf(" (p=%0.3f)", ttest.P)
			}
			out = append(out, row)
		}
	}

	max := []int{0, 0, 0, 0}
	for _, row := range out {
		for i, s := range row {
			n := utf8.RuneCountInString(s)
			if max[i] < n {
				max[i] = n
			}
		}
	}

	if flag.NArg() == 1 {
		row := out[0]
		fmt.Printf("%-*s  %s\n", max[0], row[0], row[1])
		for _, row := range out[1:] {
			fmt.Printf("%-*s  %*s\n", max[0], row[0], max[1], row[1])
		}
	} else {
		row := out[0]
		fmt.Printf("%-*s  %*s  %*s  %s\n", max[0], row[0], max[1], row[1], max[2], row[2], row[3])
		for _, row := range out[1:] {
			fmt.Printf("%-*s  %*s  %*s  %*s\n", max[0], row[0], max[1], row[1], max[2], row[2], max[3], row[3])
		}
	}
}

func (b *Benchmark) Scaler() func(*Benchmark) string {
	var format string
	var scale float64
	switch x := b.Mean / 1e9; {
	case x >= 99.5:
		format, scale = "%.0fs", 1
	case x >= 9.95:
		format, scale = "%.1fs", 1
	case x >= 0.995:
		format, scale = "%.2fs", 1
	case x >= 0.0995:
		format, scale = "%.0fms", 1000
	case x >= 0.00995:
		format, scale = "%.1fms", 1000
	case x >= 0.000995:
		format, scale = "%.2fms", 1000
	case x >= 0.0000995:
		format, scale = "%.0fµs", 1000*1000
	case x >= 0.00000995:
		format, scale = "%.1fµs", 1000*1000
	case x >= 0.000000995:
		format, scale = "%.2fµs", 1000*1000
	case x >= 0.0000000995:
		format, scale = "%.0fns", 1000*1000*1000
	case x >= 0.00000000995:
		format, scale = "%.1fns", 1000*1000*1000
	default:
		format, scale = "%.2fns", 1000*1000*1000
	}
	return func(b *Benchmark) string {
		return fmt.Sprintf(format, b.Mean/1e9*scale)
	}
}

func (b *Benchmark) Time(scaler func(*Benchmark) string) string {
	return fmt.Sprintf("%s × (%.2f,%.2f)", scaler(b), b.Min/b.Mean, b.Max/b.Mean)
}

// A Group is a collection of benchmark results for a specific version of the code.
type Group struct {
	Benchmarks []*Benchmark
	ByName     map[string]*Benchmark
}

// A Benchmark is a collection of runs of a specific benchmark.
type Benchmark struct {
	Name   string
	Runs   []Run
	Times  []float64 // Times of runs in nanoseconds
	RTimes []float64 // Times with outliers removed
	Min    float64   // minimum run
	Mean   float64
	Max    float64 // maximum run
}

// A Run records the result of a single benchmark run.
type Run struct {
	N  int
	Ns float64
}

// readFile reads a benchmark group from a file.
func readFile(file string) *Group {
	g := &Group{ByName: make(map[string]*Benchmark)}

	text, err := ioutil.ReadFile(file)
	if err != nil {
		log.Fatal(err)
	}
	for _, line := range strings.Split(string(text), "\n") {
		f := strings.Fields(line)
		if len(f) < 4 {
			continue
		}
		name := f[0]
		if !strings.HasPrefix(name, "Benchmark") {
			continue
		}
		name = strings.TrimPrefix(name, "Benchmark")
		n, _ := strconv.Atoi(f[1])
		var ns float64
		for i := 2; i+2 <= len(f); i += 2 {
			if f[i+1] == "ns/op" {
				ns, _ = strconv.ParseFloat(f[i], 64)
				break
			}
		}
		if n == 0 || ns == 0 {
			continue
		}
		b := g.ByName[name]
		if b == nil {
			b = &Benchmark{Name: name}
			g.ByName[name] = b
			g.Benchmarks = append(g.Benchmarks, b)
		}
		b.Runs = append(b.Runs, Run{n, ns})
		b.Times = append(b.Times, ns)
	}

	for _, b := range g.Benchmarks {
		// Discard outliers.
		times := stats.Sample{Xs: b.Times}
		q1, q3 := times.Percentile(0.25), times.Percentile(0.75)
		lo, hi := q1-1.5*(q3-q1), q3+1.5*(q3-q1)
		for _, time := range b.Times {
			if lo <= time && time <= hi {
				b.RTimes = append(b.RTimes, time)
			}
		}

		// Compute statistics of remaining data.
		b.Min, b.Max = stats.Bounds(b.RTimes)
		b.Mean = stats.Mean(b.RTimes)
	}

	return g
}
<|endoftext|>