<commit_before>package equalfile

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"os"
)

// Only the first 10^10 bytes of io.Reader are compared.  Ignored when using io.LimitedReader
const defaultMaxSize = 10000000000
const defaultBufSize = 20000

type Options struct {
	Debug         bool // enable debugging to stdout
	ForceFileRead bool // prevent shortcut at filesystem level (link, pathname, etc)

	// MaxSize is a safely limit to prevent forever reading from an infinite
	// reader.  If left unset, will default to 1OGBytes. Ignored when
	// CompareReader() is given one or more io.LimitedReader.
	MaxSize int64
}

type Cmp struct {
	Opt Options

	readCount int
	readMin   int
	readMax   int
	readSum   int64

	hashType         hash.Hash
	hashMatchCompare bool
	hashTable        map[string]hashSum

	buf []byte
}

type hashSum struct {
	result []byte
	err    error
}

// New creates Cmp for multiple comparison mode.
func NewMultiple(buf []byte, options Options, h hash.Hash, compareOnMatch bool) *Cmp {
	c := &Cmp{
		Opt:              options,
		hashType:         h,
		hashMatchCompare: compareOnMatch,
		hashTable:        map[string]hashSum{},
		buf:              buf,
	}
	if c.buf == nil || len(c.buf) == 0 {
		c.buf = make([]byte, defaultBufSize)
	}
	c.debugf("New: bufSize=%d\n", len(c.buf))
	return c
}

// New creates Cmp for single comparison mode.
func New(buf []byte, options Options) *Cmp {
	return NewMultiple(buf, options, nil, true)
}

func (c *Cmp) getHash(path string, maxSize int64) ([]byte, error) {
	h, found := c.hashTable[path]
	if found {
		return h.result, h.err
	}

	f, openErr := os.Open(path)
	if openErr != nil {
		return nil, openErr
	}
	defer f.Close()

	sum := make([]byte, c.hashType.Size())
	c.hashType.Reset()
	n, copyErr := io.CopyN(c.hashType, f, maxSize)
	copy(sum, c.hashType.Sum(nil))

	if copyErr == io.EOF && n < maxSize {
		copyErr = nil
	}

	return c.newHash(path, sum, copyErr)
}

func (c *Cmp) newHash(path string, sum []byte, e error) ([]byte, error) {

	c.hashTable[path] = hashSum{sum, e}

	c.debugf("newHash[%s]=%v: error=[%v]\n", path, hex.EncodeToString(sum), e)

	return sum, e
}

func (c *Cmp) multipleMode() bool {
	return c.hashType != nil
}

// CompareFile verifies that files with names path1, path2 have same contents.
func (c *Cmp) CompareFile(path1, path2 string) (bool, error) {

	if c.Opt.MaxSize < 0 {
		return false, fmt.Errorf("negative MaxSize")
	}

	r1, openErr1 := os.Open(path1)
	if openErr1 != nil {
		return false, openErr1
	}
	defer r1.Close()
	info1, statErr1 := r1.Stat()
	if statErr1 != nil {
		return false, statErr1
	}

	r2, openErr2 := os.Open(path2)
	if openErr2 != nil {
		return false, openErr2
	}
	defer r2.Close()
	info2, statErr2 := r2.Stat()
	if statErr2 != nil {
		return false, statErr2
	}

	if !c.Opt.ForceFileRead {
		// shortcut: ask the filesystem: are these files the same? (link, pathname, etc)
		if os.SameFile(info1, info2) {
			c.debugf("CompareFile(%s,%s): os reported same file\n", path1, path2)
			return true, nil
		}
	}

	if info1.Mode().IsRegular() && info2.Mode().IsRegular() {
		if info1.Size() != info2.Size() {
			c.debugf("CompareFile(%s,%s): distinct file sizes\n", path1, path2)
			return false, nil
		}
	}

	// If Opt.MaxSize not initialized, set maxSize to an appropriate value
	// for comparing regular files or streams (pipes, devices), etc.
	// Pass maxSize to getHash to ensure the hash is computed only up to
	// the user specified MaxSize amount.
	maxSize := c.Opt.MaxSize
	if maxSize == 0 {
		// If comparing regular to non-regular file, sizes may not
		// agree...  Use the larger value.
		maxSize = info1.Size()
		if maxSize < info2.Size() {
			maxSize = info2.Size()
		}
		if maxSize == 0 { // possible non-regular files
			maxSize = defaultMaxSize
		}
	}

	if c.multipleMode() {
		h1, err1 := c.getHash(path1, maxSize)
		if err1 != nil {
			return false, err1
		}
		h2, err2 := c.getHash(path2, maxSize)
		if err2 != nil {
			return false, err2
		}
		if !bytes.Equal(h1, h2) {
			return false, nil // hashes mismatch
		}
		// hashes match
		if !c.hashMatchCompare {
			return true, nil // accept hash match without byte-by-byte comparison
		}
		// do byte-by-byte comparison
		c.debugf("CompareFile(%s,%s): hash match, will compare bytes\n", path1, path2)
	}

	// Use our maxSize to avoid triggering the defaultMaxSize for files.
	// We still need to preserve the error returning properties of the
	// input amount exceeding MaxSize, so we can't use LimitedReader.
	c.resetDebugging()

	eq, err := c.compareReader(r1, r2, maxSize)

	c.printDebugCompareReader()

	return eq, err
}

func (c *Cmp) read(r io.Reader, buf []byte) (int, error) {
	n, err := r.Read(buf)

	if err == io.EOF {
		c.debugf("read: EOF found\n")
	}

	if c.Opt.Debug {
		c.readCount++
		c.readSum += int64(n)
		if n < c.readMin {
			c.readMin = n
		}
		if n > c.readMax {
			c.readMax = n
		}
	}

	return n, err
}

// CompareReader verifies that two readers provide same content.
//
// Reading more than MaxSize will return an error (along with the comparison
// value up to MaxSize bytes), unless one or both Readers are LimitedReaders,
// in which case MaxSize is ignored.
func (c *Cmp) CompareReader(r1, r2 io.Reader) (bool, error) {

	c.resetDebugging()

	equal, err := c.compareReader(r1, r2, c.Opt.MaxSize)

	c.printDebugCompareReader()

	return equal, err
}

func (c *Cmp) resetDebugging() {
	if c.Opt.Debug {
		c.readCount = 0
		c.readMin = 2000000000
		c.readMax = 0
		c.readSum = 0
	}
}

func (c *Cmp) printDebugCompareReader() {
	c.debugf("CompareReader(%d,%d): readCount=%d readMin=%d readMax=%d readSum=%d\n",
		len(c.buf), c.Opt.MaxSize, c.readCount, c.readMin, c.readMax, c.readSum)
}

// readPartial keeps reading from reader into provided buffer,
// until buffer size reaches exactly n2. n1 is initial buffer size.
// useful to ensure we get an specific buffer size from reader,
// withstanding partial reads.
func readPartial(c *Cmp, r io.Reader, buf []byte, n1, n2 int) (int, error) {
	for n1 < n2 {
		n, err := c.read(r, buf[n1:n2])
		n1 += n
		if err != nil {
			return n1, err
		}
	}
	return n1, nil
}

func (c *Cmp) compareReader(r1, r2 io.Reader, maxSize int64) (bool, error) {

	// Use LimitedReaders to ensure no data beyond MaxSize or LimitedReader limit
	// (when only one LimitedReader is given) other than at most a single byte.
	var lr1, lr2 io.Reader
	var checkAfterEOF1 bool
	var checkAfterEOF2 bool

	tmpLR1, isLR1 := r1.(*io.LimitedReader)
	tmpLR2, isLR2 := r2.(*io.LimitedReader)

	if isLR1 && isLR2 {
		lr1 = r1
		lr2 = r2
	} else if isLR1 && !isLR2 {
		lr1 = r1
		lr2 = io.LimitReader(r2, tmpLR1.N)
		checkAfterEOF2 = true
	} else if isLR2 && !isLR1 {
		lr2 = r2
		lr1 = io.LimitReader(r1, tmpLR2.N)
		checkAfterEOF1 = true
	} else {
		// Neither Reader is a LimitedReader.  Setup LimitedReaders w/ validated
		// maxSize limit, so that bytes compared will not exceed maxSize.
		if maxSize == 0 {
			maxSize = defaultMaxSize
		}

		if maxSize < 1 {
			return false, fmt.Errorf("nonpositive max size")
		}

		lr1 = io.LimitReader(r1, maxSize)
		lr2 = io.LimitReader(r2, maxSize)
		checkAfterEOF1 = true
		checkAfterEOF2 = true
	}

	buf := c.buf

	size := len(buf) / 2
	if size < 1 {
		return false, fmt.Errorf("insufficient buffer size")
	}

	buf1 := buf[:size]
	buf2 := buf[size : 2*size] // must force same size as buf1

	if len(buf1) != len(buf2) {
		return false, fmt.Errorf("buffer size mismatch buf1=%d buf2=%d", len(buf1), len(buf2))
	}

	eof1 := false
	eof2 := false

	for !eof1 && !eof2 {
		n1, err1 := c.read(lr1, buf1)
		switch err1 {
		case io.EOF:
			eof1 = true
		case nil:
		default:
			return false, err1
		}

		n2, err2 := c.read(lr2, buf2)
		switch err2 {
		case io.EOF:
			eof2 = true
		case nil:
		default:
			return false, err2
		}

		switch {
		case n1 < n2:
			n, errPart := readPartial(c, lr1, buf1, n1, n2)
			switch errPart {
			case io.EOF:
				eof1 = true
			case nil:
			default:
				return false, errPart
			}
			n1 = n
		case n2 < n1:
			n, errPart := readPartial(c, lr2, buf2, n2, n1)
			switch errPart {
			case io.EOF:
				eof2 = true
			case nil:
			default:
				return false, errPart
			}
			n2 = n
		}

		if n1 != n2 {
			c.debugf("compareReader: distinct buffer sizes\n")
			return false, nil
		}

		if !bytes.Equal(buf1[:n1], buf2[:n2]) {
			c.debugf("compareReader: found byte mismatch\n")
			return false, nil
		}
	}

	if !eof1 || !eof2 {
		c.debugf("compareReader: EOF for only one input\n")
		return false, nil
	}

	// Check the EOF status of the original readers. If neither was a
	// LimitedReader, and there is more readable data on either, then
	// return true with an error for exceeding MaxSize.  If one was a
	// LimitedReader and has more data, then return false.  If both were
	// LimitedReaders, then we've reached the desired EOF and return true.
	if checkAfterEOF1 && checkAfterEOF2 {
		// If both original readers need to be checked after EOF, then return
		// 'true' with an error if there is more data in either.
		eof1 = postEOFCheck(c, lr1, buf1[:1])
		eof2 = postEOFCheck(c, lr2, buf2[:1])
		switch {
		case eof1 && eof2:
			return true, nil
		default:
			c.debugf("compareReader: partial match, but max size exceeded\n")
			return true, fmt.Errorf("max read size reached")
		}
	}
	// Return false if only one reader is a LimitedReader, and the other
	// still has data to be read.  Else return true.
	if checkAfterEOF1 {
		return postEOFCheck(c, lr1, buf1[:1]), nil
	}
	if checkAfterEOF2 {
		return postEOFCheck(c, lr2, buf2[:1]), nil
	}

	return true, nil
}

// postEOFCheck returns false if there is more data in a LimitedReader after
// hitting EOF
func postEOFCheck(c *Cmp, r io.Reader, buf []byte) bool {
	tmpLR, isLR := r.(*io.LimitedReader)
	if isLR {
		r = tmpLR.R
	} else {
		c.debugf("compareReader: A type assertion of LimitedReader unexpectedly failed\n")
	}

	// Attempt to read more bytes from the original readers, to determine
	// if we should return an error for exceeding the MaxSize read limit.
	n, _ := readPartial(c, r, buf, 0, len(buf))
	return n == 0
}

func (c *Cmp) debugf(format string, v ...interface{}) {
	if c.Opt.Debug {
		fmt.Printf("DEBUG "+format, v...)
	}
}
<commit_msg>Use LimitReader count to avoid extra Read<commit_after>package equalfile

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"hash"
	"io"
	"os"
)

// Only the first 10^10 bytes of io.Reader are compared.  Ignored when using io.LimitedReader
const defaultMaxSize = 10000000000
const defaultBufSize = 20000

type Options struct {
	Debug         bool // enable debugging to stdout
	ForceFileRead bool // prevent shortcut at filesystem level (link, pathname, etc)

	// MaxSize is a safely limit to prevent forever reading from an infinite
	// reader.  If left unset, will default to 1OGBytes. Ignored when
	// CompareReader() is given one or more io.LimitedReader.
	MaxSize int64
}

type Cmp struct {
	Opt Options

	readCount int
	readMin   int
	readMax   int
	readSum   int64

	hashType         hash.Hash
	hashMatchCompare bool
	hashTable        map[string]hashSum

	buf []byte
}

type hashSum struct {
	result []byte
	err    error
}

// New creates Cmp for multiple comparison mode.
func NewMultiple(buf []byte, options Options, h hash.Hash, compareOnMatch bool) *Cmp {
	c := &Cmp{
		Opt:              options,
		hashType:         h,
		hashMatchCompare: compareOnMatch,
		hashTable:        map[string]hashSum{},
		buf:              buf,
	}
	if c.buf == nil || len(c.buf) == 0 {
		c.buf = make([]byte, defaultBufSize)
	}
	c.debugf("New: bufSize=%d\n", len(c.buf))
	return c
}

// New creates Cmp for single comparison mode.
func New(buf []byte, options Options) *Cmp {
	return NewMultiple(buf, options, nil, true)
}

func (c *Cmp) getHash(path string, maxSize int64) ([]byte, error) {
	h, found := c.hashTable[path]
	if found {
		return h.result, h.err
	}

	f, openErr := os.Open(path)
	if openErr != nil {
		return nil, openErr
	}
	defer f.Close()

	sum := make([]byte, c.hashType.Size())
	c.hashType.Reset()
	n, copyErr := io.CopyN(c.hashType, f, maxSize)
	copy(sum, c.hashType.Sum(nil))

	if copyErr == io.EOF && n < maxSize {
		copyErr = nil
	}

	return c.newHash(path, sum, copyErr)
}

func (c *Cmp) newHash(path string, sum []byte, e error) ([]byte, error) {

	c.hashTable[path] = hashSum{sum, e}

	c.debugf("newHash[%s]=%v: error=[%v]\n", path, hex.EncodeToString(sum), e)

	return sum, e
}

func (c *Cmp) multipleMode() bool {
	return c.hashType != nil
}

// CompareFile verifies that files with names path1, path2 have same contents.
func (c *Cmp) CompareFile(path1, path2 string) (bool, error) {

	if c.Opt.MaxSize < 0 {
		return false, fmt.Errorf("negative MaxSize")
	}

	r1, openErr1 := os.Open(path1)
	if openErr1 != nil {
		return false, openErr1
	}
	defer r1.Close()
	info1, statErr1 := r1.Stat()
	if statErr1 != nil {
		return false, statErr1
	}

	r2, openErr2 := os.Open(path2)
	if openErr2 != nil {
		return false, openErr2
	}
	defer r2.Close()
	info2, statErr2 := r2.Stat()
	if statErr2 != nil {
		return false, statErr2
	}

	if !c.Opt.ForceFileRead {
		// shortcut: ask the filesystem: are these files the same? (link, pathname, etc)
		if os.SameFile(info1, info2) {
			c.debugf("CompareFile(%s,%s): os reported same file\n", path1, path2)
			return true, nil
		}
	}

	if info1.Mode().IsRegular() && info2.Mode().IsRegular() {
		if info1.Size() != info2.Size() {
			c.debugf("CompareFile(%s,%s): distinct file sizes\n", path1, path2)
			return false, nil
		}
	}

	// If Opt.MaxSize not initialized, set maxSize to an appropriate value
	// for comparing regular files or streams (pipes, devices), etc.
	// Pass maxSize to getHash to ensure the hash is computed only up to
	// the user specified MaxSize amount.
	maxSize := c.Opt.MaxSize
	if maxSize == 0 {
		// If comparing regular to non-regular file, sizes may not
		// agree...  Use the larger value.
		maxSize = info1.Size()
		if maxSize < info2.Size() {
			maxSize = info2.Size()
		}
		if maxSize == 0 { // possible non-regular files
			maxSize = defaultMaxSize
		}
	}

	if c.multipleMode() {
		h1, err1 := c.getHash(path1, maxSize)
		if err1 != nil {
			return false, err1
		}
		h2, err2 := c.getHash(path2, maxSize)
		if err2 != nil {
			return false, err2
		}
		if !bytes.Equal(h1, h2) {
			return false, nil // hashes mismatch
		}
		// hashes match
		if !c.hashMatchCompare {
			return true, nil // accept hash match without byte-by-byte comparison
		}
		// do byte-by-byte comparison
		c.debugf("CompareFile(%s,%s): hash match, will compare bytes\n", path1, path2)
	}

	// Use our maxSize to avoid triggering the defaultMaxSize for files.
	// We still need to preserve the error returning properties of the
	// input amount exceeding MaxSize, so we can't use LimitedReader.
	c.resetDebugging()

	eq, err := c.compareReader(r1, r2, maxSize)

	c.printDebugCompareReader()

	return eq, err
}

func (c *Cmp) read(r io.Reader, buf []byte) (int, error) {
	n, err := r.Read(buf)

	if err == io.EOF {
		c.debugf("read: EOF found\n")
	}

	if c.Opt.Debug {
		c.readCount++
		c.readSum += int64(n)
		if n < c.readMin {
			c.readMin = n
		}
		if n > c.readMax {
			c.readMax = n
		}
	}

	return n, err
}

// CompareReader verifies that two readers provide same content.
//
// Reading more than MaxSize will return an error (along with the comparison
// value up to MaxSize bytes), unless one or both Readers are LimitedReaders,
// in which case MaxSize is ignored.
func (c *Cmp) CompareReader(r1, r2 io.Reader) (bool, error) {

	c.resetDebugging()

	equal, err := c.compareReader(r1, r2, c.Opt.MaxSize)

	c.printDebugCompareReader()

	return equal, err
}

func (c *Cmp) resetDebugging() {
	if c.Opt.Debug {
		c.readCount = 0
		c.readMin = 2000000000
		c.readMax = 0
		c.readSum = 0
	}
}

func (c *Cmp) printDebugCompareReader() {
	c.debugf("CompareReader(%d,%d): readCount=%d readMin=%d readMax=%d readSum=%d\n",
		len(c.buf), c.Opt.MaxSize, c.readCount, c.readMin, c.readMax, c.readSum)
}

// readPartial keeps reading from reader into provided buffer,
// until buffer size reaches exactly n2. n1 is initial buffer size.
// useful to ensure we get an specific buffer size from reader,
// withstanding partial reads.
func readPartial(c *Cmp, r io.Reader, buf []byte, n1, n2 int) (int, error) {
	for n1 < n2 {
		n, err := c.read(r, buf[n1:n2])
		n1 += n
		if err != nil {
			return n1, err
		}
	}
	return n1, nil
}

func (c *Cmp) compareReader(r1, r2 io.Reader, maxSize int64) (bool, error) {

	// Use LimitedReaders to ensure no data beyond MaxSize or LimitedReader limit
	// (when only one LimitedReader is given) other than at most a single byte.
	var lr1, lr2 io.Reader
	var checkAfterEOF1 bool
	var checkAfterEOF2 bool

	tmpLR1, isLR1 := r1.(*io.LimitedReader)
	tmpLR2, isLR2 := r2.(*io.LimitedReader)

	if isLR1 && isLR2 {
		lr1 = r1
		lr2 = r2
	} else if isLR1 && !isLR2 {
		lr1 = r1
		lr2 = io.LimitReader(r2, tmpLR1.N)
		checkAfterEOF2 = true
	} else if isLR2 && !isLR1 {
		lr2 = r2
		lr1 = io.LimitReader(r1, tmpLR2.N)
		checkAfterEOF1 = true
	} else {
		// Neither Reader is a LimitedReader.  Setup LimitedReaders w/ validated
		// maxSize limit, so that bytes compared will not exceed maxSize.
		if maxSize == 0 {
			maxSize = defaultMaxSize
		}

		if maxSize < 1 {
			return false, fmt.Errorf("nonpositive max size")
		}

		lr1 = io.LimitReader(r1, maxSize)
		lr2 = io.LimitReader(r2, maxSize)
		checkAfterEOF1 = true
		checkAfterEOF2 = true
	}

	buf := c.buf

	size := len(buf) / 2
	if size < 1 {
		return false, fmt.Errorf("insufficient buffer size")
	}

	buf1 := buf[:size]
	buf2 := buf[size : 2*size] // must force same size as buf1

	if len(buf1) != len(buf2) {
		return false, fmt.Errorf("buffer size mismatch buf1=%d buf2=%d", len(buf1), len(buf2))
	}

	eof1 := false
	eof2 := false

	for !eof1 && !eof2 {
		n1, err1 := c.read(lr1, buf1)
		switch err1 {
		case io.EOF:
			eof1 = true
		case nil:
		default:
			return false, err1
		}

		n2, err2 := c.read(lr2, buf2)
		switch err2 {
		case io.EOF:
			eof2 = true
		case nil:
		default:
			return false, err2
		}

		switch {
		case n1 < n2:
			n, errPart := readPartial(c, lr1, buf1, n1, n2)
			switch errPart {
			case io.EOF:
				eof1 = true
			case nil:
			default:
				return false, errPart
			}
			n1 = n
		case n2 < n1:
			n, errPart := readPartial(c, lr2, buf2, n2, n1)
			switch errPart {
			case io.EOF:
				eof2 = true
			case nil:
			default:
				return false, errPart
			}
			n2 = n
		}

		if n1 != n2 {
			c.debugf("compareReader: distinct buffer sizes\n")
			return false, nil
		}

		if !bytes.Equal(buf1[:n1], buf2[:n2]) {
			c.debugf("compareReader: found byte mismatch\n")
			return false, nil
		}
	}

	if !eof1 || !eof2 {
		c.debugf("compareReader: EOF for only one input\n")
		return false, nil
	}

	// Check the EOF status of the original readers. If neither was a
	// LimitedReader, and there is more readable data on either, then
	// return true with an error for exceeding MaxSize.  If one was a
	// LimitedReader and has more data, then return false.  If both were
	// LimitedReaders, then we've reached the desired EOF and return true.
	if checkAfterEOF1 && checkAfterEOF2 {
		// If both original readers need to be checked after EOF, then return
		// 'true' with an error if there is more data in either.
		eof1 = postEOFCheck(c, lr1, buf1[:1])
		eof2 = postEOFCheck(c, lr2, buf2[:1])
		switch {
		case eof1 && eof2:
			return true, nil
		default:
			c.debugf("compareReader: partial match, but max size exceeded\n")
			return true, fmt.Errorf("max read size reached")
		}
	}
	// Return false if only one reader is a LimitedReader, and the other
	// still has data to be read.  Else return true.
	if checkAfterEOF1 {
		return postEOFCheck(c, lr1, buf1[:1]), nil
	}
	if checkAfterEOF2 {
		return postEOFCheck(c, lr2, buf2[:1]), nil
	}

	return true, nil
}

// postEOFCheck returns false if there is more data in a LimitedReader after
// hitting EOF
func postEOFCheck(c *Cmp, r io.Reader, buf []byte) bool {
	tmpLR, isLR := r.(*io.LimitedReader)
	if isLR {
		// If the limit wasn't reached, then we don't need to check for
		// more data after the EOF
		if tmpLR.N > 0 {
			return true
		}

		// Use the internal Reader for checking for more data
		r = tmpLR.R
	} else {
		c.debugf("compareReader: A type assertion of LimitedReader unexpectedly failed\n")
	}

	// Attempt to read more bytes from the original readers, to determine
	// if we should return an error for exceeding the MaxSize read limit.
	n, _ := readPartial(c, r, buf, 0, len(buf))
	return n == 0
}

func (c *Cmp) debugf(format string, v ...interface{}) {
	if c.Opt.Debug {
		fmt.Printf("DEBUG "+format, v...)
	}
}
<|endoftext|>


<commit_before>package checkdisk

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

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

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

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

type unit struct {
	Name string
	Size float64
}

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

		freePct := float64(100) - disk.UsedPercent
		if thresholdPct > freePct {
			current = status
		}
	} else {
		thresholdVal, err := strconv.ParseFloat(threshold, 64)
		if err != nil {
			return checkers.UNKNOWN, err
		}

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

	return current, nil
}

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

		// Skip checking if disk.InodesTotal == 0 since inodesFreePct is meaningless.
		if disk.InodesTotal == 0 {
			return current, nil
		}

		inodesFreePct := float64(100) - disk.InodesUsedPercent
		if thresholdPct > inodesFreePct {
			current = status
		}

		return current, nil
	} else {
		return checkers.UNKNOWN, errors.New("-W, -K value should be N%")
	}
}

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

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

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

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

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

	if !opts.All {
		// Filtering partitions by Fstype
		if opts.IncludeType != nil {
			partitions = filterPartitionsByInclusion(partitions, *opts.IncludeType, fstypeOfPartition)
			if len(partitions) == 0 {
				return checkers.Unknown(fmt.Sprintf("Failed to fetch partitions: %s", errors.New("No device found for the specified *FsType*")))
			}
		}

		if opts.ExcludeType != nil {
			partitions = filterPartitionsByExclusion(partitions, *opts.ExcludeType, fstypeOfPartition)
		}

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

			partitions = filterPartitionsByInclusion(partitions, *opts.Path, mountpointOfPartition)
			if len(partitions) == 0 {
				return checkers.Unknown(fmt.Sprintf("Failed to fetch partitions: %s", errors.New("No device found for the specified *Mountpoint*")))
			}
		}

		if opts.Path == nil && opts.Exclude != nil {
			partitions = filterPartitionsByExclusion(partitions, *opts.Exclude, mountpointOfPartition)
		}
	}

	if len(partitions) == 0 {
		return checkers.Unknown(fmt.Sprintf("Failed to fetch partitions: %s", errors.New("No device found")))
	}

	var disks []*gpud.UsageStat

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

		if disk.Total != 0 {
			disks = append(disks, disk)
		}
	}

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

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

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

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

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

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

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

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

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

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

	return checkers.NewChecker(checkSt, msgss)
}

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

	return partitions, nil
}

func mountpointOfPartition(partition gpud.PartitionStat) string {
	return partition.Mountpoint
}

func fstypeOfPartition(partition gpud.PartitionStat) string {
	return partition.Fstype
}

func filterPartitionsByInclusion(partitions []gpud.PartitionStat, list []string, key func(_ gpud.PartitionStat) string) []gpud.PartitionStat {
	newPartitions := make([]gpud.PartitionStat, 0, len(partitions))
	for _, partition := range partitions {
		var ok = false
		for _, l := range list {
			if l == key(partition) {
				ok = true
				break
			}
		}
		if ok {
			newPartitions = append(newPartitions, partition)
		}
	}

	return newPartitions
}

func filterPartitionsByExclusion(partitions []gpud.PartitionStat, list []string, key func(_ gpud.PartitionStat) string) []gpud.PartitionStat {
	newPartitions := make([]gpud.PartitionStat, 0, len(partitions))
	for _, partition := range partitions {
		var ok = true
		for _, l := range list {
			if l == key(partition) {
				ok = false
				break
			}
		}
		if ok {
			newPartitions = append(newPartitions, partition)
		}
	}

	return newPartitions
}
<commit_msg>[check-disk] refactor<commit_after>package checkdisk

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

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

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

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

type unit struct {
	Name string
	Size float64
}

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

		freePct := float64(100) - disk.UsedPercent
		if thresholdPct > freePct {
			current = status
		}
	} else {
		thresholdVal, err := strconv.ParseFloat(threshold, 64)
		if err != nil {
			return checkers.UNKNOWN, err
		}

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

	return current, nil
}

func checkInodes(current checkers.Status, threshold string, disk *gpud.UsageStat, status checkers.Status) (checkers.Status, error) {
	if !strings.HasSuffix(threshold, "%") {
		return checkers.UNKNOWN, errors.New("-W, -K value should be N%")
	}

	thresholdPct, err := strconv.ParseFloat(strings.TrimRight(threshold, "%"), 64)
	if err != nil {
		return checkers.UNKNOWN, err
	}

	// Skip checking if disk.InodesTotal == 0 since inodesFreePct is meaningless.
	if disk.InodesTotal == 0 {
		return current, nil
	}

	inodesFreePct := float64(100) - disk.InodesUsedPercent
	if thresholdPct > inodesFreePct {
		current = status
	}

	return current, nil
}

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

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

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

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

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

	if !opts.All {
		// Filtering partitions by Fstype
		if opts.IncludeType != nil {
			partitions = filterPartitionsByInclusion(partitions, *opts.IncludeType, fstypeOfPartition)
			if len(partitions) == 0 {
				return checkers.Unknown(fmt.Sprintf("Failed to fetch partitions: %s", errors.New("No device found for the specified *FsType*")))
			}
		}

		if opts.ExcludeType != nil {
			partitions = filterPartitionsByExclusion(partitions, *opts.ExcludeType, fstypeOfPartition)
		}

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

			partitions = filterPartitionsByInclusion(partitions, *opts.Path, mountpointOfPartition)
			if len(partitions) == 0 {
				return checkers.Unknown(fmt.Sprintf("Failed to fetch partitions: %s", errors.New("No device found for the specified *Mountpoint*")))
			}
		}

		if opts.Path == nil && opts.Exclude != nil {
			partitions = filterPartitionsByExclusion(partitions, *opts.Exclude, mountpointOfPartition)
		}
	}

	if len(partitions) == 0 {
		return checkers.Unknown(fmt.Sprintf("Failed to fetch partitions: %s", errors.New("No device found")))
	}

	var disks []*gpud.UsageStat

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

		if disk.Total != 0 {
			disks = append(disks, disk)
		}
	}

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

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

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

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

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

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

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

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

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

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

	return checkers.NewChecker(checkSt, msgss)
}

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

	return partitions, nil
}

func mountpointOfPartition(partition gpud.PartitionStat) string {
	return partition.Mountpoint
}

func fstypeOfPartition(partition gpud.PartitionStat) string {
	return partition.Fstype
}

func filterPartitionsByInclusion(partitions []gpud.PartitionStat, list []string, key func(_ gpud.PartitionStat) string) []gpud.PartitionStat {
	newPartitions := make([]gpud.PartitionStat, 0, len(partitions))
	for _, partition := range partitions {
		var ok = false
		for _, l := range list {
			if l == key(partition) {
				ok = true
				break
			}
		}
		if ok {
			newPartitions = append(newPartitions, partition)
		}
	}

	return newPartitions
}

func filterPartitionsByExclusion(partitions []gpud.PartitionStat, list []string, key func(_ gpud.PartitionStat) string) []gpud.PartitionStat {
	newPartitions := make([]gpud.PartitionStat, 0, len(partitions))
	for _, partition := range partitions {
		var ok = true
		for _, l := range list {
			if l == key(partition) {
				ok = false
				break
			}
		}
		if ok {
			newPartitions = append(newPartitions, partition)
		}
	}

	return newPartitions
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/jpfairbanks/cse6140/project/hashes"
	"log"
	"math/rand"
	"time"
)

//CMSketch: A streaming data structure for tracking the elements of
//a long vector in a space efficient manner.
//Provides update and query operations
type CMSketch struct {
	Depth   int64
	Width   int64
	Hash    []hashes.Hash
	Counter Matrix
}

//NewCMSketch: Allocates a new sketch
//Does not create the hashes. We leave that up to you so that
//the consumer can pick hashes according to different schema.
func NewCMSketch(Depth, Width int64) CMSketch {
	Hash := make([]hashes.Hash, Depth)
	Counter := NewMatrix(Depth, Width)
	cms := CMSketch{Depth, Width, Hash, Counter}
	return cms
}

//Clone: Return a copy of the skecth with the same parameters and hashes,
//but a fresh set of counters. Suggested use is to compare two different
//manipulations for accuracy.
func (cms *CMSketch) Clone() CMSketch {
	Counter := NewMatrix(cms.Depth, cms.Width)
	out := CMSketch{cms.Depth, cms.Width, cms.Hash, Counter}
	return out
}

//Equal: Tell if two cms instances have the same parameters and data.
func (cms *CMSketch) Equal(other *CMSketch) bool {
	depths := cms.Depth == other.Depth
	widths := cms.Width == other.Width
	var hash bool
	hash = true
	for i, h := range cms.Hash {
		hash = other.Hash[i].Equal(h) && hash
		if !hash {
			break
		}
	}
	counters := cms.Counter.Equal(&other.Counter)
	out := depths && widths && hash && counters
	return out
}

//RandomHashes: Make uniformly random hashes from a user provided
//stream of random numbers.
func RandomHashes(r *rand.Rand, Depth int64) []hashes.Hash {
	hslice := make([]hashes.Hash, Depth)
	for i, _ := range hslice {
		hslice[i] = hashes.Rand(r)
	}
	return hslice
}

//UpdateSerial: Insert a single item into the sketch with a count.
//If counts can be negative, then you must estimate differently
func (cms *CMSketch) UpdateSerial(position int64, count int64) {
	for i, h := range cms.Hash {
		cms.Counter.Add(int64(i), h.Apply(position)%cms.Width, count)
	}
}

//AddSignalSlow: Updates the counter and signals completion on a channel
func (cms *CMSketch) AddSignalSlow(row int64, position int64,
	count int64, ch chan int64) {
	h := cms.Hash[row]
	cms.Counter.Add(row, h.Apply(position)%cms.Width, count)
	ch <- 1
}

//UpdateDepthParallelSlow: Insert a single item into the sketch with a count.
func (cms *CMSketch) UpdateDepthParallelSlow(position int64, count int64) {
	ch := make(chan int64, cms.Depth)
	var i int64
	for i = 0; i < cms.Depth; i++ {
		go cms.AddSignalSlow(i, position, count, ch)
	}
	for i = 0; i < cms.Depth; i++ {
		<-ch
	}
}

//AddSignal: Updates the counter and signals completion on a channel
func (cms *CMSketch) AddSignal(start_row int64, end_row int64, position int64,
	count int64, ch chan int64) {
	var h hashes.Hash
	for row := start_row; row < end_row; row++ {
		h = cms.Hash[row]
		cms.Counter.Add(row, h.Apply(position)%cms.Width, count)
	}
	ch <- 1
}

//UpdateDepthParallel: Insert a single item into the sketch with a count.
func (cms *CMSketch) UpdateDepthParallel(position int64, count int64, numProcs int64) {
	var i int64
	NbyP := cms.Depth / numProcs
	ch := make(chan int64, numProcs)
	for i = 0; i < numProcs; i++ {
		go cms.AddSignal(i*NbyP, (i+1)*NbyP, position, count, ch)
	}
	for i = 0; i < numProcs; i++ {
		<-ch
	}
}

//BatchUpdate: insert a batch of edges all at once.
//You must wait for a signal on the channel in order to ensure correct results
func (cms *CMSketch) BatchUpdate(elements []int64, ch chan int64, numProcs int64) {
	for _, z := range elements {
		cms.UpdateDepthParallel(z, 1, numProcs)
	}
	ch <- 1
}

//BatchUpdateSort: insert a batch of edges all at once.
//You must wait for a signal on the channel in order to ensure correct results
//This method uses a sort to improve cache locality.
func (cms *CMSketch) BatchUpdateSort(elements []int64, ch chan int64, numProcs int64) {
	//TODO: this method does not actually do any sorting
	for _, z := range elements {
		cms.UpdateDepthParallel(z, 1, numProcs)
	}
	ch <- 1
}

//PointQuery: query the value at position
func (cms *CMSketch) PointQuery(position int64) int64 {
	temp := make([]int64, cms.Depth)
	var h hashes.Hash
	var i int64
	for i = 0; i < cms.Depth; i++ {
		h = cms.Hash[i]
		temp[i] = cms.Counter.Read(i, h.Apply(position)%cms.Width)
	}
	mini := temp[0]
	for i = 0; i < cms.Depth; i++ {
		if mini > temp[i] {
			mini = temp[i]
		}
	}
	return mini
}

var depthPtr, widthPtr, efactorPtr *int64
var numProcs int64

func init() {
	depthPtr = flag.Int64("depth", 50, "sets the number of rows (and hash functions) in the CMSketch")
	widthPtr = flag.Int64("width", 80, "sets the number of columns in the CMSketch")
	efactorPtr = flag.Int64("efactor", 10, "number of elements = depth*width*efactor")
	NumProcsPtr := flag.Int64("procs", 1, "number of concurrent worker processes")
	flag.Parse()
	numProcs = *NumProcsPtr
}

func main() {
	//Handling command line parameters
	log.Printf("starting main\n")
	src := rand.NewSource(0)
	r := rand.New(src)
	var Depth, Width, efactor, numElements int64
	Depth = *depthPtr
	Width = *widthPtr
	efactor = *efactorPtr
	numElements = Depth * Width * efactor
	log.Printf("params:Depth:%d\n", Depth)
	log.Printf("params:Width:%d\n", Width)
	log.Printf("params:efactor:%d\n", efactor)
	log.Printf("params:numElements:%d\n", numElements)

	//Initialize Data Structures
	hslice := RandomHashes(r, Depth)
	cms := NewCMSketch(Depth, Width)
	cms.Hash = hslice
	//Make the zipf distribution of random input
	var j, z int64
	var s, v float64
	var imax uint64
	s = 1.2
	v = 1.0
	imax = 2 << 10
	zipfer := rand.NewZipf(r, s, v, imax)

	//Use set to store the exact answers
	set := make(map[int64]int64)
	log.Printf("Inserting\n")
	ts := time.Now()
	for j = 0; j < numElements; j++ {
		z = int64(zipfer.Uint64())
		//set[z] += 1
		//fmt.Println(z)
		cms.UpdateSerial(z, 1)
	}
	te := time.Now().Sub(ts)
	fmt.Printf("time: %v\n", te)
	fmt.Printf("%s\n", cms.Counter.String())
	var qj int64 //approximate answers
	var totalLoss float64
	loss := func(cj, qj float64) float64 {
		return (cj - qj) * (cj - qj)
	}
	for j, cj := range set {
		qj = cms.PointQuery(j)
		//fmt.Printf("results:%d %d %d %f\n", j, qj, cj, float64(qj)/float64(cj))
		totalLoss += loss(float64(cj), float64(qj))
	}
	fmt.Printf("Total Loss: %f/%d\n", totalLoss, numElements)
}
<commit_msg>fixing the batch methods so they actually do something. Without tests<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/jpfairbanks/cse6140/project/hashes"
	"log"
	"math/rand"
	"time"
)

//CMSketch: A streaming data structure for tracking the elements of
//a long vector in a space efficient manner.
//Provides update and query operations
type CMSketch struct {
	Depth   int64
	Width   int64
	Hash    []hashes.Hash
	Counter Matrix
}

//NewCMSketch: Allocates a new sketch
//Does not create the hashes. We leave that up to you so that
//the consumer can pick hashes according to different schema.
func NewCMSketch(Depth, Width int64) CMSketch {
	Hash := make([]hashes.Hash, Depth)
	Counter := NewMatrix(Depth, Width)
	cms := CMSketch{Depth, Width, Hash, Counter}
	return cms
}

//Clone: Return a copy of the skecth with the same parameters and hashes,
//but a fresh set of counters. Suggested use is to compare two different
//manipulations for accuracy.
func (cms *CMSketch) Clone() CMSketch {
	Counter := NewMatrix(cms.Depth, cms.Width)
	out := CMSketch{cms.Depth, cms.Width, cms.Hash, Counter}
	return out
}

//Equal: Tell if two cms instances have the same parameters and data.
func (cms *CMSketch) Equal(other *CMSketch) bool {
	depths := cms.Depth == other.Depth
	widths := cms.Width == other.Width
	var hash bool
	hash = true
	for i, h := range cms.Hash {
		hash = other.Hash[i].Equal(h) && hash
		if !hash {
			break
		}
	}
	counters := cms.Counter.Equal(&other.Counter)
	out := depths && widths && hash && counters
	return out
}

//RandomHashes: Make uniformly random hashes from a user provided
//stream of random numbers.
func RandomHashes(r *rand.Rand, Depth int64) []hashes.Hash {
	hslice := make([]hashes.Hash, Depth)
	for i, _ := range hslice {
		hslice[i] = hashes.Rand(r)
	}
	return hslice
}

//UpdateSerial: Insert a single item into the sketch with a count.
//If counts can be negative, then you must estimate differently
func (cms *CMSketch) UpdateSerial(position int64, count int64) {
	for i, h := range cms.Hash {
		cms.Counter.Add(int64(i), h.Apply(position)%cms.Width, count)
	}
}

//AddSignalSlow: Updates the counter and signals completion on a channel
func (cms *CMSketch) AddSignalSlow(row int64, position int64,
	count int64, ch chan int64) {
	h := cms.Hash[row]
	cms.Counter.Add(row, h.Apply(position)%cms.Width, count)
	ch <- 1
}

//UpdateDepthParallelSlow: Insert a single item into the sketch with a count.
func (cms *CMSketch) UpdateDepthParallelSlow(position int64, count int64) {
	ch := make(chan int64, cms.Depth)
	var i int64
	for i = 0; i < cms.Depth; i++ {
		go cms.AddSignalSlow(i, position, count, ch)
	}
	for i = 0; i < cms.Depth; i++ {
		<-ch
	}
}

//AddSignal: Updates the counter and signals completion on a channel
func (cms *CMSketch) AddSignal(start_row int64, end_row int64, position int64,
	count int64, ch chan int64) {
	var h hashes.Hash
	for row := start_row; row < end_row; row++ {
		h = cms.Hash[row]
		cms.Counter.Add(row, h.Apply(position)%cms.Width, count)
	}
	ch <- 1
}

//UpdateDepthParallel: Insert a single item into the sketch with a count.
func (cms *CMSketch) UpdateDepthParallel(position int64, count int64, numProcs int64) {
	var i int64
	NbyP := cms.Depth / numProcs
	ch := make(chan int64, numProcs)
	for i = 0; i < numProcs; i++ {
		go cms.AddSignal(i*NbyP, (i+1)*NbyP, position, count, ch)
	}
	for i = 0; i < numProcs; i++ {
		<-ch
	}
}

//BatchUpdate: insert a batch of edges all at once.
//You must wait for a signal on the channel in order to ensure correct results
func (cms *CMSketch) BatchUpdate(elements []int64, ch chan int64, numProcs int64) {
	batchSize := len(elements)
	for i, h := range cms.Hash {
		yarr := make([]int64, batchSize)
		h.ApplyBatch(elements, yarr)
		for k, y := range yarr {
			yarr[k] = y % cms.Width
		}
		for _, y := range yarr {
			cms.Counter.Add(int64(i), y, 1)
		}
	}
	ch <- 1
}

//BatchUpdateSort: insert a batch of edges all at once.
//You must wait for a signal on the channel in order to ensure correct results
//This method uses a sort to improve cache locality.
func (cms *CMSketch) BatchUpdateSort(elements []int64, ch chan int64, numProcs int64) {
	//TODO: this method needs to be optimized separately from BatchUpdate
	batchSize := len(elements)
	for i, h := range cms.Hash {
		yarr := make([]int64, batchSize)
		h.ApplyBatchSort(elements, yarr)
		for k, y := range yarr {
			yarr[k] = y % cms.Width
		}
		for _, y := range yarr {
			cms.Counter.Add(int64(i), y, 1)
		}
	}
	ch <- 1
}

//PointQuery: query the value at position
func (cms *CMSketch) PointQuery(position int64) int64 {
	temp := make([]int64, cms.Depth)
	var h hashes.Hash
	var i int64
	for i = 0; i < cms.Depth; i++ {
		h = cms.Hash[i]
		temp[i] = cms.Counter.Read(i, h.Apply(position)%cms.Width)
	}
	mini := temp[0]
	for i = 0; i < cms.Depth; i++ {
		if mini > temp[i] {
			mini = temp[i]
		}
	}
	return mini
}

var depthPtr, widthPtr, efactorPtr *int64
var numProcs int64

func init() {
	depthPtr = flag.Int64("depth", 50, "sets the number of rows (and hash functions) in the CMSketch")
	widthPtr = flag.Int64("width", 80, "sets the number of columns in the CMSketch")
	efactorPtr = flag.Int64("efactor", 10, "number of elements = depth*width*efactor")
	NumProcsPtr := flag.Int64("procs", 1, "number of concurrent worker processes")
	flag.Parse()
	numProcs = *NumProcsPtr
}

func main() {
	//Handling command line parameters
	log.Printf("starting main\n")
	src := rand.NewSource(0)
	r := rand.New(src)
	var Depth, Width, efactor, numElements int64
	Depth = *depthPtr
	Width = *widthPtr
	efactor = *efactorPtr
	numElements = Depth * Width * efactor
	log.Printf("params:Depth:%d\n", Depth)
	log.Printf("params:Width:%d\n", Width)
	log.Printf("params:efactor:%d\n", efactor)
	log.Printf("params:numElements:%d\n", numElements)

	//Initialize Data Structures
	hslice := RandomHashes(r, Depth)
	cms := NewCMSketch(Depth, Width)
	cms.Hash = hslice
	//Make the zipf distribution of random input
	var j, z int64
	var s, v float64
	var imax uint64
	s = 1.2
	v = 1.0
	imax = 2 << 10
	zipfer := rand.NewZipf(r, s, v, imax)

	//Use set to store the exact answers
	set := make(map[int64]int64)
	log.Printf("Inserting\n")
	ts := time.Now()
	for j = 0; j < numElements; j++ {
		z = int64(zipfer.Uint64())
		//set[z] += 1
		//fmt.Println(z)
		cms.UpdateSerial(z, 1)
	}
	te := time.Now().Sub(ts)
	fmt.Printf("time: %v\n", te)
	fmt.Printf("%s\n", cms.Counter.String())
	var qj int64 //approximate answers
	var totalLoss float64
	loss := func(cj, qj float64) float64 {
		return (cj - qj) * (cj - qj)
	}
	for j, cj := range set {
		qj = cms.PointQuery(j)
		//fmt.Printf("results:%d %d %d %f\n", j, qj, cj, float64(qj)/float64(cj))
		totalLoss += loss(float64(cj), float64(qj))
	}
	fmt.Printf("Total Loss: %f/%d\n", totalLoss, numElements)
}
<|endoftext|>


<commit_before>package chyle

import (
	"bytes"
	"fmt"
	"net/http"
	"net/url"

	"github.com/tidwall/gjson"

	"github.com/antham/envh"
)

// JiraIssueExpander fetch data using jira issue api
type JiraIssueExpander struct {
	client   http.Client
	username string
	password string
	URL      string
	keys     map[string]string
}

// NewJiraIssueExpanderFromPasswordAuth create a new JiraIssueExpander
func NewJiraIssueExpanderFromPasswordAuth(client http.Client, username string, password string, URL string, keys map[string]string) (JiraIssueExpander, error) {
	return JiraIssueExpander{client, username, password, URL, keys}, nil
}

// Expand fetch remote jira service if a jiraIssueId is defined to fetch issue datas
func (j JiraIssueExpander) Expand(commitMap *map[string]interface{}) (*map[string]interface{}, error) {
	var ID string

	if data, ok := (*commitMap)["jiraIssueId"]; ok {
		if data, ok := data.(string); ok {
			ID = data
		}
	}

	if ID == "" {
		return commitMap, nil
	}

	req, err := http.NewRequest("GET", j.URL+"/rest/api/2/issue/"+ID, nil)

	if err != nil {
		return commitMap, err
	}

	req.SetBasicAuth(j.username, j.password)
	req.Header.Set("Content-Type", "application/json")

	rep, err := j.client.Do(req)

	if err != nil {
		return commitMap, err
	}

	buf := bytes.NewBuffer([]byte{})
	err = rep.Write(buf)

	if err != nil {
		return commitMap, err
	}

	for identifier, key := range j.keys {
		if gjson.Get(buf.String(), key).Exists() {
			(*commitMap)[identifier] = gjson.Get(buf.String(), key).Value()
		} else {
			(*commitMap)[identifier] = nil
		}
	}

	return commitMap, nil
}

func buildJiraExpander(config *envh.EnvTree) (Expander, error) {
	datas := map[string]string{}
	keyValues := map[string]string{}

	for _, k := range []string{"USERNAME", "PASSWORD", "URL"} {
		v, err := config.FindString("CREDENTIALS", k)

		if err != nil {
			return nil, fmt.Errorf(`"%s" variable not found in "JIRA" config`, k)
		}

		datas[k] = v
	}

	_, err := url.ParseRequestURI(datas["URL"])

	if err != nil {
		return nil, fmt.Errorf(`"%s" is not a valid absolute URL defined in "JIRA" config`, datas["URL"])
	}

	debug(`Expander "USERNAME" defined with value "%s"`, datas["USERNAME"])
	debug(`Expander "PASSWORD" defined`)
	debug(`Expander "URL" defined with value "%s"`, datas["URL"])

	keys, err := config.FindChildrenKeys("KEYS")

	if err != nil {
		return nil, fmt.Errorf(`No "EXPANDERS_JIRA_KEYS" key found`)
	}

	for _, k := range keys {
		key, err := config.FindString("KEYS", k, "DESTKEY")

		if err != nil {
			return nil, fmt.Errorf(`An environment variable suffixed with "DESTKEY" must be defined with "%s", like EXPANDERS_JIRA_KEYS_%s_DESTKEY`, k, k)
		}

		value, err := config.FindString("KEYS", k, "FIELD")

		if err != nil {
			return nil, fmt.Errorf(`An environment variable suffixed with "FIELD" must be defined with "%s", like EXPANDERS_JIRA_KEYS_%s_FIELD`, k, k)
		}

		debug(`Expander KEY "%s" defined with value "%s"`, key, value)

		keyValues[key] = value
	}

	return NewJiraIssueExpanderFromPasswordAuth(http.Client{}, datas["USERNAME"], datas["PASSWORD"], datas["URL"], keyValues)
}
<commit_msg>ref(expander/jira) : simplify<commit_after>package chyle

import (
	"bytes"
	"fmt"
	"net/http"
	"net/url"

	"github.com/tidwall/gjson"

	"github.com/antham/envh"
)

// JiraIssueExpander fetch data using jira issue api
type JiraIssueExpander struct {
	client   http.Client
	username string
	password string
	URL      string
	keys     map[string]string
}

// NewJiraIssueExpanderFromPasswordAuth create a new JiraIssueExpander
func NewJiraIssueExpanderFromPasswordAuth(client http.Client, username string, password string, URL string, keys map[string]string) (JiraIssueExpander, error) {
	return JiraIssueExpander{client, username, password, URL, keys}, nil
}

// Expand fetch remote jira service if a jiraIssueId is defined to fetch issue datas
func (j JiraIssueExpander) Expand(commitMap *map[string]interface{}) (*map[string]interface{}, error) {
	var ID string

	if data, ok := (*commitMap)["jiraIssueId"]; true {
		if !ok {
			return commitMap, nil
		}

		if data, ok := data.(string); ok {
			ID = data
		}
	}

	req, err := http.NewRequest("GET", j.URL+"/rest/api/2/issue/"+ID, nil)

	if err != nil {
		return commitMap, err
	}

	req.SetBasicAuth(j.username, j.password)
	req.Header.Set("Content-Type", "application/json")

	rep, err := j.client.Do(req)

	if err != nil {
		return commitMap, err
	}

	buf := bytes.NewBuffer([]byte{})
	err = rep.Write(buf)

	if err != nil {
		return commitMap, err
	}

	for identifier, key := range j.keys {
		(*commitMap)[identifier] = nil

		if gjson.Get(buf.String(), key).Exists() {
			(*commitMap)[identifier] = gjson.Get(buf.String(), key).Value()
		}
	}

	return commitMap, nil
}

func buildJiraExpander(config *envh.EnvTree) (Expander, error) {
	datas := map[string]string{}
	keyValues := map[string]string{}

	for _, k := range []string{"USERNAME", "PASSWORD", "URL"} {
		v, err := config.FindString("CREDENTIALS", k)

		if err != nil {
			return nil, fmt.Errorf(`"%s" variable not found in "JIRA" config`, k)
		}

		datas[k] = v
	}

	_, err := url.ParseRequestURI(datas["URL"])

	if err != nil {
		return nil, fmt.Errorf(`"%s" is not a valid absolute URL defined in "JIRA" config`, datas["URL"])
	}

	debug(`Expander "USERNAME" defined with value "%s"`, datas["USERNAME"])
	debug(`Expander "PASSWORD" defined`)
	debug(`Expander "URL" defined with value "%s"`, datas["URL"])

	keys, err := config.FindChildrenKeys("KEYS")

	if err != nil {
		return nil, fmt.Errorf(`No "EXPANDERS_JIRA_KEYS" key found`)
	}

	for _, k := range keys {
		key, err := config.FindString("KEYS", k, "DESTKEY")

		if err != nil {
			return nil, fmt.Errorf(`An environment variable suffixed with "DESTKEY" must be defined with "%s", like EXPANDERS_JIRA_KEYS_%s_DESTKEY`, k, k)
		}

		value, err := config.FindString("KEYS", k, "FIELD")

		if err != nil {
			return nil, fmt.Errorf(`An environment variable suffixed with "FIELD" must be defined with "%s", like EXPANDERS_JIRA_KEYS_%s_FIELD`, k, k)
		}

		debug(`Expander KEY "%s" defined with value "%s"`, key, value)

		keyValues[key] = value
	}

	return NewJiraIssueExpanderFromPasswordAuth(http.Client{}, datas["USERNAME"], datas["PASSWORD"], datas["URL"], keyValues)
}
<|endoftext|>


<commit_before>package narc

import (
	"code.google.com/p/go.crypto/ssh"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"log"
)

type ProxyServer struct {
	agent   *Agent
	hostKey []byte

	listener *ssh.Listener
}

func NewProxyServer(agent *Agent) (*ProxyServer, error) {
	key, err := generateHostKey()
	if err != nil {
		return nil, err
	}

	return &ProxyServer{
		agent:   agent,
		hostKey: key,
	}, nil
}

func (p *ProxyServer) Start(port int) error {
	config := &ssh.ServerConfig{
		PasswordCallback: p.verifyTaskAccess,
	}

	err := config.SetRSAPrivateKey(p.hostKey)
	if err != nil {
		return err
	}

	l, err := ssh.Listen("tcp", fmt.Sprintf(":%d", port), config)
	if err != nil {
		return err
	}

	p.listener = l

	log.Println("listening on port", port)

	go p.serveConnections()

	return nil
}

func (p *ProxyServer) Stop() error {
	if p.listener != nil {
		return p.listener.Close()
	}

	return nil
}

func (p *ProxyServer) verifyTaskAccess(conn *ssh.ServerConn, user, password string) bool {
	log.Println("verifying:", user, password)

	task, found := p.agent.Registry.Lookup(user)
	if !found {
		log.Println("verify failed: task not found")
		return false
	}

	return task.SecureToken == password
}

func (p *ProxyServer) serveConnections() {
	for {
		conn, err := p.listener.Accept()
		if err != nil {
			log.Println("error accepting connection:", err)
			break
		}

		log.Println("accepted connection")

		err = conn.Handshake()
		if err != nil {
			log.Println("handshake failed:", err)
			continue
		}

		go p.handleSession(conn)
	}
}

func (p *ProxyServer) handleSession(conn *ssh.ServerConn) {
	defer conn.Close()

	for {
		channel, err := conn.Accept()
		if err == io.EOF {
			return
		}

		if err != nil {
			log.Println("failed to accept ssh channel", err)
			return
		}

		if channel.ChannelType() != "session" {
			channel.Reject(ssh.UnknownChannelType, "unknown channel type")
			break
		}

		go p.handleChannel(channel, conn.User)
	}
}

func (p *ProxyServer) handleChannel(channel ssh.Channel, taskID string) {
	err := channel.Accept()
	if err != nil {
		log.Println("failed to accept channel request:", err)
		return
	}

	task, found := p.agent.Registry.Lookup(taskID)
	if !found {
		log.Println("unknown task:", task)
		return
	}

	commandDone, channelClosed, err := task.Attach(channel)
	if err != nil {
		log.Println("failed to execute task:", err)
		return
	}

	select {
	case <-commandDone:
		err := p.agent.StopTask(taskID)
		if err != nil {
			panic(err)
		}

		channel.Close()

		<-channelClosed

	case <-channelClosed:
	}
}

func generateHostKey() ([]byte, error) {
	key, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return nil, err
	}

	blk := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}

	return pem.EncodeToMemory(blk), nil
}
<commit_msg>log when the task exits vs. channel is closed<commit_after>package narc

import (
	"code.google.com/p/go.crypto/ssh"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io"
	"log"
)

type ProxyServer struct {
	agent   *Agent
	hostKey []byte

	listener *ssh.Listener
}

func NewProxyServer(agent *Agent) (*ProxyServer, error) {
	key, err := generateHostKey()
	if err != nil {
		return nil, err
	}

	return &ProxyServer{
		agent:   agent,
		hostKey: key,
	}, nil
}

func (p *ProxyServer) Start(port int) error {
	config := &ssh.ServerConfig{
		PasswordCallback: p.verifyTaskAccess,
	}

	err := config.SetRSAPrivateKey(p.hostKey)
	if err != nil {
		return err
	}

	l, err := ssh.Listen("tcp", fmt.Sprintf(":%d", port), config)
	if err != nil {
		return err
	}

	p.listener = l

	log.Println("listening on port", port)

	go p.serveConnections()

	return nil
}

func (p *ProxyServer) Stop() error {
	if p.listener != nil {
		return p.listener.Close()
	}

	return nil
}

func (p *ProxyServer) verifyTaskAccess(conn *ssh.ServerConn, user, password string) bool {
	log.Println("verifying:", user, password)

	task, found := p.agent.Registry.Lookup(user)
	if !found {
		log.Println("verify failed: task not found")
		return false
	}

	return task.SecureToken == password
}

func (p *ProxyServer) serveConnections() {
	for {
		conn, err := p.listener.Accept()
		if err != nil {
			log.Println("error accepting connection:", err)
			break
		}

		log.Println("accepted connection")

		err = conn.Handshake()
		if err != nil {
			log.Println("handshake failed:", err)
			continue
		}

		go p.handleSession(conn)
	}
}

func (p *ProxyServer) handleSession(conn *ssh.ServerConn) {
	defer conn.Close()

	for {
		channel, err := conn.Accept()
		if err == io.EOF {
			return
		}

		if err != nil {
			log.Println("failed to accept ssh channel", err)
			return
		}

		if channel.ChannelType() != "session" {
			channel.Reject(ssh.UnknownChannelType, "unknown channel type")
			break
		}

		go p.handleChannel(channel, conn.User)
	}
}

func (p *ProxyServer) handleChannel(channel ssh.Channel, taskID string) {
	err := channel.Accept()
	if err != nil {
		log.Println("failed to accept channel request:", err)
		return
	}

	task, found := p.agent.Registry.Lookup(taskID)
	if !found {
		log.Println("unknown task:", task)
		return
	}

	commandDone, channelClosed, err := task.Attach(channel)
	if err != nil {
		log.Println("failed to execute task:", err)
		return
	}

	select {
	case <-commandDone:
		log.Println("command exited")

		err = p.agent.StopTask(taskID)
		if err != nil {
			panic(err)
		}

		log.Println("stopped task")

		channel.Close()

		<-channelClosed

	case <-channelClosed:
		log.Println("channel closed")
	}
}

func generateHostKey() ([]byte, error) {
	key, err := rsa.GenerateKey(rand.Reader, 1024)
	if err != nil {
		return nil, err
	}

	blk := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(key),
	}

	return pem.EncodeToMemory(blk), nil
}
<|endoftext|>


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

/*
#include <utmp.h>
*/
import "C"

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/locksmith/Godeps/_workspace/src/github.com/coreos/go-systemd/dbus"
	"github.com/coreos/locksmith/Godeps/_workspace/src/github.com/coreos/go-systemd/login1"
	"github.com/coreos/locksmith/Godeps/_workspace/src/github.com/coreos/pkg/capnslog"

	"github.com/coreos/locksmith/lock"
	"github.com/coreos/locksmith/pkg/machineid"
	"github.com/coreos/locksmith/pkg/timeutil"
	"github.com/coreos/locksmith/updateengine"
)

const (
	initialInterval   = time.Second * 5
	maxInterval       = time.Minute * 5
	loginsRebootDelay = time.Minute * 5
)

var (
	etcdServices = []string{
		"etcd.service",
		"etcd2.service",
	}

	// TODO(mischief): daemon is not really a seperate package. it probably should be.
	dlog = capnslog.NewPackageLogger("github.com/coreos/locksmith", "daemon")
)

const (
	StrategyReboot     = "reboot"
	StrategyEtcdLock   = "etcd-lock"
	StrategyBestEffort = "best-effort"
	StrategyOff        = "off"
)

// attempt to broadcast msg to all lines registered in utmp
// returns count of lines successfully opened (and likely broadcasted to)
func broadcast(msg string) uint {
	var cnt uint
	C.setutent()

	for {
		var utmp *C.struct_utmp
		utmp = C.getutent()
		if utmp == nil {
			break
		}

		line := C.GoString(&utmp.ut_line[0])
		tty, _ := os.OpenFile("/dev/"+line, os.O_WRONLY, 0)
		if tty == nil {
			// ignore errors, this is just a best-effort courtesy notice
			continue
		}
		cnt++
		go func() {
			fmt.Fprintf(tty, "\r%79s\r\n", " ")
			fmt.Fprintf(tty, "%-79.79s\007\007\r\n", fmt.Sprintf("Broadcast message from locksmithd at %s:", time.Now()))
			fmt.Fprintf(tty, "%-79.79s\r\n", msg) // msg is assumed to be short and not require wrapping
			fmt.Fprintf(tty, "\r%79s\r\n", " ")
			tty.Close()
		}()
	}

	return cnt
}

func expBackoff(interval time.Duration) time.Duration {
	interval = interval * 2
	if interval > maxInterval {
		interval = maxInterval
	}
	return interval
}

func rebootAndSleep(lgn *login1.Conn) {
	// Broadcast a notice, if broadcast found lines to notify, delay the reboot.
	delaymins := loginsRebootDelay / time.Minute
	lines := broadcast(fmt.Sprintf("System reboot in %d minutes!", delaymins))
	if 0 != lines {
		dlog.Noticef("Logins detected, delaying reboot for %d minutes.", delaymins)
		time.Sleep(loginsRebootDelay)
	}
	lgn.Reboot(false)
	dlog.Info("Reboot sent. Going to sleep.")

	// Wait a really long time for the reboot to occur.
	time.Sleep(time.Hour * 24 * 7)
}

// lockAndReboot attempts to acquire the lock and reboot the machine in an
// infinite loop. Returns if the reboot failed.
func (r rebooter) lockAndReboot(lck *lock.Lock) {
	interval := initialInterval
	for {
		err := lck.Lock()
		if err != nil && err != lock.ErrExist {
			interval = expBackoff(interval)
			dlog.Warningf("Failed to acquire lock: %v. Retrying in %v.", err, interval)
			time.Sleep(interval)

			continue
		}

		rebootAndSleep(r.lgn)

		return
	}
}

func setupLock() (lck *lock.Lock, err error) {
	elc, err := getClient()
	if err != nil {
		return nil, fmt.Errorf("Error initializing etcd client: %v", err)
	}

	mID := machineid.MachineID("/")
	if mID == "" {
		return nil, fmt.Errorf("Cannot read machine-id")
	}

	lck = lock.New(mID, elc)

	return lck, nil
}

// etcdActive returns true if etcd is not in an inactive state according to systemd.
func etcdActive() (active bool, name string, err error) {
	active = false
	name = ""

	sys, err := dbus.New()
	if err != nil {
		return
	}
	defer sys.Close()

	for _, service := range etcdServices {
		prop, err := sys.GetUnitProperty(service, "ActiveState")
		if err != nil {
			continue
		}

		switch prop.Value.Value().(string) {
		case "inactive":
			continue
		default:
			active = true
			name = service
			break
		}
	}

	return
}

type rebooter struct {
	strategy string
	lgn      *login1.Conn
}

func (r rebooter) reboot() int {
	useLock, err := useLock(r.strategy)
	if err != nil {
		dlog.Errorf("Failed to figure out if locksmithd needs to use a lock: %v", err)
		return 1
	}

	if useLock {
		lck, err := setupLock()
		if err != nil {
			dlog.Errorf("Failed to set up lock: %v", err)
			return 1
		}

		err = unlockIfHeld(lck)
		if err != nil {
			dlog.Errorf("Failed to unlock held lock: %v", err)
			return 1
		}

		r.lockAndReboot(lck)
	}

	rebootAndSleep(r.lgn)
	dlog.Fatal("Tried to reboot but did not!")
	return 1
}

// useLock returns whether locksmith should attempt to take a lock before
// rebooting or release a lock afterwards, based on the given strategy.
// If strategy is set to best effort, this will be dependent on whether the
// local instance of etcd is active. Otherwise, the lock will always be
// attempted (in the case of strategy = etcd lock) or never be attempted (in
// the case of strategy = reboot)
func useLock(strategy string) (useLock bool, err error) {
	switch strategy {
	case StrategyBestEffort:
		active, name, err := etcdActive()
		if err != nil {
			return false, err
		}
		if active {
			dlog.Infof("%s is active", name)
			useLock = true
		} else {
			dlog.Infof("%v are inactive", etcdServices)
			useLock = false
		}
	case StrategyEtcdLock:
		useLock = true
	case StrategyReboot:
		useLock = false
	default:
		return false, fmt.Errorf("unknown strategy: %s", strategy)
	}

	return useLock, nil
}

// unlockIfHeld will unlock a lock, if it is held by this machine, or return an error.
func unlockIfHeld(lck *lock.Lock) error {
	err := lck.Unlock()
	if err == lock.ErrNotExist {
		return nil
	} else if err == nil {
		dlog.Info("Unlocked existing lock for this machine")
		return nil
	}

	return err
}

// unlockHeldLock will loop until it can confirm that any held locks are
// released or a stop signal is sent.
func unlockHeldLocks(strategy string, stop chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()
	interval := initialInterval
	for {
		var reason string
		select {
		case <-stop:
			return
		case <-time.After(interval):
			// Here we assume that the strategy used by locksmith
			// is consistent with that used before the previous
			// reboot (if any), and use that to decide whether to
			// attempt to unlock.
			shouldUnlock, err := useLock(strategy)
			if err != nil {
				reason = fmt.Sprintf("error checking whether lock should be released: %v", err)
				break
			}
			if !shouldUnlock {
				reason = fmt.Sprintf("%v are inactive", etcdServices)
				break
			}

			lck, err := setupLock()
			if err != nil {
				reason = "error setting up lock: " + err.Error()
				break
			}

			err = unlockIfHeld(lck)
			if err == nil {
				return
			}
			reason = err.Error()
		}

		interval = expBackoff(interval)
		dlog.Errorf("Unlocking old locks failed: %v. Retrying in %v.", reason, interval)
	}
}

// runDaemon waits for the reboot needed signal coming out of update engine and
// attempts to acquire the reboot lock. If the reboot lock is acquired then the
// machine will reboot.
func runDaemon() int {
	var period *timeutil.Periodic

	strategy := os.Getenv("REBOOT_STRATEGY")

	if strategy == "" {
		strategy = StrategyBestEffort
	}

	if strategy == StrategyOff {
		dlog.Noticef("Reboot strategy is %q - shutting down.", strategy)
		return 0
	}

	// XXX: REBOOT_WINDOW_* are deprecated in favor of variables with LOCKSMITHD_ prefix,
	// but the old ones are read for compatibility.
	startw := os.Getenv("LOCKSMITHD_REBOOT_WINDOW_START")
	if startw == "" {
		startw = os.Getenv("REBOOT_WINDOW_START")
	}

	lengthw := os.Getenv("LOCKSMITHD_REBOOT_WINDOW_LENGTH")
	if lengthw == "" {
		lengthw = os.Getenv("REBOOT_WINDOW_LENGTH")
	}

	if (startw == "") != (lengthw == "") {
		dlog.Fatal("Either both or neither $REBOOT_WINDOW_START and $REBOOT_WINDOW_LENGTH must be set")
	}

	if startw != "" && lengthw != "" {
		p, err := timeutil.ParsePeriodic(startw, lengthw)
		if err != nil {
			dlog.Fatalf("Error parsing reboot window: %s", err)
		}

		period = p
	}

	shutdown := make(chan os.Signal, 1)
	stop := make(chan struct{}, 1)

	go func() {
		<-shutdown
		dlog.Notice("Received interrupt/termination signal - shutting down.")
		os.Exit(0)
	}()
	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)

	ue, err := updateengine.New()
	if err != nil {
		dlog.Fatalf("Error initializing update1 client: %v", err)
	}

	lgn, err := login1.New()
	if err != nil {
		dlog.Fatalf("Error initializing login1 client: %v", err)
	}

	var wg sync.WaitGroup
	if strategy != StrategyReboot {
		wg.Add(1)
		go unlockHeldLocks(strategy, stop, &wg)
	}

	ch := make(chan updateengine.Status, 1)
	go ue.RebootNeededSignal(ch, stop)

	r := rebooter{strategy, lgn}

	result, err := ue.GetStatus()
	if err != nil {
		dlog.Fatalf("Cannot get update engine status: %v", err)
	}

	dlog.Infof("locksmithd starting currentOperation=%q strategy=%q", result.CurrentOperation, strategy)

	if result.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		<-ch
	}

	close(stop)
	wg.Wait()

	if period != nil {
		now := time.Now()
		sleeptime := period.DurationToStart(now)
		if sleeptime > 0 {
			dlog.Infof("Waiting for %s to reboot.", sleeptime)
			time.Sleep(sleeptime)
		}
	}

	return r.reboot()
}
<commit_msg>locksmithd: complain loudly if besteffort is used<commit_after>// Copyright 2015 CoreOS, Inc.
//
// 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

/*
#include <utmp.h>
*/
import "C"

import (
	"fmt"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"github.com/coreos/locksmith/Godeps/_workspace/src/github.com/coreos/go-systemd/dbus"
	"github.com/coreos/locksmith/Godeps/_workspace/src/github.com/coreos/go-systemd/login1"
	"github.com/coreos/locksmith/Godeps/_workspace/src/github.com/coreos/pkg/capnslog"

	"github.com/coreos/locksmith/lock"
	"github.com/coreos/locksmith/pkg/machineid"
	"github.com/coreos/locksmith/pkg/timeutil"
	"github.com/coreos/locksmith/updateengine"
)

const (
	initialInterval   = time.Second * 5
	maxInterval       = time.Minute * 5
	loginsRebootDelay = time.Minute * 5
)

var (
	etcdServices = []string{
		"etcd.service",
		"etcd2.service",
	}

	// TODO(mischief): daemon is not really a seperate package. it probably should be.
	dlog = capnslog.NewPackageLogger("github.com/coreos/locksmith", "daemon")
)

const (
	StrategyReboot     = "reboot"
	StrategyEtcdLock   = "etcd-lock"
	StrategyBestEffort = "best-effort"
	StrategyOff        = "off"
)

// attempt to broadcast msg to all lines registered in utmp
// returns count of lines successfully opened (and likely broadcasted to)
func broadcast(msg string) uint {
	var cnt uint
	C.setutent()

	for {
		var utmp *C.struct_utmp
		utmp = C.getutent()
		if utmp == nil {
			break
		}

		line := C.GoString(&utmp.ut_line[0])
		tty, _ := os.OpenFile("/dev/"+line, os.O_WRONLY, 0)
		if tty == nil {
			// ignore errors, this is just a best-effort courtesy notice
			continue
		}
		cnt++
		go func() {
			fmt.Fprintf(tty, "\r%79s\r\n", " ")
			fmt.Fprintf(tty, "%-79.79s\007\007\r\n", fmt.Sprintf("Broadcast message from locksmithd at %s:", time.Now()))
			fmt.Fprintf(tty, "%-79.79s\r\n", msg) // msg is assumed to be short and not require wrapping
			fmt.Fprintf(tty, "\r%79s\r\n", " ")
			tty.Close()
		}()
	}

	return cnt
}

func expBackoff(interval time.Duration) time.Duration {
	interval = interval * 2
	if interval > maxInterval {
		interval = maxInterval
	}
	return interval
}

func rebootAndSleep(lgn *login1.Conn) {
	// Broadcast a notice, if broadcast found lines to notify, delay the reboot.
	delaymins := loginsRebootDelay / time.Minute
	lines := broadcast(fmt.Sprintf("System reboot in %d minutes!", delaymins))
	if 0 != lines {
		dlog.Noticef("Logins detected, delaying reboot for %d minutes.", delaymins)
		time.Sleep(loginsRebootDelay)
	}
	lgn.Reboot(false)
	dlog.Info("Reboot sent. Going to sleep.")

	// Wait a really long time for the reboot to occur.
	time.Sleep(time.Hour * 24 * 7)
}

// lockAndReboot attempts to acquire the lock and reboot the machine in an
// infinite loop. Returns if the reboot failed.
func (r rebooter) lockAndReboot(lck *lock.Lock) {
	interval := initialInterval
	for {
		err := lck.Lock()
		if err != nil && err != lock.ErrExist {
			interval = expBackoff(interval)
			dlog.Warningf("Failed to acquire lock: %v. Retrying in %v.", err, interval)
			time.Sleep(interval)

			continue
		}

		rebootAndSleep(r.lgn)

		return
	}
}

func setupLock() (lck *lock.Lock, err error) {
	elc, err := getClient()
	if err != nil {
		return nil, fmt.Errorf("Error initializing etcd client: %v", err)
	}

	mID := machineid.MachineID("/")
	if mID == "" {
		return nil, fmt.Errorf("Cannot read machine-id")
	}

	lck = lock.New(mID, elc)

	return lck, nil
}

// etcdActive returns true if etcd is not in an inactive state according to systemd.
func etcdActive() (active bool, name string, err error) {
	active = false
	name = ""

	sys, err := dbus.New()
	if err != nil {
		return
	}
	defer sys.Close()

	for _, service := range etcdServices {
		prop, err := sys.GetUnitProperty(service, "ActiveState")
		if err != nil {
			continue
		}

		switch prop.Value.Value().(string) {
		case "inactive":
			continue
		default:
			active = true
			name = service
			break
		}
	}

	return
}

type rebooter struct {
	strategy string
	lgn      *login1.Conn
}

func (r rebooter) reboot() int {
	useLock, err := useLock(r.strategy)
	if err != nil {
		dlog.Errorf("Failed to figure out if locksmithd needs to use a lock: %v", err)
		return 1
	}

	if useLock {
		lck, err := setupLock()
		if err != nil {
			dlog.Errorf("Failed to set up lock: %v", err)
			return 1
		}

		err = unlockIfHeld(lck)
		if err != nil {
			dlog.Errorf("Failed to unlock held lock: %v", err)
			return 1
		}

		r.lockAndReboot(lck)
	}

	rebootAndSleep(r.lgn)
	dlog.Fatal("Tried to reboot but did not!")
	return 1
}

// useLock returns whether locksmith should attempt to take a lock before
// rebooting or release a lock afterwards, based on the given strategy.
// If strategy is set to best effort, this will be dependent on whether the
// local instance of etcd is active. Otherwise, the lock will always be
// attempted (in the case of strategy = etcd lock) or never be attempted (in
// the case of strategy = reboot)
func useLock(strategy string) (useLock bool, err error) {
	switch strategy {
	case StrategyBestEffort:
		active, name, err := etcdActive()
		if err != nil {
			return false, err
		}
		if active {
			dlog.Infof("%s is active", name)
			useLock = true
		} else {
			dlog.Infof("%v are inactive", etcdServices)
			useLock = false
		}
	case StrategyEtcdLock:
		useLock = true
	case StrategyReboot:
		useLock = false
	default:
		return false, fmt.Errorf("unknown strategy: %s", strategy)
	}

	return useLock, nil
}

// unlockIfHeld will unlock a lock, if it is held by this machine, or return an error.
func unlockIfHeld(lck *lock.Lock) error {
	err := lck.Unlock()
	if err == lock.ErrNotExist {
		return nil
	} else if err == nil {
		dlog.Info("Unlocked existing lock for this machine")
		return nil
	}

	return err
}

// unlockHeldLock will loop until it can confirm that any held locks are
// released or a stop signal is sent.
func unlockHeldLocks(strategy string, stop chan struct{}, wg *sync.WaitGroup) {
	defer wg.Done()
	interval := initialInterval
	for {
		var reason string
		select {
		case <-stop:
			return
		case <-time.After(interval):
			// Here we assume that the strategy used by locksmith
			// is consistent with that used before the previous
			// reboot (if any), and use that to decide whether to
			// attempt to unlock.
			shouldUnlock, err := useLock(strategy)
			if err != nil {
				reason = fmt.Sprintf("error checking whether lock should be released: %v", err)
				break
			}
			if !shouldUnlock {
				reason = fmt.Sprintf("%v are inactive", etcdServices)
				break
			}

			lck, err := setupLock()
			if err != nil {
				reason = "error setting up lock: " + err.Error()
				break
			}

			err = unlockIfHeld(lck)
			if err == nil {
				return
			}
			reason = err.Error()
		}

		interval = expBackoff(interval)
		dlog.Errorf("Unlocking old locks failed: %v. Retrying in %v.", reason, interval)
	}
}

// runDaemon waits for the reboot needed signal coming out of update engine and
// attempts to acquire the reboot lock. If the reboot lock is acquired then the
// machine will reboot.
func runDaemon() int {
	var period *timeutil.Periodic

	strategy := os.Getenv("REBOOT_STRATEGY")

	if strategy == "" {
		strategy = StrategyBestEffort
	}

	// XXX: complain loudly if besteffort is used
	if strategy == StrategyBestEffort {
		dlog.Errorf("Reboot strategy %q is deprecated and will be removed in the future.", strategy)
		dlog.Errorf("Please explicitly set the reboot strategy to one of %v", []string{StrategyOff, StrategyReboot, StrategyEtcdLock})
		dlog.Error("See https://coreos.com/os/docs/latest/update-strategies.html for details on configuring reboot strategies.")
	}

	if strategy == StrategyOff {
		dlog.Noticef("Reboot strategy is %q - shutting down.", strategy)
		return 0
	}

	// XXX: REBOOT_WINDOW_* are deprecated in favor of variables with LOCKSMITHD_ prefix,
	// but the old ones are read for compatibility.
	startw := os.Getenv("LOCKSMITHD_REBOOT_WINDOW_START")
	if startw == "" {
		startw = os.Getenv("REBOOT_WINDOW_START")
	}

	lengthw := os.Getenv("LOCKSMITHD_REBOOT_WINDOW_LENGTH")
	if lengthw == "" {
		lengthw = os.Getenv("REBOOT_WINDOW_LENGTH")
	}

	if (startw == "") != (lengthw == "") {
		dlog.Fatal("Either both or neither $REBOOT_WINDOW_START and $REBOOT_WINDOW_LENGTH must be set")
	}

	if startw != "" && lengthw != "" {
		p, err := timeutil.ParsePeriodic(startw, lengthw)
		if err != nil {
			dlog.Fatalf("Error parsing reboot window: %s", err)
		}

		period = p
	}

	shutdown := make(chan os.Signal, 1)
	stop := make(chan struct{}, 1)

	go func() {
		<-shutdown
		dlog.Notice("Received interrupt/termination signal - shutting down.")
		os.Exit(0)
	}()
	signal.Notify(shutdown, syscall.SIGINT, syscall.SIGTERM)

	ue, err := updateengine.New()
	if err != nil {
		dlog.Fatalf("Error initializing update1 client: %v", err)
	}

	lgn, err := login1.New()
	if err != nil {
		dlog.Fatalf("Error initializing login1 client: %v", err)
	}

	var wg sync.WaitGroup
	if strategy != StrategyReboot {
		wg.Add(1)
		go unlockHeldLocks(strategy, stop, &wg)
	}

	ch := make(chan updateengine.Status, 1)
	go ue.RebootNeededSignal(ch, stop)

	r := rebooter{strategy, lgn}

	result, err := ue.GetStatus()
	if err != nil {
		dlog.Fatalf("Cannot get update engine status: %v", err)
	}

	dlog.Infof("locksmithd starting currentOperation=%q strategy=%q", result.CurrentOperation, strategy)

	if result.CurrentOperation != updateengine.UpdateStatusUpdatedNeedReboot {
		<-ch
	}

	close(stop)
	wg.Wait()

	if period != nil {
		now := time.Now()
		sleeptime := period.DurationToStart(now)
		if sleeptime > 0 {
			dlog.Infof("Waiting for %s to reboot.", sleeptime)
			time.Sleep(sleeptime)
		}
	}

	return r.reboot()
}
<|endoftext|>


<commit_before>package log_test

import (
	"testing"

	"github.com/go-kit/kit/log"
)

func TestNopLogger(t *testing.T) {
	t.Parallel()
	logger := log.NewNopLogger()
	if err := logger.Log("abc", 123); err != nil {
		t.Error(err)
	}
	if err := log.NewContext(logger).With("def", "ghi").Log(); err != nil {
		t.Error(err)
	}
}
<commit_msg>log: add nop_logger_test for benchmark completeness<commit_after>package log_test

import (
	"testing"

	"github.com/go-kit/kit/log"
)

func TestNopLogger(t *testing.T) {
	t.Parallel()
	logger := log.NewNopLogger()
	if err := logger.Log("abc", 123); err != nil {
		t.Error(err)
	}
	if err := log.NewContext(logger).With("def", "ghi").Log(); err != nil {
		t.Error(err)
	}
}

func BenchmarkNopLoggerSimple(b *testing.B) {
	benchmarkRunner(b, log.NewNopLogger(), baseMessage)
}

func BenchmarkNopLoggerContextual(b *testing.B) {
	benchmarkRunner(b, log.NewNopLogger(), withMessage)
}
<|endoftext|>


<commit_before>// Benchmarks comparing append() to straight copy
//
// Run:
//   for i in $(seq 100 100 10000) ; do echo -- Size $i ; go test -bench . -- $i ; done > results.txt 2>&1
//
// With memory benchmarks:
//   for i in $(seq 100 100 10000) ; do echo -- Size $i ; go test -bench . -benchmem -- $i ; done > results.txt 2>&1
package main

import (
	"os"
	"strconv"
	"testing"
)

var (
	dimension = 10
)

func TestMain(m *testing.M) {
	// Yuck.  Size of the data to use for the struct is last arg.
	if d, err := strconv.Atoi(os.Args[len(os.Args)-1]); err == nil {
		dimension = d
	}
	os.Exit(m.Run())
}

// TestDummy avoids runtime warning about no tests
func TestDummy(t *testing.T) {
}

func bDoIdx(d int, b *testing.B) {
	setup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doIdx(d)
	}
}

func bDoAppend(d int, b *testing.B) {
	setup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doAppend()
	}
}

func bDoPolo(d int, b *testing.B) {
	psetup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doPolo()
	}
}

func bDoPoloAppend(d int, b *testing.B) {
	psetup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doPoloAppend()
	}
}

func BenchmarkDoIdx(b *testing.B) {
	bDoIdx(dimension, b)
}

func BenchmarkDoAppend(b *testing.B) {
	bDoAppend(dimension, b)
}
func BenchmarkPolo(b *testing.B) {
	bDoPolo(dimension, b)
}

func BenchmarkPoloAppend(b *testing.B) {
	bDoPoloAppend(dimension, b)
}
<commit_msg>Test to 10 mio.<commit_after>// Benchmarks comparing append() to straight copy
//
// Run:
//   for i in $(seq 100 100 10000000) ; do echo -- Size $i ; go test -bench . -- $i ; done > results.txt 2>&1
//
// With memory benchmarks:
//   for i in $(seq 100 100 10000) ; do echo -- Size $i ; go test -bench . -benchmem -- $i ; done > results.txt 2>&1
package main

import (
	"os"
	"strconv"
	"testing"
)

var (
	dimension = 10
)

func TestMain(m *testing.M) {
	// Yuck.  Size of the data to use for the struct is last arg.
	if d, err := strconv.Atoi(os.Args[len(os.Args)-1]); err == nil {
		dimension = d
	}
	os.Exit(m.Run())
}

// TestDummy avoids runtime warning about no tests
func TestDummy(t *testing.T) {
}

func bDoIdx(d int, b *testing.B) {
	setup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doIdx(d)
	}
}

func bDoAppend(d int, b *testing.B) {
	setup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doAppend()
	}
}

func bDoPolo(d int, b *testing.B) {
	psetup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doPolo()
	}
}

func bDoPoloAppend(d int, b *testing.B) {
	psetup(d)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		doPoloAppend()
	}
}

func BenchmarkDoIdx(b *testing.B) {
	bDoIdx(dimension, b)
}

func BenchmarkDoAppend(b *testing.B) {
	bDoAppend(dimension, b)
}
func BenchmarkPolo(b *testing.B) {
	bDoPolo(dimension, b)
}

func BenchmarkPoloAppend(b *testing.B) {
	bDoPoloAppend(dimension, b)
}
<|endoftext|>


<commit_before>package goexcel

import (
	"bytes"
	"encoding/csv"
	"errors"
	"net/http"
	"reflect"
	"strings"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"

	"golang.org/x/net/context"

	"github.com/news-ai/tabulae/models"
)

func csvGetCustomFields(r *http.Request, c context.Context, numberOfColumns int, headers []string) map[string]bool {
	customFields := make(map[string]bool, len(headers))

	for x := 0; x < numberOfColumns; x++ {
		columnName := headers[x]
		if !customOrNative(columnName) {
			customFields[columnName] = true
		}
	}
	return customFields
}

func csvRowToContact(r *http.Request, c context.Context, singleRow []string, headers []string) (models.Contact, error) {
	var (
		contact       models.Contact
		employers     []int64
		pastEmployers []int64
		customFields  []models.CustomContactField
	)

	for x := 0; x < len(singleRow); x++ {
		columnName := headers[x]
		cellName := singleRow[x]
		rowToContact(r, c, columnName, cellName, &contact, &employers, &pastEmployers, &customFields)
	}

	contact.CustomFields = customFields
	contact.Employers = employers
	contact.PastEmployers = pastEmployers
	return contact, nil
}

func CsvToContactList(r *http.Request, file []byte, headers []string) ([]models.Contact, map[string]bool, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Infof(c, "%v", err)
		return []models.Contact{}, map[string]bool{}, err
	}

	// Number of columns in sheet to compare
	numberOfColumns := len(records[0])
	if numberOfColumns != len(headers) {
		err := errors.New("Number of headers does not match the ones for the sheet")
		log.Infof(c, "%v", err)
		return []models.Contact{}, map[string]bool{}, err
	}

	// Loop through all the rows
	// Extract information
	emptyContact := models.Contact{}
	contacts := []models.Contact{}

	for i := 0; i < len(records); i++ {
		contact, err := csvRowToContact(r, c, records[0], headers)
		if err != nil {
			log.Infof(c, "%v", err)
			return []models.Contact{}, map[string]bool{}, err
		}

		// To get rid of empty contacts. We don't want to create empty contacts.
		if !reflect.DeepEqual(emptyContact, contact) {
			contacts = append(contacts, contact)
		}
	}

	// Get custom fields
	customFields := getCustomFields(r, c, numberOfColumns, headers)

	return contacts, customFields, nil
}

func CsvFileToExcelHeader(r *http.Request, file []byte) ([]Column, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Errorf(c, "%v", err)
		return []Column{}, err
	}

	numberOfRows := 15
	if len(records) < numberOfRows+1 {
		numberOfRows = len(records)
	}

	numberOfColumns := len(records[0])
	columns := make([]Column, numberOfColumns)

	for _, row := range records[0:numberOfRows] {
		for currentColumn, cell := range row {
			columns[currentColumn].Rows = append(columns[currentColumn].Rows, strings.Trim(cell, " "))
		}
	}

	return columns, nil
}
<commit_msg>Stupid error only importing a single record :fu: life over<commit_after>package goexcel

import (
	"bytes"
	"encoding/csv"
	"errors"
	"net/http"
	"reflect"
	"strings"

	"google.golang.org/appengine"
	"google.golang.org/appengine/log"

	"golang.org/x/net/context"

	"github.com/news-ai/tabulae/models"
)

func csvGetCustomFields(r *http.Request, c context.Context, numberOfColumns int, headers []string) map[string]bool {
	customFields := make(map[string]bool, len(headers))

	for x := 0; x < numberOfColumns; x++ {
		columnName := headers[x]
		if !customOrNative(columnName) {
			customFields[columnName] = true
		}
	}
	return customFields
}

func csvRowToContact(r *http.Request, c context.Context, singleRow []string, headers []string) (models.Contact, error) {
	var (
		contact       models.Contact
		employers     []int64
		pastEmployers []int64
		customFields  []models.CustomContactField
	)

	for x := 0; x < len(singleRow); x++ {
		columnName := headers[x]
		cellName := singleRow[x]
		rowToContact(r, c, columnName, cellName, &contact, &employers, &pastEmployers, &customFields)
	}

	contact.CustomFields = customFields
	contact.Employers = employers
	contact.PastEmployers = pastEmployers
	return contact, nil
}

func CsvToContactList(r *http.Request, file []byte, headers []string) ([]models.Contact, map[string]bool, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Infof(c, "%v", err)
		return []models.Contact{}, map[string]bool{}, err
	}

	// Number of columns in sheet to compare
	numberOfColumns := len(records[0])
	if numberOfColumns != len(headers) {
		err := errors.New("Number of headers does not match the ones for the sheet")
		log.Infof(c, "%v", err)
		return []models.Contact{}, map[string]bool{}, err
	}

	// Loop through all the rows
	// Extract information
	emptyContact := models.Contact{}
	contacts := []models.Contact{}

	for i := 0; i < len(records); i++ {
		contact, err := csvRowToContact(r, c, records[i], headers)
		if err != nil {
			log.Infof(c, "%v", err)
			return []models.Contact{}, map[string]bool{}, err
		}

		// To get rid of empty contacts. We don't want to create empty contacts.
		if !reflect.DeepEqual(emptyContact, contact) {
			contacts = append(contacts, contact)
		}
	}

	// Get custom fields
	customFields := getCustomFields(r, c, numberOfColumns, headers)

	return contacts, customFields, nil
}

func CsvFileToExcelHeader(r *http.Request, file []byte) ([]Column, error) {
	c := appengine.NewContext(r)

	readerFile := bytes.NewReader(file)
	incomingRecords := csv.NewReader(readerFile)
	records, err := incomingRecords.ReadAll()
	if err != nil {
		log.Errorf(c, "%v", err)
		return []Column{}, err
	}

	numberOfRows := 15
	if len(records) < numberOfRows+1 {
		numberOfRows = len(records)
	}

	numberOfColumns := len(records[0])
	columns := make([]Column, numberOfColumns)

	for _, row := range records[0:numberOfRows] {
		for currentColumn, cell := range row {
			columns[currentColumn].Rows = append(columns[currentColumn].Rows, strings.Trim(cell, " "))
		}
	}

	return columns, nil
}
<|endoftext|>