<commit_before>package csvtrans

import (
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"
)

// RowTransformer is a function that transforms a CSV row.
type RowTransformer func(i int, inRow []string) (outRow []string, err error)

// BufRowTransformer is similar to RowTransformer, but also takes a buffer which
// should be returned. This allows the output CSV to reuse a buffer between
// calls, to reduce garbage
type BufRowTransformer func(i int, inRow []string, outRow []string) ([]string, error)

// Creates a RowTransformer from a BufRowTransformer which reuses the same buffer.
// Note that when using a BufRowTransformer with this that every field must be
// set each iteration, or values will be repeated between rows.
// length is the number of columns in the output CSV
func MakeRowTransformer(length int, f BufRowTransformer) RowTransformer {
	buf := make([]string, length, length)
	return func(i int, inRow []string) ([]string, error) {
		return f(i, inRow, buf)
	}
}

// Run performs the transformation on the CSV given an input reader, output
// writer and a function to transform each row.
func Run(in io.Reader, out io.Writer, f RowTransformer) error {
	inCsv := csv.NewReader(in)
	outCsv := csv.NewWriter(out)

	defer outCsv.Flush()
	i := 0
	for {
		row, err := inCsv.Read()

		if err == io.EOF {
			return nil
		} else if err != nil {
			return fmt.Errorf("Error reading CSV row at index %d: %v", i, err)
		}

		transformed, err := f(i, row)
		if err != nil {
			return fmt.Errorf("Error transforming row at index %d: %v", i, err)
		}

		err = outCsv.Write(transformed)
		if err != nil {
			return fmt.Errorf("Error writing CSV row at index %d: %v", i, err)
		}
		i += 1
	}
}

// RunFile is a wrapper around Run. This uses files as the input and output.
// Different filenames must be passed in to the function.
// If an output file already exists, it will be replaced.
func RunFile(inFile, outFile string, f RowTransformer) error {
	if inFile == outFile {
		return errors.New("inFile and outFile must be different")
	}

	in, err := os.Open(inFile) // Read access
	if err != nil {
		return fmt.Errorf("Error reading %s: %v", inFile, err)
	}
	defer in.Close()

	out, err := os.Create(outFile)
	if err != nil {
		return fmt.Errorf("Error opening %s for write: %v", outFile, err)
	}
	defer out.Close()

	err = Run(in, out, f)
	if err != nil {
		return err
	}
	return nil
}

<commit_msg>Added ability to skip a row when transforming CSV file<commit_after>package csvtrans

import (
	"encoding/csv"
	"errors"
	"fmt"
	"io"
	"os"
)

// RowTransformer is a function that transforms a CSV row.
// If an error is returned, the CSV transformation will terminate.
// If a row is to be skipped, return with outRow and err == nil
type RowTransformer func(i int, inRow []string) (outRow []string, err error)

// BufRowTransformer is similar to RowTransformer, but also takes a buffer which
// should be returned. This allows the output CSV to reuse a buffer between
// calls, to reduce garbage
type BufRowTransformer func(i int, inRow []string, outRow []string) ([]string, error)

// Creates a RowTransformer from a BufRowTransformer which reuses the same buffer.
// Note that when using a BufRowTransformer with this that every field must be
// set each iteration, or values will be repeated between rows.
// length is the number of columns in the output CSV
func MakeRowTransformer(length int, f BufRowTransformer) RowTransformer {
	buf := make([]string, length, length)
	return func(i int, inRow []string) ([]string, error) {
		return f(i, inRow, buf)
	}
}

// Run performs the transformation on the CSV given an input reader, output
// writer and a function to transform each row.
func Run(in io.Reader, out io.Writer, f RowTransformer) error {
	inCsv := csv.NewReader(in)
	outCsv := csv.NewWriter(out)

	defer outCsv.Flush()
	i := 0
	for {
		row, err := inCsv.Read()

		if err == io.EOF {
			return nil
		} else if err != nil {
			return fmt.Errorf("Error reading CSV row at index %d: %v", i, err)
		}

		transformed, err := f(i, row)
		if err != nil {
			return fmt.Errorf("Error transforming row at index %d: %v", i, err)
		}

		// Skip if row is nil and no error
		if transformed != nil {
			err = outCsv.Write(transformed)
			if err != nil {
				return fmt.Errorf("Error writing CSV row at index %d: %v", i, err)
			}
		}

		i += 1
	}
}

// RunFile is a wrapper around Run. This uses files as the input and output.
// Different filenames must be passed in to the function.
// If an output file already exists, it will be replaced.
func RunFile(inFile, outFile string, f RowTransformer) error {
	if inFile == outFile {
		return errors.New("inFile and outFile must be different")
	}

	in, err := os.Open(inFile) // Read access
	if err != nil {
		return fmt.Errorf("Error reading %s: %v", inFile, err)
	}
	defer in.Close()

	out, err := os.Create(outFile)
	if err != nil {
		return fmt.Errorf("Error opening %s for write: %v", outFile, err)
	}
	defer out.Close()

	err = Run(in, out, f)
	if err != nil {
		return err
	}
	return nil
}

<|endoftext|>


<commit_before>package ctxerr

import (
	"fmt"

	"github.com/fatih/color"
)

func init() {
	color.NoColor = true
}

func ExamplePoint() {
	fmt.Println(New("00100", Point(1, 3)))
	// Output:
	// 1:3:
	// 1 | 00100
	//   |   ^
}

func ExampleRange() {
	fmt.Println(New("01110", Range(1, 2, 1, 4)))
	// Output:
	// 1:2-4:
	// 1 | 01110
	//   |  ^^^
}

func ExampleRange_multiline() {
	fmt.Println(New("00001\n11110", Range(1, 5, 2, 4)))
	// Output:
	// 1:5-2:4:
	// 1 | 00001
	//   |     ^
	// 2 | 11110
	//   | ^^^^
}

func ExampleCtx_WithHint() {
	ctx := New("010101102110", Point(1, 9))
	ctx.Hint = "don't worry, bender"
	fmt.Println(ctx)
	// Output:
	// 1:9:
	// 1 | 010101102110
	//   |         ^ don't worry, bender
}

func ExampleCtx_WithContext_limited() {
	in := `1st
2nd
3rd has an error
4th
5th`
	ctx := New(in, Point(3, 12))
	ctx.Context = 1
	fmt.Println(ctx)
	// Output:
	// 3:12:
	// 2 | 2nd
	// 3 | 3rd has an error
	//   |            ^
	// 4 | 4th
}

func ExampleCtx_WithContext_limitedMultiline() {
	in := `1st
2nd
3rd has an error
4th still has an error
5th`
	ctx := New(in, Range(3, 1, 4, 22))
	ctx.Context = 1
	fmt.Println(ctx)
	// Output:
	// 3:1-4:22:
	// 2 | 2nd
	// 3 | 3rd has an error
	//   | ^^^^^^^^^^^^^^^^
	// 4 | 4th still has an error
	//   | ^^^^^^^^^^^^^^^^^^^^^^
	// 5 | 5th
}

func ExampleCtx_WithContext_all() {
	in := `1st
2nd
3rd has an error
4th
5th`
	ctx := New(in, Point(3, 12))
	ctx.Context = -1
	fmt.Println(ctx)
	// Output:
	// 3:12:
	// 1 | 1st
	// 2 | 2nd
	// 3 | 3rd has an error
	//   |            ^
	// 4 | 4th
	// 5 | 5th
}

func ExampleCtx_Path() {
	ctx := New("42", Point(1, 1))
	ctx.Path = "/tmp/ctxerr/answer.txt"
	fmt.Println(ctx)
	// Output:
	// /tmp/ctxerr/answer.txt:1:1:
	// 1 | 42
	//   | ^
}

func ExampleCtx_ToError() {
	cause := fmt.Errorf("not a letter")
	err := New("ab!cd", Point(1, 3)).ToError(cause)
	fmt.Println(err)
	// Output:
	// not a letter
	// 1:3:
	// 1 | ab!cd
	//   |   ^
}
<commit_msg>Rename example methods<commit_after>package ctxerr

import (
	"fmt"

	"github.com/fatih/color"
)

func init() {
	color.NoColor = true
}

func ExamplePoint() {
	fmt.Println(New("00100", Point(1, 3)))
	// Output:
	// 1:3:
	// 1 | 00100
	//   |   ^
}

func ExampleRange() {
	fmt.Println(New("01110", Range(1, 2, 1, 4)))
	// Output:
	// 1:2-4:
	// 1 | 01110
	//   |  ^^^
}

func ExampleRange_multiline() {
	fmt.Println(New("00001\n11110", Range(1, 5, 2, 4)))
	// Output:
	// 1:5-2:4:
	// 1 | 00001
	//   |     ^
	// 2 | 11110
	//   | ^^^^
}

func ExampleCtx_String_hint() {
	ctx := New("010101102110", Point(1, 9))
	ctx.Hint = "don't worry, bender"
	fmt.Println(ctx)
	// Output:
	// 1:9:
	// 1 | 010101102110
	//   |         ^ don't worry, bender
}

func ExampleCtx_String_contextLimited() {
	in := `1st
2nd
3rd has an error
4th
5th`
	ctx := New(in, Point(3, 12))
	ctx.Context = 1
	fmt.Println(ctx)
	// Output:
	// 3:12:
	// 2 | 2nd
	// 3 | 3rd has an error
	//   |            ^
	// 4 | 4th
}

func ExampleCtx_String_contextLimitedMultiline() {
	in := `1st
2nd
3rd has an error
4th still has an error
5th`
	ctx := New(in, Range(3, 1, 4, 22))
	ctx.Context = 1
	fmt.Println(ctx)
	// Output:
	// 3:1-4:22:
	// 2 | 2nd
	// 3 | 3rd has an error
	//   | ^^^^^^^^^^^^^^^^
	// 4 | 4th still has an error
	//   | ^^^^^^^^^^^^^^^^^^^^^^
	// 5 | 5th
}

func ExampleCtx_String_contextAll() {
	in := `1st
2nd
3rd has an error
4th
5th`
	ctx := New(in, Point(3, 12))
	ctx.Context = -1
	fmt.Println(ctx)
	// Output:
	// 3:12:
	// 1 | 1st
	// 2 | 2nd
	// 3 | 3rd has an error
	//   |            ^
	// 4 | 4th
	// 5 | 5th
}

func ExampleCtx_String_path() {
	ctx := New("42", Point(1, 1))
	ctx.Path = "/tmp/ctxerr/answer.txt"
	fmt.Println(ctx)
	// Output:
	// /tmp/ctxerr/answer.txt:1:1:
	// 1 | 42
	//   | ^
}

func ExampleCtx_ToError() {
	cause := fmt.Errorf("not a letter")
	err := New("ab!cd", Point(1, 3)).ToError(cause)
	fmt.Println(err)
	// Output:
	// not a letter
	// 1:3:
	// 1 | ab!cd
	//   |   ^
}
<|endoftext|>


<commit_before>package roaring

import (
	"encoding/binary"
	"fmt"
	"math/rand"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestCountTrailingZeros072(t *testing.T) {
	Convey("countTrailingZeros", t, func() {
		So(countTrailingZerosAsm(0), ShouldEqual, 64)
		So(countTrailingZerosAsm(8), ShouldEqual, 3)
		So(countTrailingZerosAsm(7), ShouldEqual, 0)
		So(countTrailingZerosAsm(1<<17), ShouldEqual, 17)
		So(countTrailingZerosAsm(7<<17), ShouldEqual, 17)
		So(countTrailingZerosAsm(255<<33), ShouldEqual, 33)

		So(countTrailingZerosDeBruijn(0), ShouldEqual, 64)
		So(countTrailingZerosDeBruijn(8), ShouldEqual, 3)
		So(countTrailingZerosDeBruijn(7), ShouldEqual, 0)
		So(countTrailingZerosDeBruijn(1<<17), ShouldEqual, 17)
		So(countTrailingZerosDeBruijn(7<<17), ShouldEqual, 17)
		So(countTrailingZerosDeBruijn(255<<33), ShouldEqual, 33)

	})
}

func getRandomUint64Set(n int) []uint64 {
	seed := int64(42)
	p("seed is %v", seed)
	rand.Seed(seed)

	var buf [8]byte
	var o []uint64
	for i := 0; i < n; i++ {
		rand.Read(buf[:])
		o = append(o, binary.LittleEndian.Uint64(buf[:]))
	}
	return o
}

func getAllOneBitUint64Set() []uint64 {
	o := []uint64{0}
	for i := uint(0); i < 64; i++ {
		o = append(o, 1<<i)
	}
	return o
}

func Benchmark100OrigNumberOfTrailingZeros(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			numberOfTrailingZeros(r[i])
		}
	}
}

func Benchmark100CountTrailingZerosDeBruijn(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			countTrailingZerosDeBruijn(r[i])
		}
	}
}

func Benchmark100CountTrailingZerosAsm(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			countTrailingZerosAsm(r[i])
		}
	}
}

func numberOfTrailingZeros(i uint64) int {
	if i == 0 {
		return 64
	}
	x := i
	n := int64(63)
	y := x << 32
	if y != 0 {
		n -= 32
		x = y
	}
	y = x << 16
	if y != 0 {
		n -= 16
		x = y
	}
	y = x << 8
	if y != 0 {
		n -= 8
		x = y
	}
	y = x << 4
	if y != 0 {
		n -= 4
		x = y
	}
	y = x << 2
	if y != 0 {
		n -= 2
		x = y
	}
	return int(n - int64(uint64(x<<1)>>63))
}

/*
//
// on an Intel(R) Core(TM) i7-5557U CPU @ 3.10GHz:
//
Benchmark100CountTrailingZerosDeBruijn-4   	10000000	       168 ns/op
Benchmark100CountTrailingZerosAsm-4        	 5000000	       278 ns/op
Benchmark100OrigNumberOfTrailingZeros-4    	 3000000	       592 ns/op

// and again:

Benchmark100CountTrailingZerosDeBruijn-4   	10000000	       168 ns/op
Benchmark100CountTrailingZerosAsm-4        	 5000000	       278 ns/op
Benchmark100OrigNumberOfTrailingZeros-4    	 3000000	       585 ns/op
*/
// go test -v -bench=100 -run 101
func Test101CountTrailingZerosCorrectness(t *testing.T) {
	r := getAllOneBitUint64Set()
	for i, v := range r {
		a := countTrailingZerosDeBruijn(v)
		b := countTrailingZerosAsm(v)
		if a != b {
			panic(fmt.Errorf("on r[%v]= v=%v,  a: %v, b:%v", i, v, a, b))
		}
	}
}
<commit_msg>atg. skip Asm check for 0 so we don't fail on older BSR based implementations of TZCNTQ<commit_after>package roaring

import (
	"encoding/binary"
	"fmt"
	"math/rand"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestCountTrailingZeros072(t *testing.T) {
	Convey("countTrailingZeros", t, func() {
		// undefined on older cpus, so skip this check on 0.
		//So(countTrailingZerosAsm(0), ShouldEqual, 64)

		So(countTrailingZerosAsm(8), ShouldEqual, 3)
		So(countTrailingZerosAsm(7), ShouldEqual, 0)
		So(countTrailingZerosAsm(1<<17), ShouldEqual, 17)
		So(countTrailingZerosAsm(7<<17), ShouldEqual, 17)
		So(countTrailingZerosAsm(255<<33), ShouldEqual, 33)

		So(countTrailingZerosDeBruijn(0), ShouldEqual, 64)
		So(countTrailingZerosDeBruijn(8), ShouldEqual, 3)
		So(countTrailingZerosDeBruijn(7), ShouldEqual, 0)
		So(countTrailingZerosDeBruijn(1<<17), ShouldEqual, 17)
		So(countTrailingZerosDeBruijn(7<<17), ShouldEqual, 17)
		So(countTrailingZerosDeBruijn(255<<33), ShouldEqual, 33)

	})
}

func getRandomUint64Set(n int) []uint64 {
	seed := int64(42)
	p("seed is %v", seed)
	rand.Seed(seed)

	var buf [8]byte
	var o []uint64
	for i := 0; i < n; i++ {
		rand.Read(buf[:])
		o = append(o, binary.LittleEndian.Uint64(buf[:]))
	}
	return o
}

func getAllOneBitUint64Set() []uint64 {
	o := []uint64{0}
	for i := uint(0); i < 64; i++ {
		o = append(o, 1<<i)
	}
	return o
}

func Benchmark100OrigNumberOfTrailingZeros(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			numberOfTrailingZeros(r[i])
		}
	}
}

func Benchmark100CountTrailingZerosDeBruijn(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			countTrailingZerosDeBruijn(r[i])
		}
	}
}

func Benchmark100CountTrailingZerosAsm(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			countTrailingZerosAsm(r[i])
		}
	}
}

func numberOfTrailingZeros(i uint64) int {
	if i == 0 {
		return 64
	}
	x := i
	n := int64(63)
	y := x << 32
	if y != 0 {
		n -= 32
		x = y
	}
	y = x << 16
	if y != 0 {
		n -= 16
		x = y
	}
	y = x << 8
	if y != 0 {
		n -= 8
		x = y
	}
	y = x << 4
	if y != 0 {
		n -= 4
		x = y
	}
	y = x << 2
	if y != 0 {
		n -= 2
		x = y
	}
	return int(n - int64(uint64(x<<1)>>63))
}

/*
//
// on an Intel(R) Core(TM) i7-5557U CPU @ 3.10GHz:
//
Benchmark100CountTrailingZerosDeBruijn-4   	10000000	       168 ns/op
Benchmark100CountTrailingZerosAsm-4        	 5000000	       278 ns/op
Benchmark100OrigNumberOfTrailingZeros-4    	 3000000	       592 ns/op

// and again:

Benchmark100CountTrailingZerosDeBruijn-4   	10000000	       168 ns/op
Benchmark100CountTrailingZerosAsm-4        	 5000000	       278 ns/op
Benchmark100OrigNumberOfTrailingZeros-4    	 3000000	       585 ns/op
*/
// go test -v -bench=100 -run 101
func Test101CountTrailingZerosCorrectness(t *testing.T) {
	r := getAllOneBitUint64Set()
	for i, v := range r {
		a := countTrailingZerosDeBruijn(v)
		b := countTrailingZerosAsm(v)
		if a != b {
			panic(fmt.Errorf("on r[%v]= v=%v,  a: %v, b:%v", i, v, a, b))
		}
	}
}
<|endoftext|>


<commit_before>package roaring

import (
	"encoding/binary"
	"math/rand"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestCountTrailingZeros072(t *testing.T) {
	Convey("countTrailingZeros", t, func() {
		So(numberOfTrailingZeros(0), ShouldEqual, 64)
		So(numberOfTrailingZeros(8), ShouldEqual, 3)
		So(numberOfTrailingZeros(7), ShouldEqual, 0)
		So(numberOfTrailingZeros(1<<17), ShouldEqual, 17)
		So(numberOfTrailingZeros(7<<17), ShouldEqual, 17)
		So(numberOfTrailingZeros(255<<33), ShouldEqual, 33)

		So(countTrailingZeros(0), ShouldEqual, 64)
		So(countTrailingZeros(8), ShouldEqual, 3)
		So(countTrailingZeros(7), ShouldEqual, 0)
		So(countTrailingZeros(1<<17), ShouldEqual, 17)
		So(countTrailingZeros(7<<17), ShouldEqual, 17)
		So(countTrailingZeros(255<<33), ShouldEqual, 33)

	})
}

func getRandomUint64Set(n int) []uint64 {
	seed := int64(42)
	p("seed is %v", seed)
	rand.Seed(seed)

	var buf [8]byte
	var o []uint64
	for i := 0; i < n; i++ {
		rand.Read(buf[:])
		o = append(o, binary.LittleEndian.Uint64(buf[:]))
	}
	return o
}

func getAllOneBitUint64Set() []uint64 {
	var o []uint64
	for i := uint(0); i < 64; i++ {
		o = append(o, 1<<i)
	}
	return o
}

func Benchmark100OrigNumberOfTrailingZeros(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			numberOfTrailingZeros(r[i])
		}
	}
}

func Benchmark100CountTrailingZerosDeBruijn(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			countTrailingZeros(r[i])
		}
	}
}

func numberOfTrailingZeros(i uint64) int {
	if i == 0 {
		return 64
	}
	x := i
	n := int64(63)
	y := x << 32
	if y != 0 {
		n -= 32
		x = y
	}
	y = x << 16
	if y != 0 {
		n -= 16
		x = y
	}
	y = x << 8
	if y != 0 {
		n -= 8
		x = y
	}
	y = x << 4
	if y != 0 {
		n -= 4
		x = y
	}
	y = x << 2
	if y != 0 {
		n -= 2
		x = y
	}
	return int(n - int64(uint64(x<<1)>>63))
}
<commit_msg>Remove 'DeBruijn' from Benchmark100CountTrailingZeros func name<commit_after>package roaring

import (
	"encoding/binary"
	"math/rand"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestCountTrailingZeros072(t *testing.T) {
	Convey("countTrailingZeros", t, func() {
		So(numberOfTrailingZeros(0), ShouldEqual, 64)
		So(numberOfTrailingZeros(8), ShouldEqual, 3)
		So(numberOfTrailingZeros(7), ShouldEqual, 0)
		So(numberOfTrailingZeros(1<<17), ShouldEqual, 17)
		So(numberOfTrailingZeros(7<<17), ShouldEqual, 17)
		So(numberOfTrailingZeros(255<<33), ShouldEqual, 33)

		So(countTrailingZeros(0), ShouldEqual, 64)
		So(countTrailingZeros(8), ShouldEqual, 3)
		So(countTrailingZeros(7), ShouldEqual, 0)
		So(countTrailingZeros(1<<17), ShouldEqual, 17)
		So(countTrailingZeros(7<<17), ShouldEqual, 17)
		So(countTrailingZeros(255<<33), ShouldEqual, 33)

	})
}

func getRandomUint64Set(n int) []uint64 {
	seed := int64(42)
	p("seed is %v", seed)
	rand.Seed(seed)

	var buf [8]byte
	var o []uint64
	for i := 0; i < n; i++ {
		rand.Read(buf[:])
		o = append(o, binary.LittleEndian.Uint64(buf[:]))
	}
	return o
}

func getAllOneBitUint64Set() []uint64 {
	var o []uint64
	for i := uint(0); i < 64; i++ {
		o = append(o, 1<<i)
	}
	return o
}

func Benchmark100OrigNumberOfTrailingZeros(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			numberOfTrailingZeros(r[i])
		}
	}
}

func Benchmark100CountTrailingZeros(b *testing.B) {
	b.StopTimer()

	r := getRandomUint64Set(64)
	r = append(r, getAllOneBitUint64Set()...)

	b.ResetTimer()
	b.StartTimer()
	for i := 0; i < b.N; i++ {
		for i := range r {
			countTrailingZeros(r[i])
		}
	}
}

func numberOfTrailingZeros(i uint64) int {
	if i == 0 {
		return 64
	}
	x := i
	n := int64(63)
	y := x << 32
	if y != 0 {
		n -= 32
		x = y
	}
	y = x << 16
	if y != 0 {
		n -= 16
		x = y
	}
	y = x << 8
	if y != 0 {
		n -= 8
		x = y
	}
	y = x << 4
	if y != 0 {
		n -= 4
		x = y
	}
	y = x << 2
	if y != 0 {
		n -= 2
		x = y
	}
	return int(n - int64(uint64(x<<1)>>63))
}
<|endoftext|>


<commit_before>package kafka

import (
	"bufio"
	"io"
	"sync"
	"time"
)

// A Batch is an iterator over a sequence of messages fetched from a kafka
// server.
//
// Batches are created by calling (*Conn).ReadBatch. They hold a internal lock
// on the connection, which is released when the batch is closed. Failing to
// call a batch's Close method will likely result in a dead-lock when trying to
// use the connection.
//
// Batches are safe to use concurrently from multiple goroutines.
type Batch struct {
	mutex         sync.Mutex
	conn          *Conn
	lock          *sync.Mutex
	msgs          *messageSetReader
	deadline      time.Time
	throttle      time.Duration
	topic         string
	partition     int
	offset        int64
	highWaterMark int64
	err           error
}

// Throttle gives the throttling duration applied by the kafka server on the
// connection.
func (batch *Batch) Throttle() time.Duration {
	return batch.throttle
}

// Watermark returns the current highest watermark in a partition.
func (batch *Batch) HighWaterMark() int64 {
	return batch.highWaterMark
}

// Offset returns the offset of the next message in the batch.
func (batch *Batch) Offset() int64 {
	batch.mutex.Lock()
	offset := batch.offset
	batch.mutex.Unlock()
	return offset
}

// Close closes the batch, releasing the connection lock and returning an error
// if reading the batch failed for any reason.
func (batch *Batch) Close() error {
	batch.mutex.Lock()
	err := batch.close()
	batch.mutex.Unlock()
	return err
}

func (batch *Batch) close() (err error) {
	conn := batch.conn
	lock := batch.lock

	batch.conn = nil
	batch.lock = nil
	if batch.msgs != nil {
		batch.msgs.discard()
	}

	if err = batch.err; err == io.EOF {
		err = nil
	}

	if conn != nil {
		conn.rdeadline.unsetConnReadDeadline()
		conn.mutex.Lock()
		conn.offset = batch.offset
		conn.mutex.Unlock()

		if err != nil {
			if _, ok := err.(Error); !ok && err != io.ErrShortBuffer {
				conn.Close()
			}
		}
	}

	if lock != nil {
		lock.Unlock()
	}

	return
}

// Err returns a non-nil error if the batch is broken. This is the same error
// that would be returned by Read, ReadMessage or Close (except in the case of
// io.EOF which is never returned by Close).
//
// This method is useful when building retry mechanisms for (*Conn).ReadBatch,
// the program can check whether the batch carried a error before attempting to
// read the first message.
//
// Note that checking errors on a batch is optional, calling Read or ReadMessage
// is always valid and can be used to either read a message or an error in cases
// where that's convenient.
func (batch *Batch) Err() error { return batch.err }

// Read reads the value of the next message from the batch into b, returning the
// number of bytes read, or an error if the next message couldn't be read.
//
// If an error is returned the batch cannot be used anymore and calling Read
// again will keep returning that error. All errors except io.EOF (indicating
// that the program consumed all messages from the batch) are also returned by
// Close.
//
// The method fails with io.ErrShortBuffer if the buffer passed as argument is
// too small to hold the message value.
func (batch *Batch) Read(b []byte) (int, error) {
	n := 0

	batch.mutex.Lock()
	offset := batch.offset

	_, _, _, err := batch.readMessage(
		func(r *bufio.Reader, size int, nbytes int) (int, error) {
			if nbytes < 0 {
				return size, nil
			}
			return discardN(r, size, nbytes)
		},
		func(r *bufio.Reader, size int, nbytes int) (int, error) {
			if nbytes < 0 {
				return size, nil
			}
			n = nbytes // return value
			if nbytes > cap(b) {
				nbytes = cap(b)
			}
			if nbytes > len(b) {
				b = b[:nbytes]
			}
			nbytes, err := io.ReadFull(r, b[:nbytes])
			if err != nil {
				return size - nbytes, err
			}
			return discardN(r, size-nbytes, n-nbytes)
		},
	)

	if err == nil && n > len(b) {
		n, err = len(b), io.ErrShortBuffer
		batch.err = io.ErrShortBuffer
		batch.offset = offset // rollback
	}

	batch.mutex.Unlock()
	return n, err
}

// ReadMessage reads and return the next message from the batch.
//
// Because this method allocate memory buffers for the message key and value
// it is less memory-efficient than Read, but has the advantage of never
// failing with io.ErrShortBuffer.
func (batch *Batch) ReadMessage() (Message, error) {
	msg := Message{}
	batch.mutex.Lock()

	var offset, timestamp int64
	var headers []Header
	var err error

	offset, timestamp, headers, err = batch.readMessage(
		func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
			msg.Key, remain, err = readNewBytes(r, size, nbytes)
			return
		},
		func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
			msg.Value, remain, err = readNewBytes(r, size, nbytes)
			return
		},
	)
	for batch.conn != nil && offset < batch.conn.offset {
		if err != nil {
			break
		}
		offset, timestamp, headers, err = batch.readMessage(
			func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
				msg.Key, remain, err = readNewBytes(r, size, nbytes)
				return
			},
			func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
				msg.Value, remain, err = readNewBytes(r, size, nbytes)
				return
			},
		)
	}

	batch.mutex.Unlock()
	msg.Topic = batch.topic
	msg.Partition = batch.partition
	msg.Offset = offset
	msg.Time = timestampToTime(timestamp)
	msg.Headers = headers

	return msg, err
}

func (batch *Batch) readMessage(
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, timestamp int64, headers []Header, err error) {
	if err = batch.err; err != nil {
		return
	}

	offset, timestamp, headers, err = batch.msgs.readMessage(batch.offset, key, val)
	switch err {
	case nil:
		batch.offset = offset + 1
	case errShortRead:
		// As an "optimization" kafka truncates the returned response after
		// producing MaxBytes, which could then cause the code to return
		// errShortRead.
		err = batch.msgs.discard()
		switch {
		case err != nil:
			batch.err = err
		case batch.msgs.remaining() == 0:
			// Because we use the adjusted deadline we could end up returning
			// before the actual deadline occurred. This is necessary otherwise
			// timing out the connection for real could end up leaving it in an
			// unpredictable state, which would require closing it.
			// This design decision was made to maximize the chances of keeping
			// the connection open, the trade off being to lose precision on the
			// read deadline management.
			err = checkTimeoutErr(batch.deadline)
			batch.err = err
		}
	default:
		batch.err = err
	}

	return
}

func checkTimeoutErr(deadline time.Time) (err error) {
	if !deadline.IsZero() && time.Now().After(deadline) {
		err = RequestTimedOut
	} else {
		err = io.EOF
	}
	return
}
<commit_msg>Fix a read timeout on short read in Batch.Read() (#267)<commit_after>package kafka

import (
	"bufio"
	"io"
	"sync"
	"time"
)

// A Batch is an iterator over a sequence of messages fetched from a kafka
// server.
//
// Batches are created by calling (*Conn).ReadBatch. They hold a internal lock
// on the connection, which is released when the batch is closed. Failing to
// call a batch's Close method will likely result in a dead-lock when trying to
// use the connection.
//
// Batches are safe to use concurrently from multiple goroutines.
type Batch struct {
	mutex         sync.Mutex
	conn          *Conn
	lock          *sync.Mutex
	msgs          *messageSetReader
	deadline      time.Time
	throttle      time.Duration
	topic         string
	partition     int
	offset        int64
	highWaterMark int64
	err           error
}

// Throttle gives the throttling duration applied by the kafka server on the
// connection.
func (batch *Batch) Throttle() time.Duration {
	return batch.throttle
}

// Watermark returns the current highest watermark in a partition.
func (batch *Batch) HighWaterMark() int64 {
	return batch.highWaterMark
}

// Offset returns the offset of the next message in the batch.
func (batch *Batch) Offset() int64 {
	batch.mutex.Lock()
	offset := batch.offset
	batch.mutex.Unlock()
	return offset
}

// Close closes the batch, releasing the connection lock and returning an error
// if reading the batch failed for any reason.
func (batch *Batch) Close() error {
	batch.mutex.Lock()
	err := batch.close()
	batch.mutex.Unlock()
	return err
}

func (batch *Batch) close() (err error) {
	conn := batch.conn
	lock := batch.lock

	batch.conn = nil
	batch.lock = nil
	if batch.msgs != nil {
		batch.msgs.discard()
	}

	if err = batch.err; err == io.EOF {
		err = nil
	}

	if conn != nil {
		conn.rdeadline.unsetConnReadDeadline()
		conn.mutex.Lock()
		conn.offset = batch.offset
		conn.mutex.Unlock()

		if err != nil {
			if _, ok := err.(Error); !ok && err != io.ErrShortBuffer {
				conn.Close()
			}
		}
	}

	if lock != nil {
		lock.Unlock()
	}

	return
}

// Err returns a non-nil error if the batch is broken. This is the same error
// that would be returned by Read, ReadMessage or Close (except in the case of
// io.EOF which is never returned by Close).
//
// This method is useful when building retry mechanisms for (*Conn).ReadBatch,
// the program can check whether the batch carried a error before attempting to
// read the first message.
//
// Note that checking errors on a batch is optional, calling Read or ReadMessage
// is always valid and can be used to either read a message or an error in cases
// where that's convenient.
func (batch *Batch) Err() error { return batch.err }

// Read reads the value of the next message from the batch into b, returning the
// number of bytes read, or an error if the next message couldn't be read.
//
// If an error is returned the batch cannot be used anymore and calling Read
// again will keep returning that error. All errors except io.EOF (indicating
// that the program consumed all messages from the batch) are also returned by
// Close.
//
// The method fails with io.ErrShortBuffer if the buffer passed as argument is
// too small to hold the message value.
func (batch *Batch) Read(b []byte) (int, error) {
	n := 0

	batch.mutex.Lock()
	offset := batch.offset

	_, _, _, err := batch.readMessage(
		func(r *bufio.Reader, size int, nbytes int) (int, error) {
			if nbytes < 0 {
				return size, nil
			}
			return discardN(r, size, nbytes)
		},
		func(r *bufio.Reader, size int, nbytes int) (int, error) {
			if nbytes < 0 {
				return size, nil
			}
			// make sure there are enough bytes for the message value.  return
			// errShortRead if the message is truncated.
			if nbytes > size {
				return size, errShortRead
			}
			n = nbytes // return value
			if nbytes > cap(b) {
				nbytes = cap(b)
			}
			if nbytes > len(b) {
				b = b[:nbytes]
			}
			nbytes, err := io.ReadFull(r, b[:nbytes])
			if err != nil {
				return size - nbytes, err
			}
			return discardN(r, size-nbytes, n-nbytes)
		},
	)

	if err == nil && n > len(b) {
		n, err = len(b), io.ErrShortBuffer
		batch.err = io.ErrShortBuffer
		batch.offset = offset // rollback
	}

	batch.mutex.Unlock()
	return n, err
}

// ReadMessage reads and return the next message from the batch.
//
// Because this method allocate memory buffers for the message key and value
// it is less memory-efficient than Read, but has the advantage of never
// failing with io.ErrShortBuffer.
func (batch *Batch) ReadMessage() (Message, error) {
	msg := Message{}
	batch.mutex.Lock()

	var offset, timestamp int64
	var headers []Header
	var err error

	offset, timestamp, headers, err = batch.readMessage(
		func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
			msg.Key, remain, err = readNewBytes(r, size, nbytes)
			return
		},
		func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
			msg.Value, remain, err = readNewBytes(r, size, nbytes)
			return
		},
	)
	for batch.conn != nil && offset < batch.conn.offset {
		if err != nil {
			break
		}
		offset, timestamp, headers, err = batch.readMessage(
			func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
				msg.Key, remain, err = readNewBytes(r, size, nbytes)
				return
			},
			func(r *bufio.Reader, size int, nbytes int) (remain int, err error) {
				msg.Value, remain, err = readNewBytes(r, size, nbytes)
				return
			},
		)
	}

	batch.mutex.Unlock()
	msg.Topic = batch.topic
	msg.Partition = batch.partition
	msg.Offset = offset
	msg.Time = timestampToTime(timestamp)
	msg.Headers = headers

	return msg, err
}

func (batch *Batch) readMessage(
	key func(*bufio.Reader, int, int) (int, error),
	val func(*bufio.Reader, int, int) (int, error),
) (offset int64, timestamp int64, headers []Header, err error) {
	if err = batch.err; err != nil {
		return
	}

	offset, timestamp, headers, err = batch.msgs.readMessage(batch.offset, key, val)
	switch err {
	case nil:
		batch.offset = offset + 1
	case errShortRead:
		// As an "optimization" kafka truncates the returned response after
		// producing MaxBytes, which could then cause the code to return
		// errShortRead.
		err = batch.msgs.discard()
		switch {
		case err != nil:
			batch.err = err
		case batch.msgs.remaining() == 0:
			// Because we use the adjusted deadline we could end up returning
			// before the actual deadline occurred. This is necessary otherwise
			// timing out the connection for real could end up leaving it in an
			// unpredictable state, which would require closing it.
			// This design decision was made to maximize the chances of keeping
			// the connection open, the trade off being to lose precision on the
			// read deadline management.
			err = checkTimeoutErr(batch.deadline)
			batch.err = err
		}
	default:
		batch.err = err
	}

	return
}

func checkTimeoutErr(deadline time.Time) (err error) {
	if !deadline.IsZero() && time.Now().After(deadline) {
		err = RequestTimedOut
	} else {
		err = io.EOF
	}
	return
}
<|endoftext|>


<commit_before>package keywords

import (
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestMatchedUsers(t *testing.T) {
	Convey("When creating a new Keywords", t, func() {
		Convey("And it is successful", func() {
			kw := New()
			So(*kw, ShouldNotBeNil)
		})
	})

	Convey("When adding a user to a keyword", t, func() {
		Convey("And it's a new keywords", func() {
			kw := New()
			kw.Add("hello", 1)
			So(kw.kw["hello"], ShouldResemble, []int64{1})
		})
		Convey("And there is already another user", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("HELLo", 2)
			So(kw.kw["hello"], ShouldResemble, []int64{1, 2})
		})
		Convey("And the user already has that keyword", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("hELlO", 1)
			So(kw.kw["hello"], ShouldResemble, []int64{1})
			So(kw.kw["hELlO"], ShouldBeNil)
		})
	})

	Convey("When removing a user from a keyword", t, func() {
		Convey("And the keyword does not exist", func() {
			kw := New()
			kw.Remove("nothere", 1)
			So(kw.kw["nothere"], ShouldBeNil)
		})
		Convey("And the keyword exists but the user is not in it", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Remove("hello", 2)
			So(kw.kw["hello"], ShouldResemble, []int64{1})
		})
		Convey("And they are the only ones in the list", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Remove("hello", 1)
			So(kw.kw["hello"], ShouldBeNil)
		})
		Convey("And there are others in the list too", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("hello", 2)
			kw.Remove("hello", 1)
			So(kw.kw["hello"], ShouldResemble, []int64{2})
		})
	})

	Convey("Given a line of text", t, func() {
		Convey("And there is no matching user", func() {
			kw := New()
			users := kw.Find("This line does not match anything")
			So(users, ShouldBeNil)
		})
		Convey("And the text is empty", func() {
			kw := New()
			users := kw.Find("")
			So(users, ShouldBeNil)
		})
		Convey("And there is a single matching user", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("Keywords", 1)
			users := kw.Find("Hello, Keywords!")
			So(users, ShouldResemble, []int64{1})
		})
		Convey("And there are multiple matching users", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("keywords", 2)
			kw.Add("keywords", 3)
			users := kw.Find("Hello, Keywords!")
			So(users, ShouldResemble, []int64{1, 2, 3})
		})
	})

	Convey("When checking if a line matches or not", t, func() {
		Convey("And it does not match", func() {
			kw := New()
			matched := kw.Match("Nothing will match this")
			So(matched, ShouldBeFalse)
		})
		Convey("And it does match", func() {
			kw := New()
			kw.Add("match", 1)
			matched := kw.Match("This will match!")
			So(matched, ShouldBeTrue)
		})
		Convey("And the string is empty", func() {
			kw := New()
			matched := kw.Match("")
			So(matched, ShouldBeFalse)
		})
	})
}
<commit_msg>add examples for godocs<commit_after>package keywords

import (
	"fmt"
	"testing"

	. "github.com/smartystreets/goconvey/convey"
)

func TestMatchedUsers(t *testing.T) {
	Convey("When creating a new Keywords", t, func() {
		Convey("And it is successful", func() {
			kw := New()
			So(*kw, ShouldNotBeNil)
		})
	})

	Convey("When adding a user to a keyword", t, func() {
		Convey("And it's a new keywords", func() {
			kw := New()
			kw.Add("hello", 1)
			So(kw.kw["hello"], ShouldResemble, []int64{1})
		})
		Convey("And there is already another user", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("HELLo", 2)
			So(kw.kw["hello"], ShouldResemble, []int64{1, 2})
		})
		Convey("And the user already has that keyword", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("hELlO", 1)
			So(kw.kw["hello"], ShouldResemble, []int64{1})
			So(kw.kw["hELlO"], ShouldBeNil)
		})
	})

	Convey("When removing a user from a keyword", t, func() {
		Convey("And the keyword does not exist", func() {
			kw := New()
			kw.Remove("nothere", 1)
			So(kw.kw["nothere"], ShouldBeNil)
		})
		Convey("And the keyword exists but the user is not in it", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Remove("hello", 2)
			So(kw.kw["hello"], ShouldResemble, []int64{1})
		})
		Convey("And they are the only ones in the list", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Remove("hello", 1)
			So(kw.kw["hello"], ShouldBeNil)
		})
		Convey("And there are others in the list too", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("hello", 2)
			kw.Remove("hello", 1)
			So(kw.kw["hello"], ShouldResemble, []int64{2})
		})
	})

	Convey("Given a line of text", t, func() {
		Convey("And there is no matching user", func() {
			kw := New()
			users := kw.Find("This line does not match anything")
			So(users, ShouldBeNil)
		})
		Convey("And the text is empty", func() {
			kw := New()
			users := kw.Find("")
			So(users, ShouldBeNil)
		})
		Convey("And there is a single matching user", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("Keywords", 1)
			users := kw.Find("Hello, Keywords!")
			So(users, ShouldResemble, []int64{1})
		})
		Convey("And there are multiple matching users", func() {
			kw := New()
			kw.Add("hello", 1)
			kw.Add("keywords", 2)
			kw.Add("keywords", 3)
			users := kw.Find("Hello, Keywords!")
			So(users, ShouldResemble, []int64{1, 2, 3})
		})
	})

	Convey("When checking if a line matches or not", t, func() {
		Convey("And it does not match", func() {
			kw := New()
			matched := kw.Match("Nothing will match this")
			So(matched, ShouldBeFalse)
		})
		Convey("And it does match", func() {
			kw := New()
			kw.Add("match", 1)
			matched := kw.Match("This will match!")
			So(matched, ShouldBeTrue)
		})
		Convey("And the string is empty", func() {
			kw := New()
			matched := kw.Match("")
			So(matched, ShouldBeFalse)
		})
	})
}

// Adds the id 1 to the keyword _hello_
func ExampleKeywords_Add() {
	kw := New()
	kw.Add("hello", 1) // kw has now associated _hello_ with id 1
}

// Removes the id 1 from the keyword _hello_
func ExampleKeywords_Remove() {
	kw := New()
	kw.Add("hello", 1)
	kw.Remove("hello", 1) // kw is empty again
}

// Finds all ids interested in keywords in the given text
func ExampleKeywords_Find() {
	kw := New()
	kw.Add("hello", 1)
	kw.Add("world", 2)
	fmt.Println(kw.Find("Hello World"))
	// Output:
	// [1 2]
}

// Match returns true if at least one id is interested in a keyword
func ExampleKeywords_Match() {
	kw := New()
	kw.Add("hello", 1)
	fmt.Println(kw.Match("Hello World"))
	fmt.Println(kw.Match("So long"))
	// Output:
	// true
	// false
}
<|endoftext|>


<commit_before>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"errors"
	"github.com/tsuru/tsuru/queue"
	"sync"
	"sync/atomic"
	"time"
)

var (
	factory         = NewFakePubSubQFactory()
	noPubSubFactory = NewFakeQFactory()
)

func init() {
	queue.Register("fake", factory)
	queue.Register("noPubSubFake", noPubSubFactory)
}

type fakeHandler struct {
	running int32
}

func (h *fakeHandler) Start() {
	atomic.StoreInt32(&h.running, 1)
}

func (h *fakeHandler) Stop() error {
	if !atomic.CompareAndSwapInt32(&h.running, 1, 0) {
		return errors.New("Not running.")
	}
	return nil
}

func (h *fakeHandler) Wait() {}

type FakeQ struct {
	messages   messageQueue
	pubSubStop chan int
	name       string
}

type FakePubSubQ struct {
	FakeQ
	pubSubStop chan int
}

type SyncSet struct {
	set map[string]bool
	sync.Mutex
}

var subscribersSet = SyncSet{set: make(map[string]bool)}

func (s *SyncSet) put(val string) {
	s.Lock()
	defer s.Unlock()
	s.set[val] = true
}

func (s *SyncSet) get(val string) bool {
	s.Lock()
	defer s.Unlock()
	return s.set[val]
}

func (s *SyncSet) delete(val string) {
	s.Lock()
	defer s.Unlock()
	delete(s.set, val)
}

func (q *FakeQ) get(ch chan *queue.Message, stop chan int) {
	defer close(ch)
	for {
		select {
		case <-stop:
			return
		default:
		}
		if msg := q.messages.dequeue(); msg != nil {
			ch <- msg
			return
		}
		time.Sleep(1e3)
	}
}

func (q *FakePubSubQ) Pub(msg []byte) error {
	if !subscribersSet.get(q.name) {
		return nil
	}
	m := queue.Message{Action: string(msg)}
	q.messages.enqueue(&m)
	return nil
}

func (q *FakePubSubQ) Sub() (chan []byte, error) {
	subChan := make(chan []byte)
	q.pubSubStop = make(chan int, 1)
	go func() {
		defer close(subChan)
		for {
			select {
			case <-q.pubSubStop:
				return
			default:
			}
			if msg := q.messages.dequeue(); msg != nil {
				subChan <- []byte(msg.Action)
			}
			time.Sleep(1e3)
		}
	}()
	subscribersSet.put(q.name)
	return subChan, nil
}

func (q *FakePubSubQ) UnSub() error {
	subscribersSet.delete(q.name)
	close(q.pubSubStop)
	return nil
}

func (q *FakeQ) Get(timeout time.Duration) (*queue.Message, error) {
	ch := make(chan *queue.Message, 1)
	stop := make(chan int, 1)
	defer close(stop)
	go q.get(ch, stop)
	select {
	case msg := <-ch:
		return msg, nil
	case <-time.After(timeout):
	}
	return nil, errors.New("Timed out.")
}

func (q *FakeQ) Put(m *queue.Message, delay time.Duration) error {
	if delay > 0 {
		go func() {
			time.Sleep(delay)
			q.messages.enqueue(m)
		}()
	} else {
		q.messages.enqueue(m)
	}
	return nil
}

type FakePubSubQFactory struct {
	queues map[string]*FakePubSubQ
	sync.Mutex
}

func NewFakePubSubQFactory() *FakePubSubQFactory {
	return &FakePubSubQFactory{
		queues: make(map[string]*FakePubSubQ),
	}
}

func (f *FakePubSubQFactory) Get(name string) (queue.Q, error) {
	f.Lock()
	defer f.Unlock()
	if q, ok := f.queues[name]; ok {
		return q, nil
	}
	q := FakePubSubQ{FakeQ: FakeQ{name: name}}
	f.queues[name] = &q
	return &q, nil
}

func (f *FakePubSubQFactory) Handler(fn func(*queue.Message), names ...string) (queue.Handler, error) {
	return &fakeHandler{}, nil
}

type FakeQFactory struct {
	queues map[string]*FakeQ
	sync.Mutex
}

func NewFakeQFactory() *FakeQFactory {
	return &FakeQFactory{
		queues: make(map[string]*FakeQ),
	}
}

func (f *FakeQFactory) Get(name string) (queue.Q, error) {
	f.Lock()
	defer f.Unlock()
	if q, ok := f.queues[name]; ok {
		return q, nil
	}
	q := FakeQ{name: name}
	f.queues[name] = &q
	return &q, nil
}

func (f *FakeQFactory) Handler(fn func(*queue.Message), names ...string) (queue.Handler, error) {
	return &fakeHandler{}, nil
}

type messageNode struct {
	m    *queue.Message
	next *messageNode
	prev *messageNode
}

type messageQueue struct {
	first *messageNode
	last  *messageNode
	n     int
	sync.Mutex
}

func (q *messageQueue) enqueue(msg *queue.Message) {
	q.Lock()
	defer q.Unlock()
	if q.last == nil {
		q.last = &messageNode{m: msg}
		q.first = q.last
	} else {
		olast := q.last
		q.last = &messageNode{m: msg, prev: olast}
		olast.next = q.last
	}
	q.n++
}

func (q *messageQueue) dequeue() *queue.Message {
	q.Lock()
	defer q.Unlock()
	if q.n == 0 {
		return nil
	}
	msg := q.first.m
	q.n--
	q.first = q.first.next
	if q.n == 0 {
		q.last = q.first
	}
	return msg
}

// CleanQ deletes all messages from queues identified by the given names.
func CleanQ(names ...string) {
	var wg sync.WaitGroup
	for _, name := range names {
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			q, _ := factory.Get(name)
			for {
				_, err := q.Get(1e6)
				if err != nil {
					break
				}
			}
		}(name)
	}
	wg.Wait()
}
<commit_msg>testing/queue: make tests with go -race happy<commit_after>// Copyright 2014 tsuru authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"errors"
	"github.com/tsuru/tsuru/queue"
	"sync"
	"sync/atomic"
	"time"
)

var (
	factory         = NewFakePubSubQFactory()
	noPubSubFactory = NewFakeQFactory()
)

func init() {
	queue.Register("fake", factory)
	queue.Register("noPubSubFake", noPubSubFactory)
}

type fakeHandler struct {
	running int32
}

func (h *fakeHandler) Start() {
	atomic.StoreInt32(&h.running, 1)
}

func (h *fakeHandler) Stop() error {
	if !atomic.CompareAndSwapInt32(&h.running, 1, 0) {
		return errors.New("Not running.")
	}
	return nil
}

func (h *fakeHandler) Wait() {}

type FakeQ struct {
	messages messageQueue
	name     string
}

type FakePubSubQ struct {
	FakeQ
	pubSubStop     chan int
	pubSubStopLock sync.Mutex
}

type SyncSet struct {
	set map[string]bool
	sync.Mutex
}

var subscribersSet = SyncSet{set: make(map[string]bool)}

func (s *SyncSet) put(val string) {
	s.Lock()
	defer s.Unlock()
	s.set[val] = true
}

func (s *SyncSet) get(val string) bool {
	s.Lock()
	defer s.Unlock()
	return s.set[val]
}

func (s *SyncSet) delete(val string) {
	s.Lock()
	defer s.Unlock()
	delete(s.set, val)
}

func (q *FakeQ) get(ch chan *queue.Message, stop chan int) {
	defer close(ch)
	for {
		select {
		case <-stop:
			return
		default:
		}
		if msg := q.messages.dequeue(); msg != nil {
			ch <- msg
			return
		}
		time.Sleep(1e3)
	}
}

func (q *FakePubSubQ) Pub(msg []byte) error {
	if !subscribersSet.get(q.name) {
		return nil
	}
	m := queue.Message{Action: string(msg)}
	q.messages.enqueue(&m)
	return nil
}

func (q *FakePubSubQ) Sub() (chan []byte, error) {
	subChan := make(chan []byte)
	q.pubSubStopLock.Lock()
	q.pubSubStop = make(chan int)
	q.pubSubStopLock.Unlock()
	go func() {
		defer close(subChan)
		for {
			q.pubSubStopLock.Lock()
			select {
			case <-q.pubSubStop:
				q.pubSubStopLock.Unlock()
				return
			default:
			}
			q.pubSubStopLock.Unlock()
			if msg := q.messages.dequeue(); msg != nil {
				subChan <- []byte(msg.Action)
			}
			time.Sleep(1e3)
		}
	}()
	subscribersSet.put(q.name)
	return subChan, nil
}

func (q *FakePubSubQ) UnSub() error {
	subscribersSet.delete(q.name)
	close(q.pubSubStop)
	return nil
}

func (q *FakeQ) Get(timeout time.Duration) (*queue.Message, error) {
	ch := make(chan *queue.Message, 1)
	stop := make(chan int, 1)
	defer close(stop)
	go q.get(ch, stop)
	select {
	case msg := <-ch:
		return msg, nil
	case <-time.After(timeout):
	}
	return nil, errors.New("Timed out.")
}

func (q *FakeQ) Put(m *queue.Message, delay time.Duration) error {
	if delay > 0 {
		go func() {
			time.Sleep(delay)
			q.messages.enqueue(m)
		}()
	} else {
		q.messages.enqueue(m)
	}
	return nil
}

type FakePubSubQFactory struct {
	queues map[string]*FakePubSubQ
	sync.Mutex
}

func NewFakePubSubQFactory() *FakePubSubQFactory {
	return &FakePubSubQFactory{
		queues: make(map[string]*FakePubSubQ),
	}
}

func (f *FakePubSubQFactory) Get(name string) (queue.Q, error) {
	f.Lock()
	defer f.Unlock()
	if q, ok := f.queues[name]; ok {
		return q, nil
	}
	q := FakePubSubQ{FakeQ: FakeQ{name: name}}
	f.queues[name] = &q
	return &q, nil
}

func (f *FakePubSubQFactory) Handler(fn func(*queue.Message), names ...string) (queue.Handler, error) {
	return &fakeHandler{}, nil
}

type FakeQFactory struct {
	queues map[string]*FakeQ
	sync.Mutex
}

func NewFakeQFactory() *FakeQFactory {
	return &FakeQFactory{
		queues: make(map[string]*FakeQ),
	}
}

func (f *FakeQFactory) Get(name string) (queue.Q, error) {
	f.Lock()
	defer f.Unlock()
	if q, ok := f.queues[name]; ok {
		return q, nil
	}
	q := FakeQ{name: name}
	f.queues[name] = &q
	return &q, nil
}

func (f *FakeQFactory) Handler(fn func(*queue.Message), names ...string) (queue.Handler, error) {
	return &fakeHandler{}, nil
}

type messageNode struct {
	m    *queue.Message
	next *messageNode
	prev *messageNode
}

type messageQueue struct {
	first *messageNode
	last  *messageNode
	n     int
	sync.Mutex
}

func (q *messageQueue) enqueue(msg *queue.Message) {
	q.Lock()
	defer q.Unlock()
	if q.last == nil {
		q.last = &messageNode{m: msg}
		q.first = q.last
	} else {
		olast := q.last
		q.last = &messageNode{m: msg, prev: olast}
		olast.next = q.last
	}
	q.n++
}

func (q *messageQueue) dequeue() *queue.Message {
	q.Lock()
	defer q.Unlock()
	if q.n == 0 {
		return nil
	}
	msg := q.first.m
	q.n--
	q.first = q.first.next
	if q.n == 0 {
		q.last = q.first
	}
	return msg
}

// CleanQ deletes all messages from queues identified by the given names.
func CleanQ(names ...string) {
	var wg sync.WaitGroup
	for _, name := range names {
		wg.Add(1)
		go func(name string) {
			defer wg.Done()
			q, _ := factory.Get(name)
			for {
				_, err := q.Get(1e6)
				if err != nil {
					break
				}
			}
		}(name)
	}
	wg.Wait()
}
<|endoftext|>


<commit_before>package tests

import "github.com/revel/revel/testing"

type AppTest struct {
	testing.TestSuite
}

func (t *AppTest) Before() {
	println("Set up")
}

func (t *AppTest) TestThatIndexPageWorks() {
	t.Get("/")
	t.AssertOk()
	t.AssertContentType("text/html; charset=utf-8")
}

func (t *AppTest) After() {
	println("Tear down")
}
<commit_msg>add test for /objects<commit_after>package tests

import "github.com/revel/revel/testing"

type AppTest struct {
	testing.TestSuite
}

func (t *AppTest) Before() {
	println("Set up")
}

func (t *AppTest) TestThatIndexPageWorks() {
	t.Get("/")
	t.AssertOk()
	t.AssertContentType("text/html; charset=utf-8")
}

func (t *AppTest) TestThatObjectsPageWorks() {
	t.Get("/objects")
	t.AssertOk()
	t.AssertContentType("application/json; charset=utf-8")
}

func (t *AppTest) After() {
	println("Tear down")
}
<|endoftext|>


<commit_before>package immortal

import (
	"os"
	"os/signal"
)

func (self *Daemon) Block() {
	c := make(chan os.Signal)
	signal.Notify(c, os.Interrupt, os.Kill)
	<-c
	os.Exit(0)
}
<commit_msg>	deleted:    block.go<commit_after><|endoftext|>


<commit_before>package husky

import (
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"regexp"
	"strings"
)

var whiteChar = regexp.MustCompile("^[0-9a-zA-Z_\\-\\.]+$")

type Database struct {
	conn      *sql.DB
	trans     *sql.Tx
	limit     int
	offset    int
	where     []string
	fields    []string
	bind      []interface{}
	enableLog bool
	queryLog  []string
}

func NewDb(dsn string) *Database {
	d := &Database{}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		fmt.Printf("DataBase Connection Error: %s\n", err)
		return d
	}

	d.conn = db
	return d
}

func (d *Database) TransBegin() {
	var err error
	if d.trans, err = d.conn.Begin(); err != nil {
		fmt.Printf("Transaction beginning Error: %s\n", err)
	}
}

func (d *Database) TransCommit() {
	if d.trans != nil {
		d.trans.Commit()
		d.trans = nil
	}
}

func (d *Database) TransRollback() {
	if d.trans != nil {
		d.trans.Rollback()
		d.trans = nil
	}
}

func (d *Database) isTrans() bool {
	if d.trans != nil {
		return true
	}

	return false
}

func (d *Database) EnableLog(enable bool) {
	d.enableLog = enable
}

func (d *Database) Select(columns ...string) *Database {
	for _, c := range columns {
		if !whiteChar.MatchString(c) {
			panic("Invalid columns name specified.")
		}
		d.fields = append(d.fields, c)
	}

	return d
}

func (d *Database) Limit(limit int) *Database {
	d.limit = limit

	return d
}

func (d *Database) Offset(offset int) *Database {
	d.offset = offset

	return d
}

func (d *Database) Where(field, operator string, bind interface{}) *Database {
	d.where = append(d.where, field+" "+operator+" ?")
	d.bind = append(d.bind, bind)

	return d
}

func (d *Database) Get(table string) (rows *sql.Rows, err error) {
	query := d.buildSelectQuery(table)
	defer d.clear()

	d.log(query, d.bind)
	if len(d.bind) > 0 {
		if d.isTrans() {
			return d.trans.Query(query, d.bind...)
		} else {
			return d.conn.Query(query, d.bind...)
		}
	} else {
		if d.isTrans() {
			return d.trans.Query(query)
		} else {
			return d.conn.Query(query)
		}
	}
}

func (d *Database) GetRow(table string) (row *sql.Row) {
	query := d.buildSelectQuery(table)
	defer d.clear()

	d.log(query, d.bind)
	if len(d.bind) > 0 {
		if d.isTrans() {
			return d.trans.QueryRow(query, d.bind...)
		} else {
			return d.conn.QueryRow(query, d.bind...)
		}
	} else {
		if d.isTrans() {
			return d.trans.QueryRow(query)
		} else {
			return d.conn.QueryRow(query)
		}
	}
}

func (d *Database) Insert(table string, values map[string]interface{}) (result sql.Result, err error) {
	query, bind := d.buildInsertQuery(table, values)
	defer d.clear()

	d.log(query, d.bind)
	if d.isTrans() {
		return d.trans.Exec(query, bind...)
	} else {
		return d.conn.Exec(query, bind...)
	}
}

func (d *Database) Update(table string, values map[string]interface{}) (result sql.Result, err error) {
	query, bind := d.buildUpdateQuery(table, values)
	defer d.clear()

	d.log(query, d.bind)
	if d.isTrans() {
		return d.trans.Exec(query, bind...)
	} else {
		return d.conn.Exec(query, bind...)
	}
}

func (d *Database) buildSelectQuery(table string) (query string) {
	if !whiteChar.MatchString(table) {
		panic("Invalid table name specified.")
	}

	query = "SELECT "
	if len(d.fields) == 0 {
		query += "*"
	} else {
		query += strings.Join(d.fields, ", ")
	}
	query += " FROM " + table
	if len(d.where) > 0 {
		query += " WHERE " + strings.Join(d.where, " AND ")
	}

	if d.limit > 0 {
		query += " LIMIT " + fmt.Sprint(d.limit)
	}
	if d.offset > 0 {
		query += " OFFSET " + fmt.Sprint(d.offset)
	}
	return
}

func (d *Database) buildInsertQuery(table string, values map[string]interface{}) (query string, bind []interface{}) {
	var (
		fields    []string
		statement []string
	)

	for f, val := range values {
		fields = append(fields, f)
		statement = append(statement, "?")
		bind = append(bind, val)
	}

	query = fmt.Sprintf(
		"INSERT INTO %s (%s) VALUES (%s)",
		table,
		strings.Join(fields, ", "),
		strings.Join(statement, ", "),
	)
	return
}

func (d *Database) buildUpdateQuery(table string, values map[string]interface{}) (query string, bind []interface{}) {
	var fields []string

	for f, val := range values {
		fields = append(fields, f+" = ?")
		bind = append(bind, val)
	}

	query = "UPDATE " + table + " SET  (" + strings.Join(fields, ", ") + ")"
	if len(d.where) > 0 {
		query += " WHERE " + strings.Join(d.where, " AND ")
	}

	if d.limit > 0 {
		query += " LIMIT " + fmt.Sprint(d.limit)
	}

	return
}

func (d *Database) log(query string, params []interface{}) {
	if d.enableLog {
		log := fmt.Sprintf("%s, %v", query, params)
		d.queryLog = append(d.queryLog, log)
	}
}

func (d *Database) LastQuery() string {
	index := len(d.queryLog) - 1
	return d.queryLog[index]
}

func (d *Database) AllQuery() string {
	return strings.Join(d.queryLog, "\n")
}

func (d *Database) clear() {
	d.limit = 0
	d.offset = 0
	d.where = []string{}
	d.fields = []string{}
	d.bind = []interface{}{}
}
<commit_msg>Implement simple query syntax<commit_after>package husky

import (
	"database/sql"
	"fmt"
	_ "github.com/go-sql-driver/mysql"
	"regexp"
	"strings"
)

var whiteChar = regexp.MustCompile("^[0-9a-zA-Z_\\-\\.]+$")

type Database struct {
	conn      *sql.DB
	trans     *sql.Tx
	limit     int
	offset    int
	where     []string
	fields    []string
	bind      []interface{}
	enableLog bool
	queryLog  []string
}

func NewDb(dsn string) *Database {
	d := &Database{}
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		fmt.Printf("DataBase Connection Error: %s\n", err)
		return d
	}

	d.conn = db
	return d
}

func (d *Database) TransBegin() {
	var err error
	if d.trans, err = d.conn.Begin(); err != nil {
		fmt.Printf("Transaction beginning Error: %s\n", err)
	}
}

func (d *Database) TransCommit() {
	if d.trans != nil {
		d.trans.Commit()
		d.trans = nil
	}
}

func (d *Database) TransRollback() {
	if d.trans != nil {
		d.trans.Rollback()
		d.trans = nil
	}
}

func (d *Database) IsTrans() bool {
	if d.trans != nil {
		return true
	}

	return false
}

func (d *Database) EnableLog(enable bool) {
	d.enableLog = enable
}

func (d *Database) Query(query string, bind ...interface{}) (rows *sql.Rows, err error) {
	if d.IsTrans() {
		return d.trans.Query(query, bind...)
	} else {
		return d.conn.Query(query, bind...)
	}
}

func (d *Database) QueryRow(query string, bind ...interface{}) (rows *sql.Row) {
	if d.IsTrans() {
		return d.trans.QueryRow(query, bind...)
	} else {
		return d.conn.QueryRow(query, bind...)
	}
}

func (d *Database) Limit(limit int) *Database {
	d.limit = limit

	return d
}

func (d *Database) Offset(offset int) *Database {
	d.offset = offset

	return d
}

func (d *Database) Where(field, operator string, bind interface{}) *Database {
	d.where = append(d.where, field+" "+operator+" ?")
	d.bind = append(d.bind, bind)

	return d
}

func (d *Database) Get(table string) (rows *sql.Rows, err error) {
	query := d.buildSelectQuery(table)
	defer d.clear()

	d.log(query, d.bind)
	if len(d.bind) > 0 {
		if d.IsTrans() {
			return d.trans.Query(query, d.bind...)
		} else {
			return d.conn.Query(query, d.bind...)
		}
	} else {
		if d.IsTrans() {
			return d.trans.Query(query)
		} else {
			return d.conn.Query(query)
		}
	}
}

func (d *Database) GetRow(table string) (row *sql.Row) {
	query := d.buildSelectQuery(table)
	defer d.clear()

	d.log(query, d.bind)
	if len(d.bind) > 0 {
		if d.IsTrans() {
			return d.trans.QueryRow(query, d.bind...)
		} else {
			return d.conn.QueryRow(query, d.bind...)
		}
	} else {
		if d.IsTrans() {
			return d.trans.QueryRow(query)
		} else {
			return d.conn.QueryRow(query)
		}
	}
}

func (d *Database) Insert(table string, values map[string]interface{}) (result sql.Result, err error) {
	query, bind := d.buildInsertQuery(table, values)
	defer d.clear()

	d.log(query, d.bind)
	if d.IsTrans() {
		return d.trans.Exec(query, bind...)
	} else {
		return d.conn.Exec(query, bind...)
	}
}

func (d *Database) Update(table string, values map[string]interface{}) (result sql.Result, err error) {
	query, bind := d.buildUpdateQuery(table, values)
	defer d.clear()

	d.log(query, d.bind)
	if d.IsTrans() {
		return d.trans.Exec(query, bind...)
	} else {
		return d.conn.Exec(query, bind...)
	}
}

func (d *Database) buildSelectQuery(table string) (query string) {
	if !whiteChar.MatchString(table) {
		panic("Invalid table name specified.")
	}

	query = "SELECT "
	if len(d.fields) == 0 {
		query += "*"
	} else {
		query += strings.Join(d.fields, ", ")
	}
	query += " FROM " + table
	if len(d.where) > 0 {
		query += " WHERE " + strings.Join(d.where, " AND ")
	}

	if d.limit > 0 {
		query += " LIMIT " + fmt.Sprint(d.limit)
	}
	if d.offset > 0 {
		query += " OFFSET " + fmt.Sprint(d.offset)
	}
	return
}

func (d *Database) buildInsertQuery(table string, values map[string]interface{}) (query string, bind []interface{}) {
	var (
		fields    []string
		statement []string
	)

	for f, val := range values {
		fields = append(fields, f)
		statement = append(statement, "?")
		bind = append(bind, val)
	}

	query = fmt.Sprintf(
		"INSERT INTO %s (%s) VALUES (%s)",
		table,
		strings.Join(fields, ", "),
		strings.Join(statement, ", "),
	)
	return
}

func (d *Database) buildUpdateQuery(table string, values map[string]interface{}) (query string, bind []interface{}) {
	var fields []string

	for f, val := range values {
		fields = append(fields, f+" = ?")
		bind = append(bind, val)
	}

	query = "UPDATE " + table + " SET  (" + strings.Join(fields, ", ") + ")"
	if len(d.where) > 0 {
		query += " WHERE " + strings.Join(d.where, " AND ")
	}

	if d.limit > 0 {
		query += " LIMIT " + fmt.Sprint(d.limit)
	}

	return
}

func (d *Database) log(query string, params []interface{}) {
	if d.enableLog {
		log := fmt.Sprintf("%s, %v", query, params)
		d.queryLog = append(d.queryLog, log)
	}
}

func (d *Database) LastQuery() string {
	index := len(d.queryLog) - 1
	return d.queryLog[index]
}

func (d *Database) AllQuery() string {
	return strings.Join(d.queryLog, "\n")
}

func (d *Database) clear() {
	d.limit = 0
	d.offset = 0
	d.where = []string{}
	d.fields = []string{}
	d.bind = []interface{}{}
}
<|endoftext|>


<commit_before>package main

import (
	log "code.google.com/p/log4go"
	"database/sql"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"time"
)

const NODES_SCHEMA = `
CREATE TABLE IF NOT EXISTS nodes (
    id INTEGER NOT NULL PRIMARY KEY,
    kind INTEGER NOT NULL,
    updated_at INTEGER,
    name TEXT,
    domoticz_idx TEXT,
    temperature REAL,
    humidity INTEGER,
    light INTEGER,
    motion INTEGER,
    lowbat INTEGER,
    vcc INTEGER
);
`

const LOGS_SCHEMA = `
CREATE TABLE IF NOT EXISTS logs (
	node_id INTEGER NOT NULL,
	at INTEGER NOT NULL,
	temperature REAL,
	humidity INTEGER,
	light INTEGER,
	motion INTEGER,
	lowbat INTEGER,
	vcc INTEGER
);
`

var ColNameForSensor map[Sensor]string

// Database
type Database struct {
	driver     *sql.DB
	nodes      []*Node
	logsTicker *time.Ticker
}

// Init
func init() {
	ColNameForSensor = map[Sensor]string{
		TEMP_SENSOR:   "temperature",
		HUMI_SENSOR:   "humidity",
		LIGHT_SENSOR:  "light",
		MOTION_SENSOR: "motion",
		LOWBAT_SENSOR: "lowbat",
		VCC_SENSOR:    "vcc",
	}
}

// Setup a new database connection and load nodes
func loadDatabase(databasePath string) (*Database, error) {
	// open
	sqlDriver, err := sql.Open("sqlite3", databasePath)
	if err != nil {
		return nil, err
	}

	db := Database{driver: sqlDriver}

	// create tables if necessary
	db.createTables()

	// load nodes
	db.loadNodes()

	return &db, nil
}

func (db *Database) close() {
	db.driver.Close()
}

// Create tables
func (db *Database) createTables() {
	schemas := [2]string{NODES_SCHEMA, LOGS_SCHEMA}

	for _, schema := range schemas {
		_, err := db.driver.Exec(schema)
		if err != nil {
			panic(log.Critical("Failed to create SQL table %q: %s", err, schema))
		}
	}
}

// Load all nodes
func (db *Database) loadNodes() {
	// reset nodes
	db.nodes = make([]*Node, 0)

	// fetch nodes from db
	rows, err := db.driver.Query("SELECT * FROM nodes")
	if err != nil {
		panic(log.Critical(err))
	}
	defer rows.Close()

	for rows.Next() {
		var node *Node

		// fetch node fields
		var (
			id           int
			kind         int
			updated_at   int64
			name         sql.NullString
			domoticz_idx sql.NullString
			temperature  sql.NullFloat64
			humidity     sql.NullInt64
			light        sql.NullInt64
			motion       sql.NullBool
			lowbat       sql.NullBool
			vcc          sql.NullInt64
		)

		// @todo Use github.com/russross/meddler ?
		rows.Scan(&id, &kind, &updated_at, &name, &domoticz_idx, &temperature, &humidity, &light, &motion, &lowbat, &vcc)

		// init node
		node = &Node{
			Id:        id,
			Kind:      kind,
			UpdatedAt: time.Unix(updated_at, 0),
		}

		if name.Valid {
			node.Name = name.String
		}

		if domoticz_idx.Valid {
			node.DomoticzIdx = domoticz_idx.String
		}

		if temperature.Valid {
			node.Temperature = float64(temperature.Float64)
		}

		if humidity.Valid {
			node.Humidity = uint8(humidity.Int64)
		}

		if light.Valid {
			node.Light = uint8(light.Int64)
		}

		if motion.Valid {
			node.Motion = motion.Bool
		}

		if lowbat.Valid {
			node.LowBattery = lowbat.Bool
		}

		if vcc.Valid {
			node.Vcc = uint(vcc.Int64)
		}

		// add node to list
		db.nodes = append(db.nodes, node)
	}

	rows.Close()
}

// Get a node
func (db *Database) nodeForId(id int) *Node {
	for _, node := range db.nodes {
		if node.Id == id {
			return node
		}
	}

	return nil
}

func insertNodeQuery(id int, kind int) (string, []interface{}) {
	return "INSERT INTO nodes(id, kind) VALUES(?, ?)", []interface{}{ id, kind }
}

// Insert a new node
func (db *Database) insertNode(id int, kind int) *Node {
	// init node
	node := &Node{Id: id, Kind: kind}

	// add node to list
	db.nodes = append(db.nodes, node)

	// persist in database
	query, args := insertNodeQuery(id, kind)

	_, err := db.driver.Exec(query, args...)
	if err != nil {
		panic(log.Critical(err))
	}

	return node
}

func updateNodeQuery(node *Node) (string, []interface{}) {
	args := make([]interface{}, 0)

	query := "UPDATE nodes SET updated_at = ?"
	args = append(args, node.UpdatedAt.Unix())

	// set sensors values
	for _, sensor := range node.sensors() {
		colName := ColNameForSensor[sensor]
		if colName != "" {
			value := node.sensorValue(sensor)

			query += fmt.Sprintf(", %s = ?", colName)
			args = append(args, value)
		}
	}

	// set NULL for absent sensors
	for _, sensor := range node.absentSensors() {
		colName := ColNameForSensor[sensor]
		if colName != "" {
			query += fmt.Sprintf(", %s = NULL", colName)
		}
	}

	query += " WHERE id = ?"
	args = append(args, node.Id)

	return query, args
}

// Update node
func (db *Database) updateNode(node *Node) {
	if len(node.sensors()) > 0 {
		node.UpdatedAt = time.Now().UTC()

		// persist in database
		query, args := updateNodeQuery(node)

		_, err := db.driver.Exec(query, args...)
		if err != nil {
			panic(log.Critical(err))
		}
	}
}

func insertLogQuery(node *Node) (string, []interface{}) {
	args := make([]interface{}, 0)

	query := "INSERT INTO logs(node_id, at"
	args = append(args, node.Id)
	args = append(args, time.Now().UTC().Unix())

	nbSensors := 0

	for _, sensor := range node.sensors() {
		colName := ColNameForSensor[sensor]
		if colName != "" {
			query += fmt.Sprintf(", %s", colName)
			args = append(args, node.sensorValue(sensor))

			nbSensors += 1
		}
	}

	query += ") VALUES(?, ?"
	for i := 0; i < nbSensors; i++ {
		query += ", ?"
	}
	query += ")"

	return query, args
}

// Insert log for given node
func (db *Database) insertLog(node *Node) {
	if len(node.sensors()) > 0 {
		// persist in database
		query, args := insertLogQuery(node)

		_, err := db.driver.Exec(query, args...)
		if err != nil {
			panic(log.Critical(err))
		}
	}
}

// Insert logs for all nodes
func (db *Database) insertLogs() {
	for _, node := range db.nodes {
		db.insertLog(node)
	}
}

// Add a log entry every 5 minutes
func (db *Database) startLogsTicker(period time.Duration, history time.Duration) {
	db.logsTicker = time.NewTicker(period)

	// do it right now
	db.insertLogs()

	go func() {
		for _ = range db.logsTicker.C {
			db.insertLogs()

			db.trimLogs(history)
		}
	}()
}

func trimLogsQuery(history time.Duration) (string, []interface{}) {
	return "DELETE FROM logs WHERE (at < ?)", []interface{}{ time.Now().UTC().Add(-history).Unix() }
}

// Delete old logs
func (db *Database) trimLogs(history time.Duration) {
	// persist in database
	query, args := trimLogsQuery(history)

	_, err := db.driver.Exec(query, args...)
	if err != nil {
		panic(log.Critical(err))
	}
}
<commit_msg>fixes the infamous Database Locked sqlite issue<commit_after>package main

import (
	log "code.google.com/p/log4go"
	"database/sql"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"time"
)

const NODES_SCHEMA = `
CREATE TABLE IF NOT EXISTS nodes (
    id INTEGER NOT NULL PRIMARY KEY,
    kind INTEGER NOT NULL,
    updated_at INTEGER,
    name TEXT,
    domoticz_idx TEXT,
    temperature REAL,
    humidity INTEGER,
    light INTEGER,
    motion INTEGER,
    lowbat INTEGER,
    vcc INTEGER
);
`

const LOGS_SCHEMA = `
CREATE TABLE IF NOT EXISTS logs (
	node_id INTEGER NOT NULL,
	at INTEGER NOT NULL,
	temperature REAL,
	humidity INTEGER,
	light INTEGER,
	motion INTEGER,
	lowbat INTEGER,
	vcc INTEGER
);
`

var ColNameForSensor map[Sensor]string

// Database
type Database struct {
	driver      *sql.DB
	queryWriter chan *DatabaseQuery
	nodes       []*Node
	logsTicker  *time.Ticker
}

// Database Query
type DatabaseQuery struct {
	query string
	args  []interface{}
}

// Init
func init() {
	ColNameForSensor = map[Sensor]string{
		TEMP_SENSOR:   "temperature",
		HUMI_SENSOR:   "humidity",
		LIGHT_SENSOR:  "light",
		MOTION_SENSOR: "motion",
		LOWBAT_SENSOR: "lowbat",
		VCC_SENSOR:    "vcc",
	}
}

// Setup a new database connection and load nodes
func loadDatabase(databasePath string) (*Database, error) {
	// open
	sqlDriver, err := sql.Open("sqlite3", databasePath)
	if err != nil {
		return nil, err
	}

	db := Database{driver: sqlDriver}

	// create tables if necessary
	db.createTables()

	// load nodes
	db.loadNodes()

	// run query writer
	db.runQueryWriter()

	return &db, nil
}

// Start RF12demo handler
func (db *Database) runQueryWriter() {
	inputChan := make(chan *DatabaseQuery)

	go func() {
		var dbQuery *DatabaseQuery

		// loop forever
		for {
			dbQuery = <-inputChan

			// log.Debug("Exec DB write query: %s", dbQuery.query)

			_, err := db.driver.Exec(dbQuery.query, dbQuery.args...)
			if err != nil {
				panic(log.Critical(err))
			}
		}
	}()

	db.queryWriter = inputChan
}

func (db *Database) writeQuery(dbQuery *DatabaseQuery) {
	db.queryWriter <- dbQuery
}

func (db *Database) close() {
	db.driver.Close()
}

// Create tables
func (db *Database) createTables() {
	schemas := [2]string{NODES_SCHEMA, LOGS_SCHEMA}

	for _, schema := range schemas {
		_, err := db.driver.Exec(schema)
		if err != nil {
			panic(log.Critical("Failed to create SQL table %q: %s", err, schema))
		}
	}
}

// Load all nodes
func (db *Database) loadNodes() {
	// reset nodes
	db.nodes = make([]*Node, 0)

	// fetch nodes from db
	rows, err := db.driver.Query("SELECT * FROM nodes")
	if err != nil {
		panic(log.Critical(err))
	}
	defer rows.Close()

	for rows.Next() {
		var node *Node

		// fetch node fields
		var (
			id           int
			kind         int
			updated_at   int64
			name         sql.NullString
			domoticz_idx sql.NullString
			temperature  sql.NullFloat64
			humidity     sql.NullInt64
			light        sql.NullInt64
			motion       sql.NullBool
			lowbat       sql.NullBool
			vcc          sql.NullInt64
		)

		// @todo Use github.com/russross/meddler ?
		rows.Scan(&id, &kind, &updated_at, &name, &domoticz_idx, &temperature, &humidity, &light, &motion, &lowbat, &vcc)

		// init node
		node = &Node{
			Id:        id,
			Kind:      kind,
			UpdatedAt: time.Unix(updated_at, 0),
		}

		if name.Valid {
			node.Name = name.String
		}

		if domoticz_idx.Valid {
			node.DomoticzIdx = domoticz_idx.String
		}

		if temperature.Valid {
			node.Temperature = float64(temperature.Float64)
		}

		if humidity.Valid {
			node.Humidity = uint8(humidity.Int64)
		}

		if light.Valid {
			node.Light = uint8(light.Int64)
		}

		if motion.Valid {
			node.Motion = motion.Bool
		}

		if lowbat.Valid {
			node.LowBattery = lowbat.Bool
		}

		if vcc.Valid {
			node.Vcc = uint(vcc.Int64)
		}

		// add node to list
		db.nodes = append(db.nodes, node)
	}

	rows.Close()
}

// Get a node
func (db *Database) nodeForId(id int) *Node {
	for _, node := range db.nodes {
		if node.Id == id {
			return node
		}
	}

	return nil
}

// Insert a new node
func (db *Database) insertNode(id int, kind int) *Node {
	// init node
	node := &Node{Id: id, Kind: kind}

	// add node to list
	db.nodes = append(db.nodes, node)

	// persist in database
	db.writeQuery(&DatabaseQuery{
		query: "INSERT INTO nodes(id, kind) VALUES(?, ?)",
		args:  []interface{}{ id, kind },
	})

	return node
}

func updateNodeQuery(node *Node) *DatabaseQuery {
	args := make([]interface{}, 0)

	query := "UPDATE nodes SET updated_at = ?"
	args = append(args, node.UpdatedAt.Unix())

	// set sensors values
	for _, sensor := range node.sensors() {
		colName := ColNameForSensor[sensor]
		if colName != "" {
			value := node.sensorValue(sensor)

			query += fmt.Sprintf(", %s = ?", colName)
			args = append(args, value)
		}
	}

	// set NULL for absent sensors
	for _, sensor := range node.absentSensors() {
		colName := ColNameForSensor[sensor]
		if colName != "" {
			query += fmt.Sprintf(", %s = NULL", colName)
		}
	}

	query += " WHERE id = ?"
	args = append(args, node.Id)

	return &DatabaseQuery{ query: query, args: args }
}

// Update node
func (db *Database) updateNode(node *Node) {
	if len(node.sensors()) > 0 {
		node.UpdatedAt = time.Now().UTC()

		// persist in database
		db.writeQuery(updateNodeQuery(node))
	}
}

func insertLogQuery(node *Node) *DatabaseQuery {
	args := make([]interface{}, 0)

	query := "INSERT INTO logs(node_id, at"
	args = append(args, node.Id)
	args = append(args, time.Now().UTC().Unix())

	nbSensors := 0

	for _, sensor := range node.sensors() {
		colName := ColNameForSensor[sensor]
		if colName != "" {
			query += fmt.Sprintf(", %s", colName)
			args = append(args, node.sensorValue(sensor))

			nbSensors += 1
		}
	}

	query += ") VALUES(?, ?"
	for i := 0; i < nbSensors; i++ {
		query += ", ?"
	}
	query += ")"

	return &DatabaseQuery{ query: query, args: args }
}

// Insert log for given node
func (db *Database) insertLog(node *Node) {
	if len(node.sensors()) > 0 {
		// persist in database
		db.writeQuery(insertLogQuery(node))
	}
}

// Insert logs for all nodes
func (db *Database) insertLogs() {
	for _, node := range db.nodes {
		db.insertLog(node)
	}
}

// Add a log entry every 5 minutes
func (db *Database) startLogsTicker(period time.Duration, history time.Duration) {
	db.logsTicker = time.NewTicker(period)

	// do it right now
	db.insertLogs()

	go func() {
		for _ = range db.logsTicker.C {
			db.insertLogs()

			db.trimLogs(history)
		}
	}()
}

// Delete old logs
func (db *Database) trimLogs(history time.Duration) {
	// persist in database
	db.writeQuery(&DatabaseQuery{
		query: "DELETE FROM logs WHERE (at < ?)",
		args:  []interface{}{ time.Now().UTC().Add(-history).Unix() },
	})
}
<|endoftext|>


<commit_before>// Copyright 2014 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

// File database.go contains code strictly related to the database, including
// setting up the database with given config, generating unique,
// random ids, and creating and managing a connection pool. There
// are also convenience functions for (e.g.) checking if a key exists
// in redis.

package zoom

import (
	"github.com/dchest/uniuri"
	"github.com/garyburd/redigo/redis"
	"strconv"
	"time"
)

// Configuration contains various options. It should be created once
// and passed in to the Init function during application startup.
type Configuration struct {
	Address  string // Address to connect to. Default: "localhost:6379"
	Network  string // Network to use. Default: "tcp"
	Database int    // Database id to use (using SELECT). Default: 0
}

var pool *redis.Pool

var defaultConfiguration = Configuration{
	Address:  "localhost:6379",
	Network:  "tcp",
	Database: 0,
}

// GetConn gets a connection from the connection pool and returns it.
// It can be used for directly interacting with the database. Check out
// http://godoc.org/github.com/garyburd/redigo/redis for full documentation
// on the redis.Conn type.
func GetConn() redis.Conn {
	return pool.Get()
}

// Init starts the Zoom library and creates a connection pool. It accepts
// a Configuration struct as an argument. Any zero values in the configuration
// will fallback to their default values. Init should be called once during
// application startup.
func Init(passedConfig *Configuration) {
	config := getConfiguration(passedConfig)
	pool = &redis.Pool{
		MaxIdle:     10,
		MaxActive:   0,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {
			c, err := redis.Dial(config.Network, config.Address)
			if err != nil {
				return nil, err
			}
			if _, err := c.Do("select", strconv.Itoa(config.Database)); err != nil {
				c.Close()
				return nil, err
			}
			return c, err
		},
	}
}

// Close closes the connection pool and shuts down the Zoom library.
// It should be run when application exits, e.g. using defer.
func Close() {
	pool.Close()
}

// KeyExists returns true iff a given key exists in redis.
// If conn is nil, a new connection will be created and
// closed before the end of the function.
func KeyExists(key string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("exists", key))
}

// SetContains returns true iff the redis set identified by key contains
// member.  If conn is nil, a new connection will be created and
// closed before the end of the function.
func SetContains(key, member string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("sismember", key, member))
}

// generateRandomId generates a random string that is more or less
// garunteed to be unique. Used as Ids for records where an Id is
// not otherwise provided.
func generateRandomId() string {
	timeInt := time.Now().Unix()
	timeString := strconv.FormatInt(timeInt, 36)
	randomString := uniuri.NewLen(16)
	return randomString + timeString
}

// getConfiguration returns a well-formed configuration struct.
// If the passedConfig is nil, returns defaultConfiguration.
// Else, for each zero value field in passedConfig,
// use the default value for that field.
func getConfiguration(passedConfig *Configuration) Configuration {
	if passedConfig == nil {
		return defaultConfiguration
	}

	// copy the passedConfig
	newConfig := *passedConfig

	if newConfig.Address == "" {
		newConfig.Address = defaultConfiguration.Address
	}
	if newConfig.Network == "" {
		newConfig.Network = defaultConfiguration.Network
	}
	// since the zero value for int is 0, we can skip config.Database

	return newConfig
}
<commit_msg>accepts config.Address like redis://user:pass@localhost:123<commit_after>// Copyright 2014 Alex Browne.  All rights reserved.
// Use of this source code is governed by the MIT
// license, which can be found in the LICENSE file.

// File database.go contains code strictly related to the database, including
// setting up the database with given config, generating unique,
// random ids, and creating and managing a connection pool. There
// are also convenience functions for (e.g.) checking if a key exists
// in redis.

package zoom

import (
	"net/url"
	"strconv"
	"time"

	"github.com/dchest/uniuri"
	"github.com/garyburd/redigo/redis"
)

// Configuration contains various options. It should be created once
// and passed in to the Init function during application startup.
type Configuration struct {
	Address  string // Address to connect to. Default: "localhost:6379"
	Network  string // Network to use. Default: "tcp"
	Database int    // Database id to use (using SELECT). Default: 0
}

var pool *redis.Pool

var defaultConfiguration = Configuration{
	Address:  "localhost:6379",
	Network:  "tcp",
	Database: 0,
}

// GetConn gets a connection from the connection pool and returns it.
// It can be used for directly interacting with the database. Check out
// http://godoc.org/github.com/garyburd/redigo/redis for full documentation
// on the redis.Conn type.
func GetConn() redis.Conn {
	return pool.Get()
}

// Init starts the Zoom library and creates a connection pool. It accepts
// a Configuration struct as an argument. Any zero values in the configuration
// will fallback to their default values. Init should be called once during
// application startup.
func Init(passedConfig *Configuration) {
	config := getConfiguration(passedConfig)
	pool = &redis.Pool{
		MaxIdle:     10,
		MaxActive:   0,
		IdleTimeout: 240 * time.Second,
		Dial: func() (redis.Conn, error) {

			u, err := url.Parse(config.Address)
			if err != nil {
				return nil, err
			}

			address := config.Address

			if u.Host != "" {
				address = u.Host
			}

			c, err := redis.Dial(config.Network, address)
			if err != nil {
				return nil, err
			}

			if u.User != nil {
				pw, ok := u.User.Password()
				if !ok {
					return nil, err
				}
				_, err = c.Do("AUTH", pw)
				if err != nil {
					return nil, err
				}

			}

			if _, err := c.Do("select", strconv.Itoa(config.Database)); err != nil {
				c.Close()
				return nil, err
			}
			return c, err
		},
	}
}

// Close closes the connection pool and shuts down the Zoom library.
// It should be run when application exits, e.g. using defer.
func Close() {
	pool.Close()
}

// KeyExists returns true iff a given key exists in redis.
// If conn is nil, a new connection will be created and
// closed before the end of the function.
func KeyExists(key string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("exists", key))
}

// SetContains returns true iff the redis set identified by key contains
// member.  If conn is nil, a new connection will be created and
// closed before the end of the function.
func SetContains(key, member string, conn redis.Conn) (bool, error) {
	if conn == nil {
		conn = pool.Get()
		defer conn.Close()
	}
	return redis.Bool(conn.Do("sismember", key, member))
}

// generateRandomId generates a random string that is more or less
// garunteed to be unique. Used as Ids for records where an Id is
// not otherwise provided.
func generateRandomId() string {
	timeInt := time.Now().Unix()
	timeString := strconv.FormatInt(timeInt, 36)
	randomString := uniuri.NewLen(16)
	return randomString + timeString
}

// getConfiguration returns a well-formed configuration struct.
// If the passedConfig is nil, returns defaultConfiguration.
// Else, for each zero value field in passedConfig,
// use the default value for that field.
func getConfiguration(passedConfig *Configuration) Configuration {
	if passedConfig == nil {
		return defaultConfiguration
	}

	// copy the passedConfig
	newConfig := *passedConfig

	if newConfig.Address == "" {
		newConfig.Address = defaultConfiguration.Address
	}
	if newConfig.Network == "" {
		newConfig.Network = defaultConfiguration.Network
	}
	// since the zero value for int is 0, we can skip config.Database

	return newConfig
}
<|endoftext|>


<commit_before>// multpart upload for box

package box

import (
	"bytes"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/ncw/rclone/backend/box/api"
	"github.com/ncw/rclone/fs"
	"github.com/ncw/rclone/lib/rest"
	"github.com/pkg/errors"
)

// createUploadSession creates an upload session for the object
func (o *Object) createUploadSession(leaf, directoryID string, size int64) (response *api.UploadSessionResponse, err error) {
	opts := rest.Opts{
		Method:  "POST",
		Path:    "/files/upload_sessions",
		RootURL: uploadURL,
	}
	request := api.UploadSessionRequest{
		FileSize: size,
	}
	// If object has an ID then it is existing so create a new version
	if o.id != "" {
		opts.Path = "/files/" + o.id + "/upload_sessions"
	} else {
		opts.Path = "/files/upload_sessions"
		request.FolderID = directoryID
		request.FileName = replaceReservedChars(leaf)
	}
	var resp *http.Response
	err = o.fs.pacer.Call(func() (bool, error) {
		resp, err = o.fs.srv.CallJSON(&opts, &request, &response)
		return shouldRetry(resp, err)
	})
	return
}

// sha1Digest produces a digest using sha1 as per RFC3230
func sha1Digest(digest []byte) string {
	return "sha=" + base64.StdEncoding.EncodeToString(digest)
}

// uploadPart uploads a part in an upload session
func (o *Object) uploadPart(SessionID string, offset, totalSize int64, chunk []byte) (response *api.UploadPartResponse, err error) {
	chunkSize := int64(len(chunk))
	in := bytes.NewReader(chunk)
	sha1sum := sha1.Sum(chunk)
	opts := rest.Opts{
		Method:        "PUT",
		Path:          "/files/upload_sessions/" + SessionID,
		RootURL:       uploadURL,
		ContentType:   "application/octet-stream",
		ContentLength: &chunkSize,
		ContentRange:  fmt.Sprintf("bytes %d-%d/%d", offset, offset+chunkSize-1, totalSize),
		ExtraHeaders: map[string]string{
			"Digest": sha1Digest(sha1sum[:]),
		},
		Body: in,
	}
	var resp *http.Response
	err = o.fs.pacer.Call(func() (bool, error) {
		_, _ = in.Seek(0, 0)
		resp, err = o.fs.srv.CallJSON(&opts, nil, &response)
		return shouldRetry(resp, err)
	})
	if err != nil {
		return nil, err
	}
	return response, nil
}

// commitUpload finishes an upload session
func (o *Object) commitUpload(SessionID string, parts []api.Part, modTime time.Time, sha1sum []byte) (result *api.FolderItems, err error) {
	opts := rest.Opts{
		Method:  "POST",
		Path:    "/files/upload_sessions/" + SessionID + "/commit",
		RootURL: uploadURL,
		ExtraHeaders: map[string]string{
			"Digest": sha1Digest(sha1sum),
		},
	}
	request := api.CommitUpload{
		Parts: parts,
	}
	request.Attributes.ContentModifiedAt = api.Time(modTime)
	request.Attributes.ContentCreatedAt = api.Time(modTime)
	var body []byte
	var resp *http.Response
	maxTries := fs.Config.LowLevelRetries
	const defaultDelay = 10
	var tries int
outer:
	for tries = 0; tries < maxTries; tries++ {
		err = o.fs.pacer.Call(func() (bool, error) {
			resp, err = o.fs.srv.CallJSON(&opts, &request, nil)
			if err != nil {
				return shouldRetry(resp, err)
			}
			body, err = rest.ReadBody(resp)
			return shouldRetry(resp, err)
		})
		delay := defaultDelay
		why := "unknown"
		if err != nil {
			// Sometimes we get 400 Error with
			// parts_mismatch immediately after uploading
			// the last part.  Ignore this error and wait.
			if boxErr, ok := err.(*api.Error); ok && boxErr.Code == "parts_mismatch" {
				why = err.Error()
			} else {
				return nil, err
			}
		} else {
			switch resp.StatusCode {
			case http.StatusOK, http.StatusCreated:
				break outer
			case http.StatusAccepted:
				why = "not ready yet"
				delayString := resp.Header.Get("Retry-After")
				if delayString != "" {
					delay, err = strconv.Atoi(delayString)
					if err != nil {
						fs.Debugf(o, "Couldn't decode Retry-After header %q: %v", delayString, err)
						delay = defaultDelay
					}
				}
			default:
				return nil, errors.Errorf("unknown HTTP status return %q (%d)", resp.Status, resp.StatusCode)
			}
		}
		fs.Debugf(o, "commit multipart upload failed %d/%d - trying again in %d seconds (%s)", tries+1, maxTries, delay, why)
		time.Sleep(time.Duration(delay) * time.Second)
	}
	if tries >= maxTries {
		return nil, errors.New("too many tries to commit multipart upload - increase --low-level-retries")
	}
	err = json.Unmarshal(body, &result)
	if err != nil {
		return nil, errors.Wrapf(err, "couldn't decode commit response: %q", body)
	}
	return result, nil
}

// abortUpload cancels an upload session
func (o *Object) abortUpload(SessionID string) (err error) {
	opts := rest.Opts{
		Method:     "DELETE",
		Path:       "/files/upload_sessions/" + SessionID,
		RootURL:    uploadURL,
		NoResponse: true,
	}
	var resp *http.Response
	err = o.fs.pacer.Call(func() (bool, error) {
		resp, err = o.fs.srv.Call(&opts)
		return shouldRetry(resp, err)
	})
	return err
}

// uploadMultipart uploads a file using multipart upload
func (o *Object) uploadMultipart(in io.Reader, leaf, directoryID string, size int64, modTime time.Time) (err error) {
	// Create upload session
	session, err := o.createUploadSession(leaf, directoryID, size)
	if err != nil {
		return errors.Wrap(err, "multipart upload create session failed")
	}
	chunkSize := session.PartSize
	fs.Debugf(o, "Multipart upload session started for %d parts of size %v", session.TotalParts, fs.SizeSuffix(chunkSize))

	// Cancel the session if something went wrong
	defer func() {
		if err != nil {
			fs.Debugf(o, "Cancelling multipart upload: %v", err)
			cancelErr := o.abortUpload(session.ID)
			if cancelErr != nil {
				fs.Logf(o, "Failed to cancel multipart upload: %v", err)
			}
		}
	}()

	// Upload the chunks
	remaining := size
	position := int64(0)
	parts := make([]api.Part, session.TotalParts)
	hash := sha1.New()
	errs := make(chan error, 1)
	var wg sync.WaitGroup
outer:
	for part := 0; part < session.TotalParts; part++ {
		// Check any errors
		select {
		case err = <-errs:
			break outer
		default:
		}

		reqSize := remaining
		if reqSize >= int64(chunkSize) {
			reqSize = int64(chunkSize)
		}

		// Make a block of memory
		buf := make([]byte, reqSize)

		// Read the chunk
		_, err = io.ReadFull(in, buf)
		if err != nil {
			err = errors.Wrap(err, "multipart upload failed to read source")
			break outer
		}

		// Make the global hash (must be done sequentially)
		_, _ = hash.Write(buf)

		// Transfer the chunk
		wg.Add(1)
		o.fs.uploadToken.Get()
		go func(part int, position int64) {
			defer wg.Done()
			defer o.fs.uploadToken.Put()
			fs.Debugf(o, "Uploading part %d/%d offset %v/%v part size %v", part+1, session.TotalParts, fs.SizeSuffix(position), fs.SizeSuffix(size), fs.SizeSuffix(chunkSize))
			partResponse, err := o.uploadPart(session.ID, position, size, buf)
			if err != nil {
				err = errors.Wrap(err, "multipart upload failed to upload part")
				select {
				case errs <- err:
				default:
				}
				return
			}
			parts[part] = partResponse.Part
		}(part, position)

		// ready for next block
		remaining -= chunkSize
		position += chunkSize
	}
	wg.Wait()
	if err == nil {
		select {
		case err = <-errs:
		default:
		}
	}
	if err != nil {
		return err
	}

	// Finalise the upload session
	result, err := o.commitUpload(session.ID, parts, modTime, hash.Sum(nil))
	if err != nil {
		return errors.Wrap(err, "multipart upload failed to finalize")
	}

	if result.TotalCount != 1 || len(result.Entries) != 1 {
		return errors.Errorf("multipart upload failed %v - not sure why", o)
	}
	return o.setMetaData(&result.Entries[0])
}
<commit_msg>box: improve accounting for chunked uploads<commit_after>// multpart upload for box

package box

import (
	"bytes"
	"crypto/sha1"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/ncw/rclone/backend/box/api"
	"github.com/ncw/rclone/fs"
	"github.com/ncw/rclone/fs/accounting"
	"github.com/ncw/rclone/lib/rest"
	"github.com/pkg/errors"
)

// createUploadSession creates an upload session for the object
func (o *Object) createUploadSession(leaf, directoryID string, size int64) (response *api.UploadSessionResponse, err error) {
	opts := rest.Opts{
		Method:  "POST",
		Path:    "/files/upload_sessions",
		RootURL: uploadURL,
	}
	request := api.UploadSessionRequest{
		FileSize: size,
	}
	// If object has an ID then it is existing so create a new version
	if o.id != "" {
		opts.Path = "/files/" + o.id + "/upload_sessions"
	} else {
		opts.Path = "/files/upload_sessions"
		request.FolderID = directoryID
		request.FileName = replaceReservedChars(leaf)
	}
	var resp *http.Response
	err = o.fs.pacer.Call(func() (bool, error) {
		resp, err = o.fs.srv.CallJSON(&opts, &request, &response)
		return shouldRetry(resp, err)
	})
	return
}

// sha1Digest produces a digest using sha1 as per RFC3230
func sha1Digest(digest []byte) string {
	return "sha=" + base64.StdEncoding.EncodeToString(digest)
}

// uploadPart uploads a part in an upload session
func (o *Object) uploadPart(SessionID string, offset, totalSize int64, chunk []byte, wrap accounting.WrapFn) (response *api.UploadPartResponse, err error) {
	chunkSize := int64(len(chunk))
	sha1sum := sha1.Sum(chunk)
	opts := rest.Opts{
		Method:        "PUT",
		Path:          "/files/upload_sessions/" + SessionID,
		RootURL:       uploadURL,
		ContentType:   "application/octet-stream",
		ContentLength: &chunkSize,
		ContentRange:  fmt.Sprintf("bytes %d-%d/%d", offset, offset+chunkSize-1, totalSize),
		ExtraHeaders: map[string]string{
			"Digest": sha1Digest(sha1sum[:]),
		},
	}
	var resp *http.Response
	err = o.fs.pacer.Call(func() (bool, error) {
		opts.Body = wrap(bytes.NewReader(chunk))
		resp, err = o.fs.srv.CallJSON(&opts, nil, &response)
		return shouldRetry(resp, err)
	})
	if err != nil {
		return nil, err
	}
	return response, nil
}

// commitUpload finishes an upload session
func (o *Object) commitUpload(SessionID string, parts []api.Part, modTime time.Time, sha1sum []byte) (result *api.FolderItems, err error) {
	opts := rest.Opts{
		Method:  "POST",
		Path:    "/files/upload_sessions/" + SessionID + "/commit",
		RootURL: uploadURL,
		ExtraHeaders: map[string]string{
			"Digest": sha1Digest(sha1sum),
		},
	}
	request := api.CommitUpload{
		Parts: parts,
	}
	request.Attributes.ContentModifiedAt = api.Time(modTime)
	request.Attributes.ContentCreatedAt = api.Time(modTime)
	var body []byte
	var resp *http.Response
	maxTries := fs.Config.LowLevelRetries
	const defaultDelay = 10
	var tries int
outer:
	for tries = 0; tries < maxTries; tries++ {
		err = o.fs.pacer.Call(func() (bool, error) {
			resp, err = o.fs.srv.CallJSON(&opts, &request, nil)
			if err != nil {
				return shouldRetry(resp, err)
			}
			body, err = rest.ReadBody(resp)
			return shouldRetry(resp, err)
		})
		delay := defaultDelay
		why := "unknown"
		if err != nil {
			// Sometimes we get 400 Error with
			// parts_mismatch immediately after uploading
			// the last part.  Ignore this error and wait.
			if boxErr, ok := err.(*api.Error); ok && boxErr.Code == "parts_mismatch" {
				why = err.Error()
			} else {
				return nil, err
			}
		} else {
			switch resp.StatusCode {
			case http.StatusOK, http.StatusCreated:
				break outer
			case http.StatusAccepted:
				why = "not ready yet"
				delayString := resp.Header.Get("Retry-After")
				if delayString != "" {
					delay, err = strconv.Atoi(delayString)
					if err != nil {
						fs.Debugf(o, "Couldn't decode Retry-After header %q: %v", delayString, err)
						delay = defaultDelay
					}
				}
			default:
				return nil, errors.Errorf("unknown HTTP status return %q (%d)", resp.Status, resp.StatusCode)
			}
		}
		fs.Debugf(o, "commit multipart upload failed %d/%d - trying again in %d seconds (%s)", tries+1, maxTries, delay, why)
		time.Sleep(time.Duration(delay) * time.Second)
	}
	if tries >= maxTries {
		return nil, errors.New("too many tries to commit multipart upload - increase --low-level-retries")
	}
	err = json.Unmarshal(body, &result)
	if err != nil {
		return nil, errors.Wrapf(err, "couldn't decode commit response: %q", body)
	}
	return result, nil
}

// abortUpload cancels an upload session
func (o *Object) abortUpload(SessionID string) (err error) {
	opts := rest.Opts{
		Method:     "DELETE",
		Path:       "/files/upload_sessions/" + SessionID,
		RootURL:    uploadURL,
		NoResponse: true,
	}
	var resp *http.Response
	err = o.fs.pacer.Call(func() (bool, error) {
		resp, err = o.fs.srv.Call(&opts)
		return shouldRetry(resp, err)
	})
	return err
}

// uploadMultipart uploads a file using multipart upload
func (o *Object) uploadMultipart(in io.Reader, leaf, directoryID string, size int64, modTime time.Time) (err error) {
	// Create upload session
	session, err := o.createUploadSession(leaf, directoryID, size)
	if err != nil {
		return errors.Wrap(err, "multipart upload create session failed")
	}
	chunkSize := session.PartSize
	fs.Debugf(o, "Multipart upload session started for %d parts of size %v", session.TotalParts, fs.SizeSuffix(chunkSize))

	// Cancel the session if something went wrong
	defer func() {
		if err != nil {
			fs.Debugf(o, "Cancelling multipart upload: %v", err)
			cancelErr := o.abortUpload(session.ID)
			if cancelErr != nil {
				fs.Logf(o, "Failed to cancel multipart upload: %v", err)
			}
		}
	}()

	// unwrap the accounting from the input, we use wrap to put it
	// back on after the buffering
	in, wrap := accounting.UnWrap(in)

	// Upload the chunks
	remaining := size
	position := int64(0)
	parts := make([]api.Part, session.TotalParts)
	hash := sha1.New()
	errs := make(chan error, 1)
	var wg sync.WaitGroup
outer:
	for part := 0; part < session.TotalParts; part++ {
		// Check any errors
		select {
		case err = <-errs:
			break outer
		default:
		}

		reqSize := remaining
		if reqSize >= int64(chunkSize) {
			reqSize = int64(chunkSize)
		}

		// Make a block of memory
		buf := make([]byte, reqSize)

		// Read the chunk
		_, err = io.ReadFull(in, buf)
		if err != nil {
			err = errors.Wrap(err, "multipart upload failed to read source")
			break outer
		}

		// Make the global hash (must be done sequentially)
		_, _ = hash.Write(buf)

		// Transfer the chunk
		wg.Add(1)
		o.fs.uploadToken.Get()
		go func(part int, position int64) {
			defer wg.Done()
			defer o.fs.uploadToken.Put()
			fs.Debugf(o, "Uploading part %d/%d offset %v/%v part size %v", part+1, session.TotalParts, fs.SizeSuffix(position), fs.SizeSuffix(size), fs.SizeSuffix(chunkSize))
			partResponse, err := o.uploadPart(session.ID, position, size, buf, wrap)
			if err != nil {
				err = errors.Wrap(err, "multipart upload failed to upload part")
				select {
				case errs <- err:
				default:
				}
				return
			}
			parts[part] = partResponse.Part
		}(part, position)

		// ready for next block
		remaining -= chunkSize
		position += chunkSize
	}
	wg.Wait()
	if err == nil {
		select {
		case err = <-errs:
		default:
		}
	}
	if err != nil {
		return err
	}

	// Finalise the upload session
	result, err := o.commitUpload(session.ID, parts, modTime, hash.Sum(nil))
	if err != nil {
		return errors.Wrap(err, "multipart upload failed to finalize")
	}

	if result.TotalCount != 1 || len(result.Entries) != 1 {
		return errors.Errorf("multipart upload failed %v - not sure why", o)
	}
	return o.setMetaData(&result.Entries[0])
}
<|endoftext|>