<commit_before>package postgres

import (
	"database/sql"
	"fmt"
	. "github.com/aktau/gomig/db/common"
	"log"
	"strings"
)

var PG_W_VERBOSE = true

var (
	postgresInit = []string{
		"SET client_encoding = 'UTF8';",
		"SET standard_conforming_strings = off;",
		"SET check_function_bodies = false;",
		"SET client_min_messages = warning;",
	}
)

const (
	explainQuery = `
SELECT col.column_name AS field,
       CASE
        WHEN col.character_maximum_length IS NOT NULL THEN col.data_type || '(' || col.character_maximum_length || ')'
        ELSE col.data_type
       END AS type,
       col.is_nullable AS null,
       CASE
        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRI'
        ELSE ''
       END AS key,
       '' AS default,
       '' AS extra
       --kcu.constraint_name AS constraint_name
       --kcu.*,
       --tc.*
FROM   information_schema.columns col
LEFT JOIN   information_schema.key_column_usage kcu ON (kcu.table_name = col.table_name AND kcu.column_name = col.column_name)
LEFT JOIN   information_schema.table_constraints AS tc ON (kcu.constraint_name = tc.constraint_name)
WHERE  col.table_name = '%v'
ORDER BY col.ordinal_position;`
)

type genericPostgresWriter struct {
	e               Executor
	insertBulkLimit int
}

/* how to do an UPSERT/MERGE in PostgreSQL
 * http://stackoverflow.com/questions/17267417/how-do-i-do-an-upsert-merge-insert-on-duplicate-update-in-postgresq */
func (w *genericPostgresWriter) MergeTable(src *Table, dstName string, r Reader) error {
	tmpName := "gomig_tmp"
	stmts := make([]string, 0, 5)

	/* create temporary table */
	stmts = append(stmts,
		fmt.Sprintf("CREATE TEMPORARY TABLE %v (\n\t%v\n)\nON COMMIT DROP;\n", tmpName, ColumnsSql(src)))

	if PG_W_VERBOSE {
		log.Println("postgres: preparing to read values from source db")
	}

	/* bulk insert values */
	rows, err := r.Read(src)
	if err != nil {
		return err
	}
	defer rows.Close()

	if PG_W_VERBOSE {
		log.Print("postgres: query done, scanning rows...")
	}

	/* an alternate way to do this, with type assertions
	 * but possibly less accurately: http://go-database-sql.org/varcols.html */
	pointers := make([]interface{}, len(src.Columns))
	containers := make([]sql.RawBytes, len(src.Columns))
	for i, _ := range pointers {
		pointers[i] = &containers[i]
	}
	stringrep := make([]string, 0, len(src.Columns))
	insertLines := make([]string, 0, 32)
	for rows.Next() {
		err := rows.Scan(pointers...)
		if err != nil {
			log.Println("postgres: error while reading from source:", err)
			return err
		}

		for idx, val := range containers {
			if val == nil {
				stringrep = append(stringrep, "NULL")
			} else {
				switch src.Columns[idx].Type {
				case "text":
					stringrep = append(stringrep, "$$"+string(val)+"$$")
				case "boolean":
					/* ascii(48) = "0" and ascii(49) = "1" */
					switch val[0] {
					case 48:
						stringrep = append(stringrep, "f")
					case 49:
						stringrep = append(stringrep, "t")
					default:
						return fmt.Errorf("writer: did not recognize bool value: string(%v) = %v, val[0] = %v", val, string(val), val[0])
					}
				case "integer":
					stringrep = append(stringrep, string(val))
				default:
					stringrep = append(stringrep, string(val))
				}
			}
		}

		insertLines = append(insertLines, "("+strings.Join(stringrep, ",")+")")
		stringrep = stringrep[:0]

		if len(insertLines) > w.insertBulkLimit {
			stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
				tmpName, strings.Join(insertLines, ",\n\t")))

			insertLines = insertLines[:0]
		}
	}

	err = rows.Err()
	if err != nil {
		return err
	}

	if len(insertLines) > 0 {
		stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
			tmpName, strings.Join(insertLines, ",\n\t")))
	}

	if PG_W_VERBOSE {
		log.Print("postgres: rowscan done, creating merge statements")
	}

	/* analyze the temp table, for performance */
	stmts = append(stmts, fmt.Sprintf("ANALYZE %v;\n", tmpName))

	/* lock the target table */
	stmts = append(stmts, fmt.Sprintf("LOCK TABLE %v IN EXCLUSIVE MODE;", dstName))

	colnames := make([]string, 0, len(src.Columns))
	srccol := make([]string, 0, len(src.Columns))
	pkWhere := make([]string, 0, len(src.Columns))
	pkIsNull := make([]string, 0, len(src.Columns))
	colassign := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
		srccol = append(srccol, "src."+col.Name)
		if col.PrimaryKey {
			pkWhere = append(pkWhere, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
			pkIsNull = append(pkIsNull, fmt.Sprintf("dst.%[1]v IS NULL", col.Name))
		} else {
			colassign = append(colassign, fmt.Sprintf("%[1]v = src.%[1]v", col.Name))
		}
	}
	pkWherePart := strings.Join(pkWhere, "\nAND    ")
	pkIsNullPart := strings.Join(pkIsNull, "\nAND    ")
	srccolPart := strings.Join(srccol, ",\n       ")

	/* UPDATE from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
UPDATE %v AS dst
SET    %v
FROM   %v AS src
WHERE  %v;`, dstName, strings.Join(colassign, ",\n       "), tmpName, pkWherePart))

	/* INSERT from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
INSERT INTO %[1]v (%[3]v)
SELECT %[4]v
FROM   %[2]v AS src
LEFT OUTER JOIN %[1]v AS dst ON (
	   %[5]v
)
WHERE  %[6]v;
`, dstName, tmpName, strings.Join(colnames, ", "), srccolPart, pkWherePart, pkIsNullPart))

	if PG_W_VERBOSE {
		log.Print("postgres: statements completed, executing transaction")
	}

	err = w.e.Transaction(
		fmt.Sprintf("merge table %v into table %v", src.Name, dstName), stmts)
	return err
}

func (w *genericPostgresWriter) Close() error {
	return w.e.Close()
}

type PostgresWriter struct {
	genericPostgresWriter
}

func NewPostgresWriter(conf *Config) (*PostgresWriter, error) {
	db, err := openDB(conf)
	if err != nil {
		return nil, err
	}

	executor, err := NewDbExecutor(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection (WARNING: connection pooling might mess with this)", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresWriter{genericPostgresWriter{executor, 64}}, nil
}

type PostgresFileWriter struct {
	genericPostgresWriter
}

func NewPostgresFileWriter(filename string) (*PostgresFileWriter, error) {
	executor, err := NewFileExecutor(filename)
	if err != nil {
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresFileWriter{genericPostgresWriter{executor, 256}}, err
}

func PostgresType(genericType string) string {
	return genericType
}

func ColumnsSql(table *Table) string {
	colSql := make([]string, 0, len(table.Columns))

	for _, col := range table.Columns {
		colSql = append(colSql, fmt.Sprintf("%v %v", col.Name, PostgresType(col.Type)))
	}

	return strings.Join(colSql, ",\n\t")
}
<commit_msg>add a PK on the temp table<commit_after>package postgres

import (
	"database/sql"
	"fmt"
	. "github.com/aktau/gomig/db/common"
	"log"
	"strings"
)

var PG_W_VERBOSE = true

var (
	postgresInit = []string{
		"SET client_encoding = 'UTF8';",
		"SET standard_conforming_strings = off;",
		"SET check_function_bodies = false;",
		"SET client_min_messages = warning;",
	}
)

const (
	explainQuery = `
SELECT col.column_name AS field,
       CASE
        WHEN col.character_maximum_length IS NOT NULL THEN col.data_type || '(' || col.character_maximum_length || ')'
        ELSE col.data_type
       END AS type,
       col.is_nullable AS null,
       CASE
        WHEN tc.constraint_type = 'PRIMARY KEY' THEN 'PRI'
        ELSE ''
       END AS key,
       '' AS default,
       '' AS extra
       --kcu.constraint_name AS constraint_name
       --kcu.*,
       --tc.*
FROM   information_schema.columns col
LEFT JOIN   information_schema.key_column_usage kcu ON (kcu.table_name = col.table_name AND kcu.column_name = col.column_name)
LEFT JOIN   information_schema.table_constraints AS tc ON (kcu.constraint_name = tc.constraint_name)
WHERE  col.table_name = '%v'
ORDER BY col.ordinal_position;`
)

type genericPostgresWriter struct {
	e               Executor
	insertBulkLimit int
}

/* how to do an UPSERT/MERGE in PostgreSQL
 * http://stackoverflow.com/questions/17267417/how-do-i-do-an-upsert-merge-insert-on-duplicate-update-in-postgresq */
func (w *genericPostgresWriter) MergeTable(src *Table, dstName string, r Reader) error {
	tmpName := "gomig_tmp"
	stmts := make([]string, 0, 5)

	/* create temporary table */
	stmts = append(stmts,
		fmt.Sprintf("CREATE TEMPORARY TABLE %v (\n\t%v\n)\nON COMMIT DROP;\n", tmpName, ColumnsSql(src)))

	if PG_W_VERBOSE {
		log.Println("postgres: preparing to read values from source db")
	}

	/* bulk insert values */
	rows, err := r.Read(src)
	if err != nil {
		return err
	}
	defer rows.Close()

	if PG_W_VERBOSE {
		log.Print("postgres: query done, scanning rows...")
	}

	/* an alternate way to do this, with type assertions
	 * but possibly less accurately: http://go-database-sql.org/varcols.html */
	pointers := make([]interface{}, len(src.Columns))
	containers := make([]sql.RawBytes, len(src.Columns))
	for i, _ := range pointers {
		pointers[i] = &containers[i]
	}
	stringrep := make([]string, 0, len(src.Columns))
	insertLines := make([]string, 0, 32)
	for rows.Next() {
		err := rows.Scan(pointers...)
		if err != nil {
			log.Println("postgres: error while reading from source:", err)
			return err
		}

		for idx, val := range containers {
			if val == nil {
				stringrep = append(stringrep, "NULL")
			} else {
				switch src.Columns[idx].Type {
				case "text":
					stringrep = append(stringrep, "$$"+string(val)+"$$")
				case "boolean":
					/* ascii(48) = "0" and ascii(49) = "1" */
					switch val[0] {
					case 48:
						stringrep = append(stringrep, "f")
					case 49:
						stringrep = append(stringrep, "t")
					default:
						return fmt.Errorf("writer: did not recognize bool value: string(%v) = %v, val[0] = %v", val, string(val), val[0])
					}
				case "integer":
					stringrep = append(stringrep, string(val))
				default:
					stringrep = append(stringrep, string(val))
				}
			}
		}

		insertLines = append(insertLines, "("+strings.Join(stringrep, ",")+")")
		stringrep = stringrep[:0]

		if len(insertLines) > w.insertBulkLimit {
			stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
				tmpName, strings.Join(insertLines, ",\n\t")))

			insertLines = insertLines[:0]
		}
	}

	err = rows.Err()
	if err != nil {
		return err
	}

	if len(insertLines) > 0 {
		stmts = append(stmts, fmt.Sprintf("INSERT INTO %v VALUES\n\t%v;\n",
			tmpName, strings.Join(insertLines, ",\n\t")))
	}

	if PG_W_VERBOSE {
		log.Print("postgres: rowscan done, creating merge statements")
	}

	/* analyze the temp table, for performance */
	stmts = append(stmts, fmt.Sprintf("ANALYZE %v;\n", tmpName))

	/* lock the target table */
	stmts = append(stmts, fmt.Sprintf("LOCK TABLE %v IN EXCLUSIVE MODE;", dstName))

	colnames := make([]string, 0, len(src.Columns))
	srccol := make([]string, 0, len(src.Columns))
	pkWhere := make([]string, 0, len(src.Columns))
	pkIsNull := make([]string, 0, len(src.Columns))
	colassign := make([]string, 0, len(src.Columns))
	for _, col := range src.Columns {
		colnames = append(colnames, col.Name)
		srccol = append(srccol, "src."+col.Name)
		if col.PrimaryKey {
			pkWhere = append(pkWhere, fmt.Sprintf("dst.%[1]v = src.%[1]v", col.Name))
			pkIsNull = append(pkIsNull, fmt.Sprintf("dst.%[1]v IS NULL", col.Name))
		} else {
			colassign = append(colassign, fmt.Sprintf("%[1]v = src.%[1]v", col.Name))
		}
	}
	pkWherePart := strings.Join(pkWhere, "\nAND    ")
	pkIsNullPart := strings.Join(pkIsNull, "\nAND    ")
	srccolPart := strings.Join(srccol, ",\n       ")

	/* UPDATE from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
UPDATE %v AS dst
SET    %v
FROM   %v AS src
WHERE  %v;`, dstName, strings.Join(colassign, ",\n       "), tmpName, pkWherePart))

	/* INSERT from temp table to target table based on PK */
	stmts = append(stmts, fmt.Sprintf(`
INSERT INTO %[1]v (%[3]v)
SELECT %[4]v
FROM   %[2]v AS src
LEFT OUTER JOIN %[1]v AS dst ON (
	   %[5]v
)
WHERE  %[6]v;
`, dstName, tmpName, strings.Join(colnames, ", "), srccolPart, pkWherePart, pkIsNullPart))

	if PG_W_VERBOSE {
		log.Print("postgres: statements completed, executing transaction")
	}

	err = w.e.Transaction(
		fmt.Sprintf("merge table %v into table %v", src.Name, dstName), stmts)
	return err
}

func (w *genericPostgresWriter) Close() error {
	return w.e.Close()
}

type PostgresWriter struct {
	genericPostgresWriter
}

func NewPostgresWriter(conf *Config) (*PostgresWriter, error) {
	db, err := openDB(conf)
	if err != nil {
		return nil, err
	}

	executor, err := NewDbExecutor(db)
	if err != nil {
		db.Close()
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection (WARNING: connection pooling might mess with this)", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresWriter{genericPostgresWriter{executor, 64}}, nil
}

type PostgresFileWriter struct {
	genericPostgresWriter
}

func NewPostgresFileWriter(filename string) (*PostgresFileWriter, error) {
	executor, err := NewFileExecutor(filename)
	if err != nil {
		return nil, err
	}

	errors := executor.Multiple("initializing DB connection", postgresInit)
	if len(errors) > 0 {
		executor.Close()
		for _, err := range errors {
			log.Println("postgres error:", err)
		}
		return nil, errors[0]
	}

	return &PostgresFileWriter{genericPostgresWriter{executor, 256}}, err
}

func PostgresType(genericType string) string {
	return genericType
}

func ColumnsSql(table *Table) string {
	colSql := make([]string, 0, len(table.Columns))

	for _, col := range table.Columns {
		colSql = append(colSql, fmt.Sprintf("%v %v", col.Name, PostgresType(col.Type)))
	}

	pkCols := make([]string, 0, len(table.Columns))
	for _, col := range table.Columns {
		if col.PrimaryKey {
			pkCols = append(pkCols, col.Name)
		}
	}

	/* add the primary key */
	colSql = append(colSql, fmt.Sprintf("PRIMARY KEY (%v)",
		strings.Join(pkCols, ", ")))

	return strings.Join(colSql, ",\n\t")
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/*
All of the calls to the external SSH library
(code.google.com/p/go.crypto/ssh) will go through this package.  This gives us
the opportunity to adjust the interface for our needs.  More importantly, it
will allow us to more easily swap out the backend if all of our external SSH
calls are in the same place.
*/
package ssh

import (
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os/exec"
	"strconv"
	"time"
)

const DEFAULT_SSH_PORT = 22

// keychain implements the ssh.ClientKeyring interface
type keychain struct {
	keys []ssh.Signer
}

func (k *keychain) Key(i int) (ssh.PublicKey, error) {
	if i < 0 || i >= len(k.keys) {
		return nil, nil
	}

	return k.keys[i].PublicKey(), nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	return k.keys[i].Sign(rand, data)
}

func (k *keychain) add(key ssh.Signer) {
	k.keys = append(k.keys, key)
}

func (k *keychain) loadPEM(file string) error {
	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return err
	}
	key, err := ssh.ParsePrivateKey(buf)
	if err != nil {
		return err
	}
	k.add(key)
	return nil
}

// clientPassword implements the ssh.ClientPassword interface
type clientPassword string

func (p clientPassword) Password(user string) (string, error) {
	return string(p), nil
}

// Establish a code.google.com/p/go.crypto/ssh Session.
// The caller is responsible for closing the session.
func getSession(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16) (session *ssh.Session, err error) {

	var authorizers []ssh.ClientAuth = []ssh.ClientAuth{}
	if privKeyPath != "" {
		var clientKeychain *keychain = new(keychain)
		if err := clientKeychain.loadPEM(privKeyPath); err != nil {
			return session, err
		}
		authorizers = append(
			authorizers, ssh.ClientAuthKeyring(clientKeychain))
	}

	if password != nil {
		authorizers = append(
			authorizers, ssh.ClientAuthPassword(clientPassword(*password)))
	}

	if len(authorizers) == 0 {
		return session, errors.New("No authorization methods provided")
	}

	/* Try to authenticate with a public SSH key first, try a password if that fails */
	config := &ssh.ClientConfig{
		User: username,
		Auth: authorizers,
	}
	client, err := ssh.Dial(
		"tcp",
		addr+":"+strconv.FormatUint(uint64(portNum), 10),
		config)
	if err != nil {
		return session, err
	}

	session, err = client.NewSession()
	if err != nil {
		return session, err
	}
	return session, err
}

func TestConnection(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16) (err error) {

	var session *ssh.Session

	session, err = getSession(addr, username, password, privKeyPath, portNum)
	if err != nil {
		return err
	}
	defer session.Close()

	if err = session.Run("true"); err != nil {
		return err
	}
	return nil
}

func handleTimeout() {
	/* TODO: kill remote process and any other cleanup */
	fmt.Println("TIMEOUT!")
}

// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The username parameter is the username to use to SSH to the remote host.
// The password parameter is the password to use to SSH to the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum parameter is the SSH port number of the remote host.
// The command parameter is the command to run on the remote host.
// The timeout parameter is the number of seconds before abandoning the command.
// A timeout of 0 means no timeout.
func Run(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16,
	command string,
	timeout uint32) (stdout string, stderr string, err error) {

	var session *ssh.Session

	session, err = getSession(addr, username, password, privKeyPath, portNum)
	if err != nil {
		return "", "", err
	}
	defer session.Close()

	var stdout_buf bytes.Buffer
	var stderr_buf bytes.Buffer
	session.Stdout = &stdout_buf
	session.Stderr = &stderr_buf

	if timeout == 0 {
		if err = session.Run(command); err != nil {
			return "", "", err
		}
	} else {
		c := make(chan error)
		var e error
		go func() {
			c <- session.Run(command)
		}()
		select {
		case e = <-c:
			if e == nil {
				return "", "", err
			}
		case <-time.After(time.Duration(timeout) * time.Second):
			err = errors.New("timeout")
			handleTimeout()
		}
	}

	return stdout_buf.String(), stderr_buf.String(), err
}

// Secure copy (scp) from localhost to addr
// Run a separate scp process (for now) to secure copy files between hosts.
// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum is the SSH port number of the remote host.
// The incoming parameter indicates which direction to perform the copy.
// The recursive parameter indicates whether to use the -r scp option.
// Password authentication not supported for this function.
// XXX: This function should probably go away in favor of a single Scp function when Issue #1 is fixed.
func ScpTo(
	addr string,
	username string,
	portNum uint16,
	recursive bool,
	localPath string,
	remotePath string) (err error) {

	/* Unfortunately, there doesn't appear to be an SFTP or SCP library, so
	we'll just have to run a separate scp process.  This means no password
	authentication for when calling this method. */
	var stdout bytes.Buffer
	var args []string = []string{}
	var cmd *exec.Cmd
	if recursive {
		args = append(args, "-r")
	}
	args = append(args, fmt.Sprintf("-P"))
	args = append(args, strconv.FormatUint(uint64(portNum), 10))
	args = append(args, localPath)
	args = append(args, fmt.Sprintf("%s@%s:%s", username, addr, remotePath))
	cmd = exec.Command("scp", args...)
	cmd.Stderr = &stdout
	err = cmd.Run()
	if err != nil {
		return errors.New("scp to " + addr + " failed: " + err.Error())
	}
	return nil

}

// Secure copy (scp) from addr to localhost
// Run a separate scp process (for now) to secure copy files between hosts.
// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum is the SSH port number of the remote host.
// The incoming parameter indicates which direction to perform the copy.
// The recursive parameter indicates whether to use the -r scp option.
// Password authentication not supported for this function.
// XXX: This function should probably go away in favor of a single Scp function when Issue #1 is fixed.
func ScpFrom(
	addr string,
	username string,
	portNum uint16,
	recursive bool,
	remotePath string,
	localPath string) (err error) {

	/* Unfortunately, there doesn't appear to be an SFTP or SCP library, so
	we'll just have to run a separate scp process.  This means no password
	authentication for when calling this method. */
	var stdout bytes.Buffer
	var args []string = []string{}
	var cmd *exec.Cmd
	if recursive {
		args = append(args, "-r")
	}
	args = append(args, fmt.Sprintf("-P"))
	args = append(args, strconv.FormatUint(uint64(portNum), 10))
	args = append(args, fmt.Sprintf("%s@%s:%s", username, addr, remotePath))
	args = append(args, localPath)
	cmd = exec.Command("scp", args...)
	cmd.Stderr = &stdout
	err = cmd.Run()
	if err != nil {
		return errors.New("scp to " + addr + "failed: " + err.Error())
	}
	return nil
}
<commit_msg>Implement a rudimentary handleTimeout for ssh<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

/*
All of the calls to the external SSH library
(code.google.com/p/go.crypto/ssh) will go through this package.  This gives us
the opportunity to adjust the interface for our needs.  More importantly, it
will allow us to more easily swap out the backend if all of our external SSH
calls are in the same place.
*/
package ssh

import (
	"bytes"
	"code.google.com/p/go.crypto/ssh"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os/exec"
	"strconv"
	"time"
)

const DEFAULT_SSH_PORT = 22

// keychain implements the ssh.ClientKeyring interface
type keychain struct {
	keys []ssh.Signer
}

func (k *keychain) Key(i int) (ssh.PublicKey, error) {
	if i < 0 || i >= len(k.keys) {
		return nil, nil
	}

	return k.keys[i].PublicKey(), nil
}

func (k *keychain) Sign(i int, rand io.Reader, data []byte) (sig []byte, err error) {
	return k.keys[i].Sign(rand, data)
}

func (k *keychain) add(key ssh.Signer) {
	k.keys = append(k.keys, key)
}

func (k *keychain) loadPEM(file string) error {
	buf, err := ioutil.ReadFile(file)
	if err != nil {
		return err
	}
	key, err := ssh.ParsePrivateKey(buf)
	if err != nil {
		return err
	}
	k.add(key)
	return nil
}

// clientPassword implements the ssh.ClientPassword interface
type clientPassword string

func (p clientPassword) Password(user string) (string, error) {
	return string(p), nil
}

// Establish a code.google.com/p/go.crypto/ssh Session.
// The caller is responsible for closing the session.
func getSession(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16) (session *ssh.Session, err error) {

	var authorizers []ssh.ClientAuth = []ssh.ClientAuth{}
	if privKeyPath != "" {
		var clientKeychain *keychain = new(keychain)
		if err := clientKeychain.loadPEM(privKeyPath); err != nil {
			return session, err
		}
		authorizers = append(
			authorizers, ssh.ClientAuthKeyring(clientKeychain))
	}

	if password != nil {
		authorizers = append(
			authorizers, ssh.ClientAuthPassword(clientPassword(*password)))
	}

	if len(authorizers) == 0 {
		return session, errors.New("No authorization methods provided")
	}

	/* Try to authenticate with a public SSH key first, try a password if that fails */
	config := &ssh.ClientConfig{
		User: username,
		Auth: authorizers,
	}
	client, err := ssh.Dial(
		"tcp",
		addr+":"+strconv.FormatUint(uint64(portNum), 10),
		config)
	if err != nil {
		return session, err
	}

	session, err = client.NewSession()
	if err != nil {
		return session, err
	}
	return session, err
}

func TestConnection(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16) (err error) {

	var session *ssh.Session

	session, err = getSession(addr, username, password, privKeyPath, portNum)
	if err != nil {
		return err
	}
	defer session.Close()

	if err = session.Run("true"); err != nil {
		return err
	}
	return nil
}

// Connect to the remote host and attempt to kill -9 the process that timed out
// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The username parameter is the username to use to SSH to the remote host.
// The password parameter is the password to use to SSH to the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum parameter is the SSH port number of the remote host.
// The command parameter is the command that timed out on the remote host.
func handleTimeout(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16,
	expiredCmd string) (stdout string, stderr string, err error) {

	// TODO: perhaps be a bit more diplomatic in the future by trying
	// -TERM (instead of -9) first?
	// The negation of the PID is important, see kill(1)
	killCmd := fmt.Sprintf("kill -9 -$(pgrep -f \"%s\")", expiredCmd)

	session, err := getSession(addr, username, password, privKeyPath, portNum)
	if err != nil {
		log.Printf("Failed to connect for cleanup after timeout: %s", err.Error())
		return
	}
	defer session.Close()

	session.Run(killCmd)
	return
}

// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The username parameter is the username to use to SSH to the remote host.
// The password parameter is the password to use to SSH to the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum parameter is the SSH port number of the remote host.
// The command parameter is the command to run on the remote host.
// The timeout parameter is the number of seconds before abandoning the command.
// A timeout of 0 means no timeout.
func Run(
	addr string,
	username string,
	password *string,
	privKeyPath string,
	portNum uint16,
	command string,
	timeout uint32) (stdout string, stderr string, err error) {

	var session *ssh.Session

	session, err = getSession(addr, username, password, privKeyPath, portNum)
	if err != nil {
		return "", "", err
	}
	defer session.Close()

	var stdout_buf bytes.Buffer
	var stderr_buf bytes.Buffer
	session.Stdout = &stdout_buf
	session.Stderr = &stderr_buf

	if timeout == 0 {
		if err = session.Run(command); err != nil {
			return "", "", err
		}
	} else {
		c := make(chan error)
		var e error
		go func() {
			c <- session.Run(command)
		}()
		select {
		case e = <-c:
			if e == nil {
				return "", "", err
			}
		case <-time.After(time.Duration(timeout) * time.Second):
			err = errors.New("timeout")
			defer handleTimeout(
				addr,
				username,
				password,
				privKeyPath,
				portNum,
				command)
		}
	}

	return stdout_buf.String(), stderr_buf.String(), err
}

// Secure copy (scp) from localhost to addr
// Run a separate scp process (for now) to secure copy files between hosts.
// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum is the SSH port number of the remote host.
// The incoming parameter indicates which direction to perform the copy.
// The recursive parameter indicates whether to use the -r scp option.
// Password authentication not supported for this function.
// XXX: This function should probably go away in favor of a single Scp function when Issue #1 is fixed.
func ScpTo(
	addr string,
	username string,
	portNum uint16,
	recursive bool,
	localPath string,
	remotePath string) (err error) {

	/* Unfortunately, there doesn't appear to be an SFTP or SCP library, so
	we'll just have to run a separate scp process.  This means no password
	authentication for when calling this method. */
	var stdout bytes.Buffer
	var args []string = []string{}
	var cmd *exec.Cmd
	if recursive {
		args = append(args, "-r")
	}
	args = append(args, fmt.Sprintf("-P"))
	args = append(args, strconv.FormatUint(uint64(portNum), 10))
	args = append(args, localPath)
	args = append(args, fmt.Sprintf("%s@%s:%s", username, addr, remotePath))
	cmd = exec.Command("scp", args...)
	cmd.Stderr = &stdout
	err = cmd.Run()
	if err != nil {
		return errors.New("scp to " + addr + " failed: " + err.Error())
	}
	return nil

}

// Secure copy (scp) from addr to localhost
// Run a separate scp process (for now) to secure copy files between hosts.
// The addr parameter is the address (IP, hostname, etc) of the remote host.
// The privKeyPath parameter is the path to the private key of the master.
// The portNum is the SSH port number of the remote host.
// The incoming parameter indicates which direction to perform the copy.
// The recursive parameter indicates whether to use the -r scp option.
// Password authentication not supported for this function.
// XXX: This function should probably go away in favor of a single Scp function when Issue #1 is fixed.
func ScpFrom(
	addr string,
	username string,
	portNum uint16,
	recursive bool,
	remotePath string,
	localPath string) (err error) {

	/* Unfortunately, there doesn't appear to be an SFTP or SCP library, so
	we'll just have to run a separate scp process.  This means no password
	authentication for when calling this method. */
	var stdout bytes.Buffer
	var args []string = []string{}
	var cmd *exec.Cmd
	if recursive {
		args = append(args, "-r")
	}
	args = append(args, fmt.Sprintf("-P"))
	args = append(args, strconv.FormatUint(uint64(portNum), 10))
	args = append(args, fmt.Sprintf("%s@%s:%s", username, addr, remotePath))
	args = append(args, localPath)
	cmd = exec.Command("scp", args...)
	cmd.Stderr = &stdout
	err = cmd.Run()
	if err != nil {
		return errors.New("scp to " + addr + "failed: " + err.Error())
	}
	return nil
}
<|endoftext|>


<commit_before>package vegeta

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
)

// Target is an HTTP request blueprint.
type Target struct {
	Method string      `json:"method"`
	URL    string      `json:"url"`
	Body   []byte      `json:"body"`
	Header http.Header `json:"header"`
}

// Request creates an *http.Request out of Target and returns it along with an
// error in case of failure.
func (t *Target) Request() (*http.Request, error) {
	req, err := http.NewRequest(t.Method, t.URL, bytes.NewReader(t.Body))
	if err != nil {
		return nil, err
	}
	for k, vs := range t.Header {
		req.Header[k] = make([]string, len(vs))
		copy(req.Header[k], vs)
	}
	if host := req.Header.Get("Host"); host != "" {
		req.Host = host
	}
	return req, nil
}

var (
	// ErrNoTargets is returned when not enough Targets are available.
	ErrNoTargets = errors.New("no targets to attack")
	// ErrNilTarget is returned when the passed Target pointer is nil.
	ErrNilTarget = errors.New("nil target")
)

// A Targeter decodes a Target or returns an error in case of failure.
// Implementations must be safe for concurrent use.
type Targeter func(*Target) error

// NewJSONTargeter returns a new targeter that decodes one Target from the
// given io.Reader on every invocation. Each target is one JSON object in its own line.
// The body field of each target must be base64 encoded.
//
//    {"method":"POST", "url":"https://goku/1", "header":{"Content-Type":["text/plain"], "body": "Rk9P"}
//    {"method":"GET",  "url":"https://goku/2"}
//
// body will be set as the Target's body if no body is provided in each target definiton.
// hdr will be merged with the each Target's headers.
//
func NewJSONTargeter(src io.Reader, body []byte, header http.Header) Targeter {
	type decoder struct {
		*json.Decoder
		sync.Mutex
	}
	dec := decoder{Decoder: json.NewDecoder(src)}

	return func(tgt *Target) (err error) {
		if tgt == nil {
			return ErrNilTarget
		}

		dec.Lock()
		defer dec.Unlock()

		if err = dec.Decode(tgt); err == nil {
			return nil
		}

		switch err {
		case io.EOF:
			return ErrNoTargets
		default:
			return err
		}
	}
}

// NewStaticTargeter returns a Targeter which round-robins over the passed
// Targets.
func NewStaticTargeter(tgts ...Target) Targeter {
	i := int64(-1)
	return func(tgt *Target) error {
		if tgt == nil {
			return ErrNilTarget
		}
		*tgt = tgts[atomic.AddInt64(&i, 1)%int64(len(tgts))]
		return nil
	}
}

// NewEagerTargeter eagerly reads all Targets out of the provided Targeter and
// returns a NewStaticTargeter with them.
func NewEagerTargeter(t Targeter) (Targeter, error) {
	var (
		tgts []Target
		tgt  Target
		err  error
	)

	for {
		if err = t(&tgt); err == ErrNoTargets {
			break
		} else if err != nil {
			return nil, err
		}
		tgts = append(tgts, tgt)
	}

	if len(tgts) == 0 {
		return nil, ErrNoTargets
	}

	return NewStaticTargeter(tgts...), nil
}

// NewHTTPTargeter returns a new Targeter that decodes one Target from the
// given io.Reader on every invocation. The format is as follows:
//
//    GET https://foo.bar/a/b/c
//    Header-X: 123
//    Header-Y: 321
//    @/path/to/body/file
//
//    POST https://foo.bar/b/c/a
//    Header-X: 123
//
// body will be set as the Target's body if no body is provided.
// hdr will be merged with the each Target's headers.
func NewHTTPTargeter(src io.Reader, body []byte, hdr http.Header) Targeter {
	var mu sync.Mutex
	sc := peekingScanner{src: bufio.NewScanner(src)}
	return func(tgt *Target) (err error) {
		mu.Lock()
		defer mu.Unlock()

		if tgt == nil {
			return ErrNilTarget
		}

		var line string
		for {
			if !sc.Scan() {
				return ErrNoTargets
			}
			line = strings.TrimSpace(sc.Text())
			if len(line) != 0 {
				break
			}
		}

		tgt.Body = body
		tgt.Header = http.Header{}
		for k, vs := range hdr {
			tgt.Header[k] = vs
		}

		tokens := strings.SplitN(line, " ", 2)
		if len(tokens) < 2 {
			return fmt.Errorf("bad target: %s", line)
		}
		if !startsWithHTTPMethod(line) {
			return fmt.Errorf("bad method: %s", tokens[0])
		}
		tgt.Method = tokens[0]
		if _, err = url.ParseRequestURI(tokens[1]); err != nil {
			return fmt.Errorf("bad URL: %s", tokens[1])
		}
		tgt.URL = tokens[1]
		line = strings.TrimSpace(sc.Peek())
		if line == "" || startsWithHTTPMethod(line) {
			return nil
		}
		for sc.Scan() {
			if line = strings.TrimSpace(sc.Text()); line == "" {
				break
			} else if strings.HasPrefix(line, "@") {
				if tgt.Body, err = ioutil.ReadFile(line[1:]); err != nil {
					return fmt.Errorf("bad body: %s", err)
				}
				break
			}
			tokens = strings.SplitN(line, ":", 2)
			if len(tokens) < 2 {
				return fmt.Errorf("bad header: %s", line)
			}
			for i := range tokens {
				if tokens[i] = strings.TrimSpace(tokens[i]); tokens[i] == "" {
					return fmt.Errorf("bad header: %s", line)
				}
			}
			// Add key/value directly to the http.Header (map[string][]string).
			// http.Header.Add() canonicalizes keys but vegeta is used
			// to test systems that require case-sensitive headers.
			tgt.Header[tokens[0]] = append(tgt.Header[tokens[0]], tokens[1])
		}
		if err = sc.Err(); err != nil {
			return ErrNoTargets
		}
		return nil
	}
}

var httpMethodChecker = regexp.MustCompile("^[A-Z]+\\s")

// A line starts with an http method when the first word is uppercase ascii
// followed by a space.
func startsWithHTTPMethod(t string) bool {
	return httpMethodChecker.MatchString(t)
}

// Wrap a Scanner so we can cheat and look at the next value and react accordingly,
// but still have it be around the next time we Scan() + Text()
type peekingScanner struct {
	src    *bufio.Scanner
	peeked string
}

func (s *peekingScanner) Err() error {
	return s.src.Err()
}

func (s *peekingScanner) Peek() string {
	if !s.src.Scan() {
		return ""
	}
	s.peeked = s.src.Text()
	return s.peeked
}

func (s *peekingScanner) Scan() bool {
	if s.peeked == "" {
		return s.src.Scan()
	}
	return true
}

func (s *peekingScanner) Text() string {
	if s.peeked == "" {
		return s.src.Text()
	}
	t := s.peeked
	s.peeked = ""
	return t
}
<commit_msg>targets: Introduce constants for different formats<commit_after>package vegeta

import (
	"bufio"
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"sync"
	"sync/atomic"
)

// Target is an HTTP request blueprint.
type Target struct {
	Method string      `json:"method"`
	URL    string      `json:"url"`
	Body   []byte      `json:"body"`
	Header http.Header `json:"header"`
}

// Request creates an *http.Request out of Target and returns it along with an
// error in case of failure.
func (t *Target) Request() (*http.Request, error) {
	req, err := http.NewRequest(t.Method, t.URL, bytes.NewReader(t.Body))
	if err != nil {
		return nil, err
	}
	for k, vs := range t.Header {
		req.Header[k] = make([]string, len(vs))
		copy(req.Header[k], vs)
	}
	if host := req.Header.Get("Host"); host != "" {
		req.Host = host
	}
	return req, nil
}

var (
	// ErrNoTargets is returned when not enough Targets are available.
	ErrNoTargets = errors.New("no targets to attack")
	// ErrNilTarget is returned when the passed Target pointer is nil.
	ErrNilTarget = errors.New("nil target")
)

const (
	// HTTPTargetFormat is the human readable identifier for the HTTP target format.
	HTTPTargetFormat = "http"
	// JSONTargetFormat is the human readable identifier for the JSON target format.
	JSONTargetFormat = "json"
)

// A Targeter decodes a Target or returns an error in case of failure.
// Implementations must be safe for concurrent use.
type Targeter func(*Target) error

// NewJSONTargeter returns a new targeter that decodes one Target from the
// given io.Reader on every invocation. Each target is one JSON object in its own line.
// The body field of each target must be base64 encoded.
//
//    {"method":"POST", "url":"https://goku/1", "header":{"Content-Type":["text/plain"], "body": "Rk9P"}
//    {"method":"GET",  "url":"https://goku/2"}
//
// body will be set as the Target's body if no body is provided in each target definiton.
// hdr will be merged with the each Target's headers.
//
func NewJSONTargeter(src io.Reader, body []byte, header http.Header) Targeter {
	type decoder struct {
		*json.Decoder
		sync.Mutex
	}
	dec := decoder{Decoder: json.NewDecoder(src)}

	return func(tgt *Target) (err error) {
		if tgt == nil {
			return ErrNilTarget
		}

		dec.Lock()
		defer dec.Unlock()

		if err = dec.Decode(tgt); err == nil {
			return nil
		}

		switch err {
		case io.EOF:
			return ErrNoTargets
		default:
			return err
		}
	}
}

// NewStaticTargeter returns a Targeter which round-robins over the passed
// Targets.
func NewStaticTargeter(tgts ...Target) Targeter {
	i := int64(-1)
	return func(tgt *Target) error {
		if tgt == nil {
			return ErrNilTarget
		}
		*tgt = tgts[atomic.AddInt64(&i, 1)%int64(len(tgts))]
		return nil
	}
}

// NewEagerTargeter eagerly reads all Targets out of the provided Targeter and
// returns a NewStaticTargeter with them.
func NewEagerTargeter(t Targeter) (Targeter, error) {
	var (
		tgts []Target
		tgt  Target
		err  error
	)

	for {
		if err = t(&tgt); err == ErrNoTargets {
			break
		} else if err != nil {
			return nil, err
		}
		tgts = append(tgts, tgt)
	}

	if len(tgts) == 0 {
		return nil, ErrNoTargets
	}

	return NewStaticTargeter(tgts...), nil
}

// NewHTTPTargeter returns a new Targeter that decodes one Target from the
// given io.Reader on every invocation. The format is as follows:
//
//    GET https://foo.bar/a/b/c
//    Header-X: 123
//    Header-Y: 321
//    @/path/to/body/file
//
//    POST https://foo.bar/b/c/a
//    Header-X: 123
//
// body will be set as the Target's body if no body is provided.
// hdr will be merged with the each Target's headers.
func NewHTTPTargeter(src io.Reader, body []byte, hdr http.Header) Targeter {
	var mu sync.Mutex
	sc := peekingScanner{src: bufio.NewScanner(src)}
	return func(tgt *Target) (err error) {
		mu.Lock()
		defer mu.Unlock()

		if tgt == nil {
			return ErrNilTarget
		}

		var line string
		for {
			if !sc.Scan() {
				return ErrNoTargets
			}
			line = strings.TrimSpace(sc.Text())
			if len(line) != 0 {
				break
			}
		}

		tgt.Body = body
		tgt.Header = http.Header{}
		for k, vs := range hdr {
			tgt.Header[k] = vs
		}

		tokens := strings.SplitN(line, " ", 2)
		if len(tokens) < 2 {
			return fmt.Errorf("bad target: %s", line)
		}
		if !startsWithHTTPMethod(line) {
			return fmt.Errorf("bad method: %s", tokens[0])
		}
		tgt.Method = tokens[0]
		if _, err = url.ParseRequestURI(tokens[1]); err != nil {
			return fmt.Errorf("bad URL: %s", tokens[1])
		}
		tgt.URL = tokens[1]
		line = strings.TrimSpace(sc.Peek())
		if line == "" || startsWithHTTPMethod(line) {
			return nil
		}
		for sc.Scan() {
			if line = strings.TrimSpace(sc.Text()); line == "" {
				break
			} else if strings.HasPrefix(line, "@") {
				if tgt.Body, err = ioutil.ReadFile(line[1:]); err != nil {
					return fmt.Errorf("bad body: %s", err)
				}
				break
			}
			tokens = strings.SplitN(line, ":", 2)
			if len(tokens) < 2 {
				return fmt.Errorf("bad header: %s", line)
			}
			for i := range tokens {
				if tokens[i] = strings.TrimSpace(tokens[i]); tokens[i] == "" {
					return fmt.Errorf("bad header: %s", line)
				}
			}
			// Add key/value directly to the http.Header (map[string][]string).
			// http.Header.Add() canonicalizes keys but vegeta is used
			// to test systems that require case-sensitive headers.
			tgt.Header[tokens[0]] = append(tgt.Header[tokens[0]], tokens[1])
		}
		if err = sc.Err(); err != nil {
			return ErrNoTargets
		}
		return nil
	}
}

var httpMethodChecker = regexp.MustCompile("^[A-Z]+\\s")

// A line starts with an http method when the first word is uppercase ascii
// followed by a space.
func startsWithHTTPMethod(t string) bool {
	return httpMethodChecker.MatchString(t)
}

// Wrap a Scanner so we can cheat and look at the next value and react accordingly,
// but still have it be around the next time we Scan() + Text()
type peekingScanner struct {
	src    *bufio.Scanner
	peeked string
}

func (s *peekingScanner) Err() error {
	return s.src.Err()
}

func (s *peekingScanner) Peek() string {
	if !s.src.Scan() {
		return ""
	}
	s.peeked = s.src.Text()
	return s.peeked
}

func (s *peekingScanner) Scan() bool {
	if s.peeked == "" {
		return s.src.Scan()
	}
	return true
}

func (s *peekingScanner) Text() string {
	if s.peeked == "" {
		return s.src.Text()
	}
	t := s.peeked
	s.peeked = ""
	return t
}
<|endoftext|>


<commit_before>package web

import (
	"html/template"
)

var (
	dashboardTemplate *template.Template

	dashboardHTML = `<!DOCTYPE html>

    <title>Micro Debug</title>
    <meta name="application-name" content="netdata">
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    <meta charset="utf-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <meta name="apple-mobile-web-app-capable" content="yes">
    <meta name="apple-mobile-web-app-status-bar-style" content="black-translucent">
  <link href="https://fonts.googleapis.com/css?family=Source+Code+Pro&display=swap" rel="stylesheet">

     Debug

`
)
<commit_msg>update debug template<commit_after>package web

import (
	"html/template"
)

var (
	dashboardTemplate *template.Template

	dashboardHTML = `<!DOCTYPE html>

    <title>Micro Debug</title>
    <meta name="application-name" content="netdata">
    <meta http-equiv="Content-Type" content="text/html; charset=utf-8" />
    <meta charset="utf-8">
    <meta http-equiv="X-UA-Compatible" content="IE=edge,chrome=1">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <meta name="apple-mobile-web-app-capable" content="yes">
    <meta name="apple-mobile-web-app-status-bar-style" content="black-translucent">
  <link href="https://fonts.googleapis.com/css?family=Source+Code+Pro&display=swap" rel="stylesheet">

     Debug

`
)
<|endoftext|>


<commit_before>package hasher

import (
	"sync"
)

type Buffer struct {
	buf  []Chunk
	eof  chan struct{}
	mu   sync.RWMutex
	cond sync.Cond
}

func NewBuffer(chunked <-chan Chunk) *Buffer {
	b := &Buffer{eof: make(chan struct{})}
	b.cond.L = &b.mu
	go func() {
		for chunk := range chunked {
			b.mu.Lock()
			b.buf = append(b.buf, chunk)
			b.mu.Unlock()
			b.cond.Broadcast()
		}
		close(b.eof)
		b.mu.Lock()
		b.mu.Unlock()
		b.cond.Broadcast()
	}()
	return b
}

func (b *Buffer) isEof() bool {
	select {
	case <-b.eof:
		return true
	default:
		return false
	}
}

func (b *Buffer) Get(idx int) (chunk Chunk, eof bool) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	for len(b.buf) <= idx && !b.isEof() {
		b.cond.Wait()
	}

	if idx < len(b.buf) {
		return b.buf[idx], false
	} else {
		return Chunk{}, true
	}
}

func (b *Buffer) Eof() <-chan struct{} {
	return b.eof
}

func (b *Buffer) Len() (length int, finished bool) {
	finished = b.isEof()
	b.mu.Lock()
	length = len(b.buf)
	b.mu.Unlock()
	return
}

func (b *Buffer) NewReader(cancel <-chan bool) <-chan Chunk {
	reader := make(chan Chunk, 20)
	go func() {
		for idx := 0; true; idx++ {
			chunk, eof := b.Get(idx)
			if eof {
				break
			}
			select {
			case reader <- chunk:
			case <-cancel:
				break
			}
		}
		close(reader)
	}()
	return reader
}
<commit_msg>Wrong locker given to sync.Cond<commit_after>package hasher

import (
	"sync"
)

type Buffer struct {
	buf  []Chunk
	eof  chan struct{}
	mu   sync.RWMutex
	cond sync.Cond
}

func NewBuffer(chunked <-chan Chunk) *Buffer {
	b := &Buffer{eof: make(chan struct{})}
	b.cond.L = b.mu.RLocker()
	go func() {
		for chunk := range chunked {
			b.mu.Lock()
			b.buf = append(b.buf, chunk)
			b.mu.Unlock()
			b.cond.Broadcast()
		}
		close(b.eof)
		b.mu.Lock()
		b.mu.Unlock()
		b.cond.Broadcast()
	}()
	return b
}

func (b *Buffer) isEof() bool {
	select {
	case <-b.eof:
		return true
	default:
		return false
	}
}

func (b *Buffer) Get(idx int) (chunk Chunk, eof bool) {
	b.mu.RLock()
	defer b.mu.RUnlock()

	for len(b.buf) <= idx && !b.isEof() {
		b.cond.Wait()
	}

	if idx < len(b.buf) {
		return b.buf[idx], false
	} else {
		return Chunk{}, true
	}
}

func (b *Buffer) Eof() <-chan struct{} {
	return b.eof
}

func (b *Buffer) Len() (length int, finished bool) {
	finished = b.isEof()
	b.mu.Lock()
	length = len(b.buf)
	b.mu.Unlock()
	return
}

func (b *Buffer) NewReader(cancel <-chan bool) <-chan Chunk {
	reader := make(chan Chunk, 20)
	go func() {
		for idx := 0; true; idx++ {
			chunk, eof := b.Get(idx)
			if eof {
				break
			}
			select {
			case reader <- chunk:
			case <-cancel:
				break
			}
		}
		close(reader)
	}()
	return reader
}
<|endoftext|>


<commit_before>package hashstructure

import (
	"encoding/binary"
	"fmt"
	"hash"
	"hash/fnv"
	"reflect"
)

// HashOptions are options that are available for hashing.
type HashOptions struct {
	// Hasher is the hash function to use. If this isn't set, it will
	// default to FNV.
	Hasher hash.Hash64

	// TagName is the struct tag to look at when hashing the structure.
	// By default this is "hash".
	TagName string

	// ZeroNil is flag determining if nil pointer should be treated equal
	// to a zero value of pointed type. By default this is false.
	ZeroNil bool
}

// Hash returns the hash value of an arbitrary value.
//
// If opts is nil, then default options will be used. See HashOptions
// for the default values.
//
// Notes on the value:
//
//   * Unexported fields on structs are ignored and do not affect the
//     hash value.
//
//   * Adding an exported field to a struct with the zero value will change
//     the hash value.
//
// For structs, the hashing can be controlled using tags. For example:
//
//    struct {
//        Name string
//        UUID string `hash:"ignore"`
//    }
//
// The available tag values are:
//
//   * "ignore" or "-" - The field will be ignored and not affect the hash code.
//
//   * "set" - The field will be treated as a set, where ordering doesn't
//             affect the hash code. This only works for slices.
//
func Hash(v interface{}, opts *HashOptions) (uint64, error) {
	// Create default options
	if opts == nil {
		opts = &HashOptions{}
	}
	if opts.Hasher == nil {
		opts.Hasher = fnv.New64()
	}
	if opts.TagName == "" {
		opts.TagName = "hash"
	}

	// Reset the hash
	opts.Hasher.Reset()

	// Create our walker and walk the structure
	w := &walker{
		h:       opts.Hasher,
		tag:     opts.TagName,
		zeronil: opts.ZeroNil,
	}
	return w.visit(reflect.ValueOf(v), nil)
}

type walker struct {
	h       hash.Hash64
	tag     string
	zeronil bool
}

type visitOpts struct {
	// Flags are a bitmask of flags to affect behavior of this visit
	Flags visitFlag

	// Information about the struct containing this field
	Struct      interface{}
	StructField string
}

func (w *walker) visit(v reflect.Value, opts *visitOpts) (uint64, error) {
	t := reflect.TypeOf(0)

	// Loop since these can be wrapped in multiple layers of pointers
	// and interfaces.
	for {
		// If we have an interface, dereference it. We have to do this up
		// here because it might be a nil in there and the check below must
		// catch that.
		if v.Kind() == reflect.Interface {
			v = v.Elem()
			continue
		}

		if v.Kind() == reflect.Ptr {
			if w.zeronil {
				t = v.Type().Elem()
			}
			v = reflect.Indirect(v)
			continue
		}

		break
	}

	// If it is nil, treat it like a zero.
	if !v.IsValid() {
		v = reflect.Zero(t)
	}

	// Binary writing can use raw ints, we have to convert to
	// a sized-int, we'll choose the largest...
	switch v.Kind() {
	case reflect.Int:
		v = reflect.ValueOf(int64(v.Int()))
	case reflect.Uint:
		v = reflect.ValueOf(uint64(v.Uint()))
	case reflect.Bool:
		var tmp int8
		if v.Bool() {
			tmp = 1
		}
		v = reflect.ValueOf(tmp)
	}

	k := v.Kind()

	// We can shortcut numeric values by directly binary writing them
	if k >= reflect.Int && k <= reflect.Complex64 {
		// A direct hash calculation
		w.h.Reset()
		err := binary.Write(w.h, binary.LittleEndian, v.Interface())
		return w.h.Sum64(), err
	}

	switch k {
	case reflect.Array:
		var h uint64
		l := v.Len()
		for i := 0; i < l; i++ {
			current, err := w.visit(v.Index(i), nil)
			if err != nil {
				return 0, err
			}

			h = hashUpdateOrdered(w.h, h, current)
		}

		return h, nil

	case reflect.Map:
		var includeMap IncludableMap
		if opts != nil && opts.Struct != nil {
			if v, ok := opts.Struct.(IncludableMap); ok {
				includeMap = v
			}
		}

		// Build the hash for the map. We do this by XOR-ing all the key
		// and value hashes. This makes it deterministic despite ordering.
		var h uint64
		for _, k := range v.MapKeys() {
			v := v.MapIndex(k)
			if includeMap != nil {
				incl, err := includeMap.HashIncludeMap(
					opts.StructField, k.Interface(), v.Interface())
				if err != nil {
					return 0, err
				}
				if !incl {
					continue
				}
			}

			kh, err := w.visit(k, nil)
			if err != nil {
				return 0, err
			}
			vh, err := w.visit(v, nil)
			if err != nil {
				return 0, err
			}

			fieldHash := hashUpdateOrdered(w.h, kh, vh)
			h = hashUpdateUnordered(h, fieldHash)
		}

		return h, nil

	case reflect.Struct:
		var include Includable
		parent := v.Interface()
		if impl, ok := parent.(Includable); ok {
			include = impl
		}

		t := v.Type()
		h, err := w.visit(reflect.ValueOf(t.Name()), nil)
		if err != nil {
			return 0, err
		}

		l := v.NumField()
		for i := 0; i < l; i++ {
			if v := v.Field(i); v.CanSet() || t.Field(i).Name != "_" {
				var f visitFlag
				fieldType := t.Field(i)
				if fieldType.PkgPath != "" {
					// Unexported
					continue
				}

				tag := fieldType.Tag.Get(w.tag)
				if tag == "ignore" || tag == "-" {
					// Ignore this field
					continue
				}

				// Check if we implement includable and check it
				if include != nil {
					incl, err := include.HashInclude(fieldType.Name, v)
					if err != nil {
						return 0, err
					}
					if !incl {
						continue
					}
				}

				switch tag {
				case "set":
					f |= visitFlagSet
				}

				kh, err := w.visit(reflect.ValueOf(fieldType.Name), nil)
				if err != nil {
					return 0, err
				}

				vh, err := w.visit(v, &visitOpts{
					Flags:       f,
					Struct:      parent,
					StructField: fieldType.Name,
				})
				if err != nil {
					return 0, err
				}

				fieldHash := hashUpdateOrdered(w.h, kh, vh)
				h = hashUpdateUnordered(h, fieldHash)
			}
		}

		return h, nil

	case reflect.Slice:
		// We have two behaviors here. If it isn't a set, then we just
		// visit all the elements. If it is a set, then we do a deterministic
		// hash code.
		var h uint64
		var set bool
		if opts != nil {
			set = (opts.Flags & visitFlagSet) != 0
		}
		l := v.Len()
		for i := 0; i < l; i++ {
			current, err := w.visit(v.Index(i), nil)
			if err != nil {
				return 0, err
			}

			if set {
				h = hashUpdateUnordered(h, current)
			} else {
				h = hashUpdateOrdered(w.h, h, current)
			}
		}

		return h, nil

	case reflect.String:
		// Directly hash
		w.h.Reset()
		_, err := w.h.Write([]byte(v.String()))
		return w.h.Sum64(), err

	default:
		return 0, fmt.Errorf("unknown kind to hash: %s", k)
	}

	return 0, nil
}

func hashUpdateOrdered(h hash.Hash64, a, b uint64) uint64 {
	// For ordered updates, use a real hash function
	h.Reset()

	// We just panic if the binary writes fail because we are writing
	// an int64 which should never be fail-able.
	e1 := binary.Write(h, binary.LittleEndian, a)
	e2 := binary.Write(h, binary.LittleEndian, b)
	if e1 != nil {
		panic(e1)
	}
	if e2 != nil {
		panic(e2)
	}

	return h.Sum64()
}

func hashUpdateUnordered(a, b uint64) uint64 {
	return a ^ b
}

// visitFlag is used as a bitmask for affecting visit behavior
type visitFlag uint

const (
	visitFlagInvalid visitFlag = iota
	visitFlagSet               = iota << 1
)
<commit_msg>Make concurrency clear for #9<commit_after>package hashstructure

import (
	"encoding/binary"
	"fmt"
	"hash"
	"hash/fnv"
	"reflect"
)

// HashOptions are options that are available for hashing.
type HashOptions struct {
	// Hasher is the hash function to use. If this isn't set, it will
	// default to FNV.
	Hasher hash.Hash64

	// TagName is the struct tag to look at when hashing the structure.
	// By default this is "hash".
	TagName string

	// ZeroNil is flag determining if nil pointer should be treated equal
	// to a zero value of pointed type. By default this is false.
	ZeroNil bool
}

// Hash returns the hash value of an arbitrary value.
//
// If opts is nil, then default options will be used. See HashOptions
// for the default values. The same *HashOptions value cannot be used
// concurrently. None of the values within a *HashOptions struct are
// safe to read/write while hashing is being done.
//
// Notes on the value:
//
//   * Unexported fields on structs are ignored and do not affect the
//     hash value.
//
//   * Adding an exported field to a struct with the zero value will change
//     the hash value.
//
// For structs, the hashing can be controlled using tags. For example:
//
//    struct {
//        Name string
//        UUID string `hash:"ignore"`
//    }
//
// The available tag values are:
//
//   * "ignore" or "-" - The field will be ignored and not affect the hash code.
//
//   * "set" - The field will be treated as a set, where ordering doesn't
//             affect the hash code. This only works for slices.
//
func Hash(v interface{}, opts *HashOptions) (uint64, error) {
	// Create default options
	if opts == nil {
		opts = &HashOptions{}
	}
	if opts.Hasher == nil {
		opts.Hasher = fnv.New64()
	}
	if opts.TagName == "" {
		opts.TagName = "hash"
	}

	// Reset the hash
	opts.Hasher.Reset()

	// Create our walker and walk the structure
	w := &walker{
		h:       opts.Hasher,
		tag:     opts.TagName,
		zeronil: opts.ZeroNil,
	}
	return w.visit(reflect.ValueOf(v), nil)
}

type walker struct {
	h       hash.Hash64
	tag     string
	zeronil bool
}

type visitOpts struct {
	// Flags are a bitmask of flags to affect behavior of this visit
	Flags visitFlag

	// Information about the struct containing this field
	Struct      interface{}
	StructField string
}

func (w *walker) visit(v reflect.Value, opts *visitOpts) (uint64, error) {
	t := reflect.TypeOf(0)

	// Loop since these can be wrapped in multiple layers of pointers
	// and interfaces.
	for {
		// If we have an interface, dereference it. We have to do this up
		// here because it might be a nil in there and the check below must
		// catch that.
		if v.Kind() == reflect.Interface {
			v = v.Elem()
			continue
		}

		if v.Kind() == reflect.Ptr {
			if w.zeronil {
				t = v.Type().Elem()
			}
			v = reflect.Indirect(v)
			continue
		}

		break
	}

	// If it is nil, treat it like a zero.
	if !v.IsValid() {
		v = reflect.Zero(t)
	}

	// Binary writing can use raw ints, we have to convert to
	// a sized-int, we'll choose the largest...
	switch v.Kind() {
	case reflect.Int:
		v = reflect.ValueOf(int64(v.Int()))
	case reflect.Uint:
		v = reflect.ValueOf(uint64(v.Uint()))
	case reflect.Bool:
		var tmp int8
		if v.Bool() {
			tmp = 1
		}
		v = reflect.ValueOf(tmp)
	}

	k := v.Kind()

	// We can shortcut numeric values by directly binary writing them
	if k >= reflect.Int && k <= reflect.Complex64 {
		// A direct hash calculation
		w.h.Reset()
		err := binary.Write(w.h, binary.LittleEndian, v.Interface())
		return w.h.Sum64(), err
	}

	switch k {
	case reflect.Array:
		var h uint64
		l := v.Len()
		for i := 0; i < l; i++ {
			current, err := w.visit(v.Index(i), nil)
			if err != nil {
				return 0, err
			}

			h = hashUpdateOrdered(w.h, h, current)
		}

		return h, nil

	case reflect.Map:
		var includeMap IncludableMap
		if opts != nil && opts.Struct != nil {
			if v, ok := opts.Struct.(IncludableMap); ok {
				includeMap = v
			}
		}

		// Build the hash for the map. We do this by XOR-ing all the key
		// and value hashes. This makes it deterministic despite ordering.
		var h uint64
		for _, k := range v.MapKeys() {
			v := v.MapIndex(k)
			if includeMap != nil {
				incl, err := includeMap.HashIncludeMap(
					opts.StructField, k.Interface(), v.Interface())
				if err != nil {
					return 0, err
				}
				if !incl {
					continue
				}
			}

			kh, err := w.visit(k, nil)
			if err != nil {
				return 0, err
			}
			vh, err := w.visit(v, nil)
			if err != nil {
				return 0, err
			}

			fieldHash := hashUpdateOrdered(w.h, kh, vh)
			h = hashUpdateUnordered(h, fieldHash)
		}

		return h, nil

	case reflect.Struct:
		var include Includable
		parent := v.Interface()
		if impl, ok := parent.(Includable); ok {
			include = impl
		}

		t := v.Type()
		h, err := w.visit(reflect.ValueOf(t.Name()), nil)
		if err != nil {
			return 0, err
		}

		l := v.NumField()
		for i := 0; i < l; i++ {
			if v := v.Field(i); v.CanSet() || t.Field(i).Name != "_" {
				var f visitFlag
				fieldType := t.Field(i)
				if fieldType.PkgPath != "" {
					// Unexported
					continue
				}

				tag := fieldType.Tag.Get(w.tag)
				if tag == "ignore" || tag == "-" {
					// Ignore this field
					continue
				}

				// Check if we implement includable and check it
				if include != nil {
					incl, err := include.HashInclude(fieldType.Name, v)
					if err != nil {
						return 0, err
					}
					if !incl {
						continue
					}
				}

				switch tag {
				case "set":
					f |= visitFlagSet
				}

				kh, err := w.visit(reflect.ValueOf(fieldType.Name), nil)
				if err != nil {
					return 0, err
				}

				vh, err := w.visit(v, &visitOpts{
					Flags:       f,
					Struct:      parent,
					StructField: fieldType.Name,
				})
				if err != nil {
					return 0, err
				}

				fieldHash := hashUpdateOrdered(w.h, kh, vh)
				h = hashUpdateUnordered(h, fieldHash)
			}
		}

		return h, nil

	case reflect.Slice:
		// We have two behaviors here. If it isn't a set, then we just
		// visit all the elements. If it is a set, then we do a deterministic
		// hash code.
		var h uint64
		var set bool
		if opts != nil {
			set = (opts.Flags & visitFlagSet) != 0
		}
		l := v.Len()
		for i := 0; i < l; i++ {
			current, err := w.visit(v.Index(i), nil)
			if err != nil {
				return 0, err
			}

			if set {
				h = hashUpdateUnordered(h, current)
			} else {
				h = hashUpdateOrdered(w.h, h, current)
			}
		}

		return h, nil

	case reflect.String:
		// Directly hash
		w.h.Reset()
		_, err := w.h.Write([]byte(v.String()))
		return w.h.Sum64(), err

	default:
		return 0, fmt.Errorf("unknown kind to hash: %s", k)
	}

	return 0, nil
}

func hashUpdateOrdered(h hash.Hash64, a, b uint64) uint64 {
	// For ordered updates, use a real hash function
	h.Reset()

	// We just panic if the binary writes fail because we are writing
	// an int64 which should never be fail-able.
	e1 := binary.Write(h, binary.LittleEndian, a)
	e2 := binary.Write(h, binary.LittleEndian, b)
	if e1 != nil {
		panic(e1)
	}
	if e2 != nil {
		panic(e2)
	}

	return h.Sum64()
}

func hashUpdateUnordered(a, b uint64) uint64 {
	return a ^ b
}

// visitFlag is used as a bitmask for affecting visit behavior
type visitFlag uint

const (
	visitFlagInvalid visitFlag = iota
	visitFlagSet               = iota << 1
)
<|endoftext|>


<commit_before>package lib

// Version is the github-nippou version
const Version = "4.1.2"
<commit_msg>Bump version to 4.1.3<commit_after>package lib

// Version is the github-nippou version
const Version = "4.1.3"
<|endoftext|>


<commit_before>package light

import (
	"time"

	"github.com/pdf/golifx"
	"github.com/pdf/golifx/common"
	"github.com/pdf/golifx/protocol"
)

// Color is given in the degrees of a circle (see http://www.workwithcolor.com/hsl-color-picker-01.htm).
type Color uint16

const (
	Red       Color = 0 // or 360
	Yellow    Color = 60
	Green     Color = 120
	Turquoise Color = 180
	Blue      Color = 240
	Pink      Color = 300
)

type light struct {
	client *golifx.Client
}

func newLight() (*light, error) {
	// Get debug output for LIFX device
	//logger := logrus.New()
	//logger.Out = os.Stderr
	//logger.Level = logrus.DebugLevel
	//golifx.SetLogger(logger)

	client, err := golifx.NewClient(&protocol.V2{Reliable: true})
	if err != nil {
		return nil, err
	}
	client.SetDiscoveryInterval(5 * time.Minute)

	return &light{client}, nil
}

func (l *light) setColor(c Color) {
	color := common.Color{
		Hue:        65535 / 360 * (uint16(c) % 360),
		Saturation: 65535,
		Brightness: 26214,
		Kelvin:     2500,
	}
	l.client.SetColor(color, 1*time.Second)
}

func (l *light) setPower(p bool) {
	l.client.SetPower(p)
}

func (l *light) turnOff() {
	l.setPower(false)
}

func (l *light) turnOn() {
	l.setPower(true)
}
<commit_msg>Test using light instead of client to set power and color.<commit_after>package light

import (
	"time"

	"github.com/pdf/golifx"
	"github.com/pdf/golifx/common"
	"github.com/pdf/golifx/protocol"
)

// Color is given in the degrees of a circle (see http://www.workwithcolor.com/hsl-color-picker-01.htm).
type Color uint16

const (
	Red       Color = 0 // or 360
	Yellow    Color = 60
	Green     Color = 120
	Turquoise Color = 180
	Blue      Color = 240
	Pink      Color = 300
)

type light struct {
	client *golifx.Client
	device common.Light
}

func newLight() (*light, error) {
	// Get debug output for LIFX device
	//logger := logrus.New()
	//logger.Out = os.Stderr
	//logger.Level = logrus.DebugLevel
	//golifx.SetLogger(logger)

	client, err := golifx.NewClient(&protocol.V2{Reliable: true})
	if err != nil {
		return nil, err
	}
	client.SetDiscoveryInterval(5 * time.Minute)

	device, err := client.GetLightByLabel("BuildBulb")
	if err != nil {
		return nil, err
	}
	return &light{client, device}, nil
}

func (l *light) setColor(c Color) {
	color := common.Color{
		Hue:        65535 / 360 * (uint16(c) % 360),
		Saturation: 65535,
		Brightness: 26214,
		Kelvin:     2500,
	}
	l.device.SetColor(color, 1*time.Second)
	l.client.SetColor(color, 1*time.Second)
}

func (l *light) setPower(p bool) {
	l.device.SetPower(p)
	l.client.SetPower(p)
}

func (l *light) turnOff() {
	l.setPower(false)
}

func (l *light) turnOn() {
	l.setPower(true)
}
<|endoftext|>


<commit_before>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package sync_test

import (
	. "sync"

	"runtime"
	"testing"
	"time"
)

func TestCondSignal(t *testing.T) {
	var m Mutex
	c := NewCond(&m)
	n := 2
	running := make(chan bool, n)
	awake := make(chan bool, n)
	for i := 0; i < n; i++ {
		go func() {
			m.Lock()
			running <- true
			c.Wait()
			awake <- true
			m.Unlock()
		}()
	}
	for i := 0; i < n; i++ {
		<-running // Wait for everyone to run.
	}
	for n > 0 {
		select {
		case <-awake:
			t.Fatal("goroutine not asleep")
		default:
		}
		m.Lock()
		c.Signal()
		m.Unlock()
		<-awake // Will deadlock if no goroutine wakes up
		select {
		case <-awake:
			t.Fatal("too many goroutines awake")
		default:
		}
		n--
	}
	c.Signal()
}

func TestCondSignalGenerations(t *testing.T) {
	var m Mutex
	c := NewCond(&m)
	n := 100
	running := make(chan bool, n)
	awake := make(chan int, n)
	for i := 0; i < n; i++ {
		go func(i int) {
			m.Lock()
			running <- true
			c.Wait()
			awake <- i
			m.Unlock()
		}(i)
		if i > 0 {
			a := <-awake
			if a != i-1 {
				t.Fatalf("wrong goroutine woke up: want %d, got %d", i-1, a)
			}
		}
		<-running
		m.Lock()
		c.Signal()
		m.Unlock()
	}
}

func TestCondBroadcast(t *testing.T) {
	var m Mutex
	c := NewCond(&m)
	n := 200
	running := make(chan int, n)
	awake := make(chan int, n)
	exit := false
	for i := 0; i < n; i++ {
		go func(g int) {
			m.Lock()
			for !exit {
				running <- g
				c.Wait()
				awake <- g
			}
			m.Unlock()
		}(i)
	}
	for i := 0; i < n; i++ {
		for i := 0; i < n; i++ {
			<-running // Will deadlock unless n are running.
		}
		if i == n-1 {
			m.Lock()
			exit = true
			m.Unlock()
		}
		select {
		case <-awake:
			t.Fatal("goroutine not asleep")
		default:
		}
		m.Lock()
		c.Broadcast()
		m.Unlock()
		seen := make([]bool, n)
		for i := 0; i < n; i++ {
			g := <-awake
			if seen[g] {
				t.Fatal("goroutine woke up twice")
			}
			seen[g] = true
		}
	}
	select {
	case <-running:
		t.Fatal("goroutine did not exit")
	default:
	}
	c.Broadcast()
}

func TestRace(t *testing.T) {
	x := 0
	c := NewCond(&Mutex{})
	done := make(chan bool)
	go func() {
		c.L.Lock()
		x = 1
		c.Wait()
		if x != 2 {
			t.Error("want 2")
		}
		x = 3
		c.Signal()
		c.L.Unlock()
		done <- true
	}()
	go func() {
		c.L.Lock()
		for {
			if x == 1 {
				x = 2
				c.Signal()
				break
			}
			c.L.Unlock()
			runtime.Gosched()
			c.L.Lock()
		}
		c.L.Unlock()
		done <- true
	}()
	go func() {
		c.L.Lock()
		for {
			if x == 2 {
				c.Wait()
				if x != 3 {
					t.Error("want 3")
				}
				break
			}
			if x == 3 {
				break
			}
			c.L.Unlock()
			runtime.Gosched()
			c.L.Lock()
		}
		c.L.Unlock()
		done <- true
	}()
	<-done
	<-done
	<-done
}

func TestCondSignalStealing(t *testing.T) {
	for iters := 0; iters < 1000; iters++ {
		var m Mutex
		cond := NewCond(&m)

		// Start a waiter.
		ch := make(chan struct{})
		go func() {
			m.Lock()
			ch <- struct{}{}
			cond.Wait()
			m.Unlock()

			ch <- struct{}{}
		}()

		<-ch
		m.Lock()
		m.Unlock()

		// We know that the waiter is in the cond.Wait() call because we
		// synchronized with it, then acquired/released the mutex it was
		// holding when we synchronized.
		//
		// Start two goroutines that will race: one will broadcast on
		// the cond var, the other will wait on it.
		//
		// The new waiter may or may not get notified, but the first one
		// has to be notified.
		done := false
		go func() {
			cond.Broadcast()
		}()

		go func() {
			m.Lock()
			for !done {
				cond.Wait()
			}
			m.Unlock()
		}()

		// Check that the first waiter does get signaled.
		select {
		case <-ch:
		case <-time.After(2 * time.Second):
			t.Fatalf("First waiter didn't get broadcast.")
		}

		// Release the second waiter in case it didn't get the
		// broadcast.
		m.Lock()
		done = true
		m.Unlock()
		cond.Broadcast()
	}
}

func TestCondCopy(t *testing.T) {
	defer func() {
		err := recover()
		if err == nil || err.(string) != "sync.Cond is copied" {
			t.Fatalf("got %v, expect sync.Cond is copied", err)
		}
	}()
	c := Cond{L: &Mutex{}}
	c.Signal()
	c2 := c
	c2.Signal()
}

func BenchmarkCond1(b *testing.B) {
	benchmarkCond(b, 1)
}

func BenchmarkCond2(b *testing.B) {
	benchmarkCond(b, 2)
}

func BenchmarkCond4(b *testing.B) {
	benchmarkCond(b, 4)
}

func BenchmarkCond8(b *testing.B) {
	benchmarkCond(b, 8)
}

func BenchmarkCond16(b *testing.B) {
	benchmarkCond(b, 16)
}

func BenchmarkCond32(b *testing.B) {
	benchmarkCond(b, 32)
}

func benchmarkCond(b *testing.B, waiters int) {
	c := NewCond(&Mutex{})
	done := make(chan bool)
	id := 0

	for routine := 0; routine < waiters+1; routine++ {
		go func() {
			for i := 0; i < b.N; i++ {
				c.L.Lock()
				if id == -1 {
					c.L.Unlock()
					break
				}
				id++
				if id == waiters+1 {
					id = 0
					c.Broadcast()
				} else {
					c.Wait()
				}
				c.L.Unlock()
			}
			c.L.Lock()
			id = -1
			c.Broadcast()
			c.L.Unlock()
			done <- true
		}()
	}
	for routine := 0; routine < waiters+1; routine++ {
		<-done
	}
}
<commit_msg>sync: hide test of misuse of Cond from vet<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package sync_test

import (
	"reflect"
	"runtime"
	. "sync"
	"testing"
	"time"
)

func TestCondSignal(t *testing.T) {
	var m Mutex
	c := NewCond(&m)
	n := 2
	running := make(chan bool, n)
	awake := make(chan bool, n)
	for i := 0; i < n; i++ {
		go func() {
			m.Lock()
			running <- true
			c.Wait()
			awake <- true
			m.Unlock()
		}()
	}
	for i := 0; i < n; i++ {
		<-running // Wait for everyone to run.
	}
	for n > 0 {
		select {
		case <-awake:
			t.Fatal("goroutine not asleep")
		default:
		}
		m.Lock()
		c.Signal()
		m.Unlock()
		<-awake // Will deadlock if no goroutine wakes up
		select {
		case <-awake:
			t.Fatal("too many goroutines awake")
		default:
		}
		n--
	}
	c.Signal()
}

func TestCondSignalGenerations(t *testing.T) {
	var m Mutex
	c := NewCond(&m)
	n := 100
	running := make(chan bool, n)
	awake := make(chan int, n)
	for i := 0; i < n; i++ {
		go func(i int) {
			m.Lock()
			running <- true
			c.Wait()
			awake <- i
			m.Unlock()
		}(i)
		if i > 0 {
			a := <-awake
			if a != i-1 {
				t.Fatalf("wrong goroutine woke up: want %d, got %d", i-1, a)
			}
		}
		<-running
		m.Lock()
		c.Signal()
		m.Unlock()
	}
}

func TestCondBroadcast(t *testing.T) {
	var m Mutex
	c := NewCond(&m)
	n := 200
	running := make(chan int, n)
	awake := make(chan int, n)
	exit := false
	for i := 0; i < n; i++ {
		go func(g int) {
			m.Lock()
			for !exit {
				running <- g
				c.Wait()
				awake <- g
			}
			m.Unlock()
		}(i)
	}
	for i := 0; i < n; i++ {
		for i := 0; i < n; i++ {
			<-running // Will deadlock unless n are running.
		}
		if i == n-1 {
			m.Lock()
			exit = true
			m.Unlock()
		}
		select {
		case <-awake:
			t.Fatal("goroutine not asleep")
		default:
		}
		m.Lock()
		c.Broadcast()
		m.Unlock()
		seen := make([]bool, n)
		for i := 0; i < n; i++ {
			g := <-awake
			if seen[g] {
				t.Fatal("goroutine woke up twice")
			}
			seen[g] = true
		}
	}
	select {
	case <-running:
		t.Fatal("goroutine did not exit")
	default:
	}
	c.Broadcast()
}

func TestRace(t *testing.T) {
	x := 0
	c := NewCond(&Mutex{})
	done := make(chan bool)
	go func() {
		c.L.Lock()
		x = 1
		c.Wait()
		if x != 2 {
			t.Error("want 2")
		}
		x = 3
		c.Signal()
		c.L.Unlock()
		done <- true
	}()
	go func() {
		c.L.Lock()
		for {
			if x == 1 {
				x = 2
				c.Signal()
				break
			}
			c.L.Unlock()
			runtime.Gosched()
			c.L.Lock()
		}
		c.L.Unlock()
		done <- true
	}()
	go func() {
		c.L.Lock()
		for {
			if x == 2 {
				c.Wait()
				if x != 3 {
					t.Error("want 3")
				}
				break
			}
			if x == 3 {
				break
			}
			c.L.Unlock()
			runtime.Gosched()
			c.L.Lock()
		}
		c.L.Unlock()
		done <- true
	}()
	<-done
	<-done
	<-done
}

func TestCondSignalStealing(t *testing.T) {
	for iters := 0; iters < 1000; iters++ {
		var m Mutex
		cond := NewCond(&m)

		// Start a waiter.
		ch := make(chan struct{})
		go func() {
			m.Lock()
			ch <- struct{}{}
			cond.Wait()
			m.Unlock()

			ch <- struct{}{}
		}()

		<-ch
		m.Lock()
		m.Unlock()

		// We know that the waiter is in the cond.Wait() call because we
		// synchronized with it, then acquired/released the mutex it was
		// holding when we synchronized.
		//
		// Start two goroutines that will race: one will broadcast on
		// the cond var, the other will wait on it.
		//
		// The new waiter may or may not get notified, but the first one
		// has to be notified.
		done := false
		go func() {
			cond.Broadcast()
		}()

		go func() {
			m.Lock()
			for !done {
				cond.Wait()
			}
			m.Unlock()
		}()

		// Check that the first waiter does get signaled.
		select {
		case <-ch:
		case <-time.After(2 * time.Second):
			t.Fatalf("First waiter didn't get broadcast.")
		}

		// Release the second waiter in case it didn't get the
		// broadcast.
		m.Lock()
		done = true
		m.Unlock()
		cond.Broadcast()
	}
}

func TestCondCopy(t *testing.T) {
	defer func() {
		err := recover()
		if err == nil || err.(string) != "sync.Cond is copied" {
			t.Fatalf("got %v, expect sync.Cond is copied", err)
		}
	}()
	c := Cond{L: &Mutex{}}
	c.Signal()
	var c2 Cond
	reflect.ValueOf(&c2).Elem().Set(reflect.ValueOf(&c).Elem()) // c2 := c, hidden from vet
	c2.Signal()
}

func BenchmarkCond1(b *testing.B) {
	benchmarkCond(b, 1)
}

func BenchmarkCond2(b *testing.B) {
	benchmarkCond(b, 2)
}

func BenchmarkCond4(b *testing.B) {
	benchmarkCond(b, 4)
}

func BenchmarkCond8(b *testing.B) {
	benchmarkCond(b, 8)
}

func BenchmarkCond16(b *testing.B) {
	benchmarkCond(b, 16)
}

func BenchmarkCond32(b *testing.B) {
	benchmarkCond(b, 32)
}

func benchmarkCond(b *testing.B, waiters int) {
	c := NewCond(&Mutex{})
	done := make(chan bool)
	id := 0

	for routine := 0; routine < waiters+1; routine++ {
		go func() {
			for i := 0; i < b.N; i++ {
				c.L.Lock()
				if id == -1 {
					c.L.Unlock()
					break
				}
				id++
				if id == waiters+1 {
					id = 0
					c.Broadcast()
				} else {
					c.Wait()
				}
				c.L.Unlock()
			}
			c.L.Lock()
			id = -1
			c.Broadcast()
			c.L.Unlock()
			done <- true
		}()
	}
	for routine := 0; routine < waiters+1; routine++ {
		<-done
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
	"io/ioutil"
	"net/http"
	"runtime"
	"strconv"
	"time"
)

var (
	//stress
	stress = kingpin.Command("stress", "Run predefined load of requests.").Alias("s")

	//stress flags
	stressCount       = stress.Flag("num", "Number of requests to make.").Short('n').Default("1").Int()
	stressConcurrency = stress.Flag("concurrent", "Number of multiple requests to make.").Short('c').Default("1").Int()

	//request flags
	stressTimeout   = stress.Flag("timeout", "Maximum seconds to wait for response").Short('t').Default("10s").Duration()
	stressReqMethod = stress.Flag("request-method", "Request type. GET, HEAD, POST, PUT, etc.").Short('X').Default("GET").String()
	stressReqBody   = stress.Flag("body", "String to use as request body e.g. POST body.").String()
	stressHeaders   = HTTPHeader(stress.Flag("header", "Add arbitrary header line, eg. 'Accept-Encoding:gzip'").Short('H'))
	stressUserAgent = stress.Flag("user-agent", "Add User-Agent header.").Short('A').Default("pewpew").String()
	stressBasicAuth = BasicAuth(stress.Flag("basic-auth", "Add HTTP basic authentication, eg. 'user123:password456'"))
	stressHttp2     = stress.Flag("http2", "Use HTTP2.").Bool()
	stressUrl       = stress.Arg("url", "URL to stress, formatted http[s]://hostname[:port][/path]").String()

	//global flags
	verbose  = kingpin.Flag("verbose", "Print extra troubleshooting info").Short('v').Bool()
	cpuCount = kingpin.Flag("cpu", "Number of CPUs to use.").Default(strconv.Itoa(runtime.GOMAXPROCS(0))).Int()
)

func main() {
	kingpin.CommandLine.Help = "HTTP(S) & HTTP2 load tester for performance and stress testing"
	kingpin.CommandLine.HelpFlag.Short('h')

	parseArgs := kingpin.Parse()

	runtime.GOMAXPROCS(*cpuCount)
	if *cpuCount < 1 {
		kingpin.Fatalf("CPU count must be greater or equal to 1")
	}

	switch parseArgs {
	case "stress":
		kingpin.FatalIfError(runStress(), "stress failed")
	}
}

type workerDone struct{}

type requestStat struct {
	duration int64 //nanoseconds
}
type requestStatSummary struct {
	avgQps      float64 //per nanoseconds
	avgDuration int64   //nanoseconds
	maxDuration int64   //nanoseconds
	minDuration int64   //nanoseconds
}

func runStress() error {
	//checks
	if *stressUrl == "" {
		return errors.New("needs URL")
	}
	if *stressCount <= 0 {
		return errors.New("number of requests must be one or more")
	}
	if *stressConcurrency <= 0 {
		return errors.New("concurrency must be one or more")
	}
	if *stressTimeout < 0 {
		return errors.New("timeout must be zero or more")
	}
	if *stressConcurrency > *stressCount {
		return errors.New("concurrency must be higher than number of requests")
	}

	fmt.Println("Stress testing " + *stressUrl + "...")
	fmt.Printf("Running %d tests, %d at a time\n", *stressCount, *stressConcurrency)

	//setup the request
	var req *http.Request
	var err error
	if *stressReqBody != "" {
		req, err = http.NewRequest(*stressReqMethod, *stressUrl, bytes.NewBuffer([]byte(*stressReqBody)))
	} else {
		req, err = http.NewRequest(*stressReqMethod, *stressUrl, nil)
	}
	if err != nil {
		return errors.New("failed to create request: " + err.Error())
	}
	req.Header = *stressHeaders //add headers
	req.Header.Set("User-Agent", *stressUserAgent)
	if (*stressBasicAuth).String() != "" {
		req.SetBasicAuth((*stressBasicAuth).User, (*stressBasicAuth).Password)
	}

	//setup the queue of requests
	requestChan := make(chan *http.Request, *stressCount)
	for i := 0; i < *stressCount; i++ {
		requestChan <- req
	}
	close(requestChan)

	workerDoneChan := make(chan workerDone)   //workers use this to indicate they are done
	requestStatChan := make(chan requestStat) //workers communicate each requests' info

	//workers
	totalStartTime := time.Now()
	var totalEndTime time.Time
	for i := 0; i < *stressConcurrency; i++ {
		//TODO handle the returned errors from this
		go func() error {
			tr := &http.Transport{}
			if !*stressHttp2 {
				nilMap := make(map[string](func(authority string, c *tls.Conn) http.RoundTripper))
				tr = &http.Transport{TLSNextProto: nilMap}
			}
			client := &http.Client{Timeout: time.Duration(*stressTimeout) * time.Second, Transport: tr}
			for {
				select {
				case req, ok := <-requestChan:
					if !ok {
						workerDoneChan <- workerDone{}
						return nil
					}
					//run the acutal request
					reqStartTime := time.Now()
					response, err := client.Do((*http.Request)(req))
					reqEndTime := time.Now()
					if err != nil {
						return errors.New("Failed to make request:" + err.Error())
					}
					reqTimeNs := (reqEndTime.UnixNano() - reqStartTime.UnixNano())

					var requestData string
					if *verbose {
						requestData = "----Request----\n\n"

						//request timing
						requestData = requestData + fmt.Sprintf("Request took %dms\n\n", reqTimeNs/1000000)

						//reponse metadata
						requestData = requestData + fmt.Sprintf("Response:\n%+v\n\n", *response)

						//reponse body
						defer response.Body.Close()
						body, err := ioutil.ReadAll(response.Body)
						if err != nil {
							return errors.New("Failed to read response body:" + err.Error())
						}
						requestData = requestData + fmt.Sprintf("Body:\n%s\n\n", body)
					}

					if requestData != "" {
						fmt.Print(requestData)
					}
					requestStatChan <- requestStat{duration: reqTimeNs}
				}
			}
		}()
	}

	allRequestStats := make([]requestStat, *stressCount)
	requestsCompleteCount := 0
	workersDoneCount := 0
	//wait for all workers to finish
WorkerLoop:
	for {
		select {
		case <-workerDoneChan:
			workersDoneCount++
			if workersDoneCount == *stressConcurrency {
				//all workers are done
				totalEndTime = time.Now()
				break WorkerLoop
			}
		case requestStat := <-requestStatChan:
			allRequestStats[requestsCompleteCount] = requestStat
			requestsCompleteCount++
		}
	}

	fmt.Println("----Summary----\n")

	//info about the request
	fmt.Println("Method: " + req.Method)
	fmt.Println("Host: " + req.Host)

	totalTimeNs := totalEndTime.UnixNano() - totalStartTime.UnixNano()
	reqStats := createRequestsStats(allRequestStats, totalTimeNs)
	fmt.Println(createTextSummary(reqStats, totalTimeNs))

	return nil
}

//create statistical summary of all requests
func createRequestsStats(requestStats []requestStat, totalTimeNs int64) requestStatSummary {
	if len(requestStats) == 0 {
		return requestStatSummary{}
	}

	summary := requestStatSummary{maxDuration: requestStats[0].duration, minDuration: requestStats[0].duration}
	var totalDurations int64
	totalDurations = 0 //total time of all requests (concurrent is counted)
	for i := 0; i < len(requestStats); i++ {
		if requestStats[i].duration > summary.maxDuration {
			summary.maxDuration = requestStats[i].duration
		}
		if requestStats[i].duration < summary.minDuration {
			summary.minDuration = requestStats[i].duration
		}
		totalDurations += requestStats[i].duration
	}
	summary.avgDuration = totalDurations / int64(len(requestStats))
	summary.avgQps = float64(len(requestStats)) / float64(totalTimeNs)
	return summary
}

//creates nice readable summary of entire stress test
func createTextSummary(reqStatSummary requestStatSummary, totalTimeNs int64) string {
	summary := "\n"

	summary = summary + "Runtime Statistics:\n"
	summary = summary + "Total time:  " + strconv.Itoa(int(totalTimeNs/1000000)) + " ms\n"
	summary = summary + "Mean QPS:    " + fmt.Sprintf("%.2f", reqStatSummary.avgQps*1000000000) + " req/sec\n"

	summary = summary + "\nQuery Statistics\n"
	summary = summary + "Mean query:     " + strconv.Itoa(int(reqStatSummary.avgDuration/1000000)) + " ms\n"
	summary = summary + "Fastest query:  " + strconv.Itoa(int(reqStatSummary.minDuration/1000000)) + " ms\n"
	summary = summary + "Slowest query:  " + strconv.Itoa(int(reqStatSummary.maxDuration/1000000)) + " ms\n"
	return summary
}
<commit_msg>Nonfunctional flag cleanup<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"errors"
	"fmt"
	kingpin "gopkg.in/alecthomas/kingpin.v2"
	"io/ioutil"
	"net/http"
	"runtime"
	"strconv"
	"time"
)

var (
	//stress
	stress = kingpin.Command("stress", "Run predefined load of requests.").Alias("s")

	//stress flags
	stressCount       = stress.Flag("num", "Number of requests to make.").Short('n').Default("1").Int()
	stressConcurrency = stress.Flag("concurrent", "Number of multiple requests to make.").Short('c').Default("1").Int()

	//request flags
	stressTimeout   = stress.Flag("timeout", "Maximum seconds to wait for response").Short('t').Default("10s").Duration()
	stressReqMethod = stress.Flag("request-method", "Request type. GET, HEAD, POST, PUT, etc.").Short('X').Default("GET").String()
	stressReqBody   = stress.Flag("body", "String to use as request body e.g. POST body.").String()
	stressHeaders   = HTTPHeader(stress.Flag("header", "Add arbitrary header line, eg. 'Accept-Encoding:gzip'").Short('H'))
	stressUserAgent = stress.Flag("user-agent", "Add User-Agent header.").Short('A').Default("pewpew").String()
	stressBasicAuth = BasicAuth(stress.Flag("basic-auth", "Add HTTP basic authentication, eg. 'user123:password456'"))
	stressHttp2     = stress.Flag("http2", "Use HTTP2.").Bool()

	//url
	stressUrl = stress.Arg("url", "URL to stress, formatted http[s]://hostname[:port][/path]").String()

	//global flags
	verbose  = kingpin.Flag("verbose", "Print extra troubleshooting info").Short('v').Bool()
	cpuCount = kingpin.Flag("cpu", "Number of CPUs to use.").Default(strconv.Itoa(runtime.GOMAXPROCS(0))).Int()
)

func main() {
	kingpin.CommandLine.Help = "HTTP(S) & HTTP2 load tester for performance and stress testing"
	kingpin.CommandLine.HelpFlag.Short('h')

	parseArgs := kingpin.Parse()

	runtime.GOMAXPROCS(*cpuCount)
	if *cpuCount < 1 {
		kingpin.Fatalf("CPU count must be greater or equal to 1")
	}

	switch parseArgs {
	case "stress":
		kingpin.FatalIfError(runStress(), "stress failed")
	}
}

type workerDone struct{}

type requestStat struct {
	duration int64 //nanoseconds
}
type requestStatSummary struct {
	avgQps      float64 //per nanoseconds
	avgDuration int64   //nanoseconds
	maxDuration int64   //nanoseconds
	minDuration int64   //nanoseconds
}

func runStress() error {
	//checks
	if *stressUrl == "" {
		return errors.New("needs URL")
	}
	if *stressCount <= 0 {
		return errors.New("number of requests must be one or more")
	}
	if *stressConcurrency <= 0 {
		return errors.New("concurrency must be one or more")
	}
	if *stressTimeout < 0 {
		return errors.New("timeout must be zero or more")
	}
	if *stressConcurrency > *stressCount {
		return errors.New("concurrency must be higher than number of requests")
	}

	fmt.Println("Stress testing " + *stressUrl + "...")
	fmt.Printf("Running %d tests, %d at a time\n", *stressCount, *stressConcurrency)

	//setup the request
	var req *http.Request
	var err error
	if *stressReqBody != "" {
		req, err = http.NewRequest(*stressReqMethod, *stressUrl, bytes.NewBuffer([]byte(*stressReqBody)))
	} else {
		req, err = http.NewRequest(*stressReqMethod, *stressUrl, nil)
	}
	if err != nil {
		return errors.New("failed to create request: " + err.Error())
	}
	req.Header = *stressHeaders //add headers
	req.Header.Set("User-Agent", *stressUserAgent)
	if (*stressBasicAuth).String() != "" {
		req.SetBasicAuth((*stressBasicAuth).User, (*stressBasicAuth).Password)
	}

	//setup the queue of requests
	requestChan := make(chan *http.Request, *stressCount)
	for i := 0; i < *stressCount; i++ {
		requestChan <- req
	}
	close(requestChan)

	workerDoneChan := make(chan workerDone)   //workers use this to indicate they are done
	requestStatChan := make(chan requestStat) //workers communicate each requests' info

	//workers
	totalStartTime := time.Now()
	var totalEndTime time.Time
	for i := 0; i < *stressConcurrency; i++ {
		//TODO handle the returned errors from this
		go func() error {
			tr := &http.Transport{}
			if !*stressHttp2 {
				nilMap := make(map[string](func(authority string, c *tls.Conn) http.RoundTripper))
				tr = &http.Transport{TLSNextProto: nilMap}
			}
			client := &http.Client{Timeout: time.Duration(*stressTimeout) * time.Second, Transport: tr}
			for {
				select {
				case req, ok := <-requestChan:
					if !ok {
						workerDoneChan <- workerDone{}
						return nil
					}
					//run the acutal request
					reqStartTime := time.Now()
					response, err := client.Do((*http.Request)(req))
					reqEndTime := time.Now()
					if err != nil {
						return errors.New("Failed to make request:" + err.Error())
					}
					reqTimeNs := (reqEndTime.UnixNano() - reqStartTime.UnixNano())

					var requestData string
					if *verbose {
						requestData = "----Request----\n\n"

						//request timing
						requestData = requestData + fmt.Sprintf("Request took %dms\n\n", reqTimeNs/1000000)

						//reponse metadata
						requestData = requestData + fmt.Sprintf("Response:\n%+v\n\n", *response)

						//reponse body
						defer response.Body.Close()
						body, err := ioutil.ReadAll(response.Body)
						if err != nil {
							return errors.New("Failed to read response body:" + err.Error())
						}
						requestData = requestData + fmt.Sprintf("Body:\n%s\n\n", body)
					}

					if requestData != "" {
						fmt.Print(requestData)
					}
					requestStatChan <- requestStat{duration: reqTimeNs}
				}
			}
		}()
	}

	allRequestStats := make([]requestStat, *stressCount)
	requestsCompleteCount := 0
	workersDoneCount := 0
	//wait for all workers to finish
WorkerLoop:
	for {
		select {
		case <-workerDoneChan:
			workersDoneCount++
			if workersDoneCount == *stressConcurrency {
				//all workers are done
				totalEndTime = time.Now()
				break WorkerLoop
			}
		case requestStat := <-requestStatChan:
			allRequestStats[requestsCompleteCount] = requestStat
			requestsCompleteCount++
		}
	}

	fmt.Println("----Summary----\n")

	//info about the request
	fmt.Println("Method: " + req.Method)
	fmt.Println("Host: " + req.Host)

	totalTimeNs := totalEndTime.UnixNano() - totalStartTime.UnixNano()
	reqStats := createRequestsStats(allRequestStats, totalTimeNs)
	fmt.Println(createTextSummary(reqStats, totalTimeNs))

	return nil
}

//create statistical summary of all requests
func createRequestsStats(requestStats []requestStat, totalTimeNs int64) requestStatSummary {
	if len(requestStats) == 0 {
		return requestStatSummary{}
	}

	summary := requestStatSummary{maxDuration: requestStats[0].duration, minDuration: requestStats[0].duration}
	var totalDurations int64
	totalDurations = 0 //total time of all requests (concurrent is counted)
	for i := 0; i < len(requestStats); i++ {
		if requestStats[i].duration > summary.maxDuration {
			summary.maxDuration = requestStats[i].duration
		}
		if requestStats[i].duration < summary.minDuration {
			summary.minDuration = requestStats[i].duration
		}
		totalDurations += requestStats[i].duration
	}
	summary.avgDuration = totalDurations / int64(len(requestStats))
	summary.avgQps = float64(len(requestStats)) / float64(totalTimeNs)
	return summary
}

//creates nice readable summary of entire stress test
func createTextSummary(reqStatSummary requestStatSummary, totalTimeNs int64) string {
	summary := "\n"

	summary = summary + "Runtime Statistics:\n"
	summary = summary + "Total time:  " + strconv.Itoa(int(totalTimeNs/1000000)) + " ms\n"
	summary = summary + "Mean QPS:    " + fmt.Sprintf("%.2f", reqStatSummary.avgQps*1000000000) + " req/sec\n"

	summary = summary + "\nQuery Statistics\n"
	summary = summary + "Mean query:     " + strconv.Itoa(int(reqStatSummary.avgDuration/1000000)) + " ms\n"
	summary = summary + "Fastest query:  " + strconv.Itoa(int(reqStatSummary.minDuration/1000000)) + " ms\n"
	summary = summary + "Slowest query:  " + strconv.Itoa(int(reqStatSummary.maxDuration/1000000)) + " ms\n"
	return summary
}
<|endoftext|>