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

package runtime_test

import (
	"math"
	"runtime"
	"sync/atomic"
	"syscall"
	"testing"
	"time"
)

var stop = make(chan bool, 1)

func perpetuumMobile() {
	select {
	case <-stop:
	default:
		go perpetuumMobile()
	}
}

func TestStopTheWorldDeadlock(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	maxprocs := runtime.GOMAXPROCS(3)
	compl := make(chan bool, 2)
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GC()
		}
		compl <- true
	}()
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GOMAXPROCS(3)
		}
		compl <- true
	}()
	go perpetuumMobile()
	<-compl
	<-compl
	stop <- true
	runtime.GOMAXPROCS(maxprocs)
}

func TestYieldProgress(t *testing.T) {
	testYieldProgress(t, false)
}

func TestYieldLockedProgress(t *testing.T) {
	testYieldProgress(t, true)
}

func testYieldProgress(t *testing.T, locked bool) {
	c := make(chan bool)
	cack := make(chan bool)
	go func() {
		if locked {
			runtime.LockOSThread()
		}
		for {
			select {
			case <-c:
				cack <- true
				return
			default:
				runtime.Gosched()
			}
		}
	}()
	time.Sleep(10 * time.Millisecond)
	c <- true
	<-cack
}

func TestYieldLocked(t *testing.T) {
	const N = 10
	c := make(chan bool)
	go func() {
		runtime.LockOSThread()
		for i := 0; i < N; i++ {
			runtime.Gosched()
			time.Sleep(time.Millisecond)
		}
		c <- true
		// runtime.UnlockOSThread() is deliberately omitted
	}()
	<-c
}

func TestGoroutineParallelism(t *testing.T) {
	P := 4
	N := 10
	if testing.Short() {
		P = 3
		N = 3
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(P))
	for try := 0; try < N; try++ {
		done := make(chan bool)
		x := uint32(0)
		for p := 0; p < P; p++ {
			// Test that all P goroutines are scheduled at the same time
			go func(p int) {
				for i := 0; i < 3; i++ {
					expected := uint32(P*i + p)
					for atomic.LoadUint32(&x) != expected {
					}
					atomic.StoreUint32(&x, expected+1)
				}
				done <- true
			}(p)
		}
		for p := 0; p < P; p++ {
			<-done
		}
	}
}

func TestBlockLocked(t *testing.T) {
	const N = 10
	c := make(chan bool)
	go func() {
		runtime.LockOSThread()
		for i := 0; i < N; i++ {
			c <- true
		}
		runtime.UnlockOSThread()
	}()
	for i := 0; i < N; i++ {
		<-c
	}
}

func TestTimerFairness(t *testing.T) {
	done := make(chan bool)
	c := make(chan bool)
	for i := 0; i < 2; i++ {
		go func() {
			for {
				select {
				case c <- true:
				case <-done:
					return
				}
			}
		}()
	}

	timer := time.After(20 * time.Millisecond)
	for {
		select {
		case <-c:
		case <-timer:
			close(done)
			return
		}
	}
}

func TestTimerFairness2(t *testing.T) {
	done := make(chan bool)
	c := make(chan bool)
	for i := 0; i < 2; i++ {
		go func() {
			timer := time.After(20 * time.Millisecond)
			var buf [1]byte
			for {
				syscall.Read(0, buf[0:0])
				select {
				case c <- true:
				case <-c:
				case <-timer:
					done <- true
					return
				}
			}
		}()
	}
	<-done
	<-done
}

// The function is used to test preemption at split stack checks.
// Declaring a var avoids inlining at the call site.
var preempt = func() int {
	var a [128]int
	sum := 0
	for _, v := range a {
		sum += v
	}
	return sum
}

func TestPreemption(t *testing.T) {
	// Test that goroutines are preempted at function calls.
	N := 5
	if testing.Short() {
		N = 2
	}
	c := make(chan bool)
	var x uint32
	for g := 0; g < 2; g++ {
		go func(g int) {
			for i := 0; i < N; i++ {
				for atomic.LoadUint32(&x) != uint32(g) {
					preempt()
				}
				atomic.StoreUint32(&x, uint32(1-g))
			}
			c <- true
		}(g)
	}
	<-c
	<-c
}

func TestPreemptionGC(t *testing.T) {
	// Test that pending GC preempts running goroutines.
	P := 5
	N := 10
	if testing.Short() {
		P = 3
		N = 2
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(P + 1))
	var stop uint32
	for i := 0; i < P; i++ {
		go func() {
			for atomic.LoadUint32(&stop) == 0 {
				preempt()
			}
		}()
	}
	for i := 0; i < N; i++ {
		runtime.Gosched()
		runtime.GC()
	}
	atomic.StoreUint32(&stop, 1)
}

func TestGCFairness(t *testing.T) {
	output := executeTest(t, testGCFairnessSource, nil)
	want := "OK\n"
	if output != want {
		t.Fatalf("want %s, got %s\n", want, output)
	}
}

const testGCFairnessSource = `
package main

import (
	"fmt"
	"os"
	"runtime"
	"time"
)

func main() {
	runtime.GOMAXPROCS(1)
	f, err := os.Open("/dev/null")
	if os.IsNotExist(err) {
		// This test tests what it is intended to test only if writes are fast.
		// If there is no /dev/null, we just don't execute the test.
		fmt.Println("OK\n")
		return
	}
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	for i := 0; i < 2; i++ {
		go func() {
			for {
				f.Write([]byte("."))
			}
		}()
	}
	time.Sleep(10 * time.Millisecond)
	fmt.Println("OK")
}
`

func stackGrowthRecursive(i int) {
	var pad [128]uint64
	if i != 0 && pad[0] == 0 {
		stackGrowthRecursive(i - 1)
	}
}

func TestPreemptSplitBig(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in -short mode")
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(2))
	stop := make(chan int)
	go big(stop)
	for i := 0; i < 3; i++ {
		time.Sleep(10 * time.Microsecond) // let big start running
		runtime.GC()
	}
	close(stop)
}

func big(stop chan int) int {
	n := 0
	for {
		// delay so that gc is sure to have asked for a preemption
		for i := 0; i < 1e9; i++ {
			n++
		}

		// call bigframe, which used to miss the preemption in its prologue.
		bigframe(stop)

		// check if we've been asked to stop.
		select {
		case <-stop:
			return n
		}
	}
}

func bigframe(stop chan int) int {
	// not splitting the stack will overflow.
	// small will notice that it needs a stack split and will
	// catch the overflow.
	var x [8192]byte
	return small(stop, &x)
}

func small(stop chan int, x *[8192]byte) int {
	for i := range x {
		x[i] = byte(i)
	}
	sum := 0
	for i := range x {
		sum += int(x[i])
	}

	// keep small from being a leaf function, which might
	// make it not do any stack check at all.
	nonleaf(stop)

	return sum
}

func nonleaf(stop chan int) bool {
	// do something that won't be inlined:
	select {
	case <-stop:
		return true
	default:
		return false
	}
}

func TestSchedLocalQueue(t *testing.T) {
	runtime.TestSchedLocalQueue1()
}

func TestSchedLocalQueueSteal(t *testing.T) {
	runtime.TestSchedLocalQueueSteal1()
}

func benchmarkStackGrowth(b *testing.B, rec int) {
	const CallsPerSched = 1000
	procs := runtime.GOMAXPROCS(-1)
	N := int32(b.N / CallsPerSched)
	c := make(chan bool, procs)
	for p := 0; p < procs; p++ {
		go func() {
			for atomic.AddInt32(&N, -1) >= 0 {
				runtime.Gosched()
				for g := 0; g < CallsPerSched; g++ {
					stackGrowthRecursive(rec)
				}
			}
			c <- true
		}()
	}
	for p := 0; p < procs; p++ {
		<-c
	}
}

func BenchmarkStackGrowth(b *testing.B) {
	benchmarkStackGrowth(b, 10)
}

func BenchmarkStackGrowthDeep(b *testing.B) {
	benchmarkStackGrowth(b, 1024)
}

func BenchmarkCreateGoroutines(b *testing.B) {
	benchmarkCreateGoroutines(b, 1)
}

func BenchmarkCreateGoroutinesParallel(b *testing.B) {
	benchmarkCreateGoroutines(b, runtime.GOMAXPROCS(-1))
}

func benchmarkCreateGoroutines(b *testing.B, procs int) {
	c := make(chan bool)
	var f func(n int)
	f = func(n int) {
		if n == 0 {
			c <- true
			return
		}
		go f(n - 1)
	}
	for i := 0; i < procs; i++ {
		go f(b.N / procs)
	}
	for i := 0; i < procs; i++ {
		<-c
	}
}

type Matrix [][]float64

func BenchmarkMatmult(b *testing.B) {
	b.StopTimer()
	// matmult is O(N**3) but testing expects O(b.N),
	// so we need to take cube root of b.N
	n := int(math.Cbrt(float64(b.N))) + 1
	A := makeMatrix(n)
	B := makeMatrix(n)
	C := makeMatrix(n)
	b.StartTimer()
	matmult(nil, A, B, C, 0, n, 0, n, 0, n, 8)
}

func makeMatrix(n int) Matrix {
	m := make(Matrix, n)
	for i := 0; i < n; i++ {
		m[i] = make([]float64, n)
		for j := 0; j < n; j++ {
			m[i][j] = float64(i*n + j)
		}
	}
	return m
}

func matmult(done chan<- struct{}, A, B, C Matrix, i0, i1, j0, j1, k0, k1, threshold int) {
	di := i1 - i0
	dj := j1 - j0
	dk := k1 - k0
	if di >= dj && di >= dk && di >= threshold {
		// divide in two by y axis
		mi := i0 + di/2
		done1 := make(chan struct{}, 1)
		go matmult(done1, A, B, C, i0, mi, j0, j1, k0, k1, threshold)
		matmult(nil, A, B, C, mi, i1, j0, j1, k0, k1, threshold)
		<-done1
	} else if dj >= dk && dj >= threshold {
		// divide in two by x axis
		mj := j0 + dj/2
		done1 := make(chan struct{}, 1)
		go matmult(done1, A, B, C, i0, i1, j0, mj, k0, k1, threshold)
		matmult(nil, A, B, C, i0, i1, mj, j1, k0, k1, threshold)
		<-done1
	} else if dk >= threshold {
		// divide in two by "k" axis
		// deliberately not parallel because of data races
		mk := k0 + dk/2
		matmult(nil, A, B, C, i0, i1, j0, j1, k0, mk, threshold)
		matmult(nil, A, B, C, i0, i1, j0, j1, mk, k1, threshold)
	} else {
		// the matrices are small enough, compute directly
		for i := i0; i < i1; i++ {
			for j := j0; j < j1; j++ {
				for k := k0; k < k1; k++ {
					C[i][j] += A[i][k] * B[k][j]
				}
			}
		}
	}
	if done != nil {
		done <- struct{}{}
	}
}
<commit_msg>runtime: fix test on windows The test prints an excessive \n when /dev/null is not present.<commit_after>// Copyright 2011 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package runtime_test

import (
	"math"
	"runtime"
	"sync/atomic"
	"syscall"
	"testing"
	"time"
)

var stop = make(chan bool, 1)

func perpetuumMobile() {
	select {
	case <-stop:
	default:
		go perpetuumMobile()
	}
}

func TestStopTheWorldDeadlock(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping during short test")
	}
	maxprocs := runtime.GOMAXPROCS(3)
	compl := make(chan bool, 2)
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GC()
		}
		compl <- true
	}()
	go func() {
		for i := 0; i != 1000; i += 1 {
			runtime.GOMAXPROCS(3)
		}
		compl <- true
	}()
	go perpetuumMobile()
	<-compl
	<-compl
	stop <- true
	runtime.GOMAXPROCS(maxprocs)
}

func TestYieldProgress(t *testing.T) {
	testYieldProgress(t, false)
}

func TestYieldLockedProgress(t *testing.T) {
	testYieldProgress(t, true)
}

func testYieldProgress(t *testing.T, locked bool) {
	c := make(chan bool)
	cack := make(chan bool)
	go func() {
		if locked {
			runtime.LockOSThread()
		}
		for {
			select {
			case <-c:
				cack <- true
				return
			default:
				runtime.Gosched()
			}
		}
	}()
	time.Sleep(10 * time.Millisecond)
	c <- true
	<-cack
}

func TestYieldLocked(t *testing.T) {
	const N = 10
	c := make(chan bool)
	go func() {
		runtime.LockOSThread()
		for i := 0; i < N; i++ {
			runtime.Gosched()
			time.Sleep(time.Millisecond)
		}
		c <- true
		// runtime.UnlockOSThread() is deliberately omitted
	}()
	<-c
}

func TestGoroutineParallelism(t *testing.T) {
	P := 4
	N := 10
	if testing.Short() {
		P = 3
		N = 3
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(P))
	for try := 0; try < N; try++ {
		done := make(chan bool)
		x := uint32(0)
		for p := 0; p < P; p++ {
			// Test that all P goroutines are scheduled at the same time
			go func(p int) {
				for i := 0; i < 3; i++ {
					expected := uint32(P*i + p)
					for atomic.LoadUint32(&x) != expected {
					}
					atomic.StoreUint32(&x, expected+1)
				}
				done <- true
			}(p)
		}
		for p := 0; p < P; p++ {
			<-done
		}
	}
}

func TestBlockLocked(t *testing.T) {
	const N = 10
	c := make(chan bool)
	go func() {
		runtime.LockOSThread()
		for i := 0; i < N; i++ {
			c <- true
		}
		runtime.UnlockOSThread()
	}()
	for i := 0; i < N; i++ {
		<-c
	}
}

func TestTimerFairness(t *testing.T) {
	done := make(chan bool)
	c := make(chan bool)
	for i := 0; i < 2; i++ {
		go func() {
			for {
				select {
				case c <- true:
				case <-done:
					return
				}
			}
		}()
	}

	timer := time.After(20 * time.Millisecond)
	for {
		select {
		case <-c:
		case <-timer:
			close(done)
			return
		}
	}
}

func TestTimerFairness2(t *testing.T) {
	done := make(chan bool)
	c := make(chan bool)
	for i := 0; i < 2; i++ {
		go func() {
			timer := time.After(20 * time.Millisecond)
			var buf [1]byte
			for {
				syscall.Read(0, buf[0:0])
				select {
				case c <- true:
				case <-c:
				case <-timer:
					done <- true
					return
				}
			}
		}()
	}
	<-done
	<-done
}

// The function is used to test preemption at split stack checks.
// Declaring a var avoids inlining at the call site.
var preempt = func() int {
	var a [128]int
	sum := 0
	for _, v := range a {
		sum += v
	}
	return sum
}

func TestPreemption(t *testing.T) {
	// Test that goroutines are preempted at function calls.
	N := 5
	if testing.Short() {
		N = 2
	}
	c := make(chan bool)
	var x uint32
	for g := 0; g < 2; g++ {
		go func(g int) {
			for i := 0; i < N; i++ {
				for atomic.LoadUint32(&x) != uint32(g) {
					preempt()
				}
				atomic.StoreUint32(&x, uint32(1-g))
			}
			c <- true
		}(g)
	}
	<-c
	<-c
}

func TestPreemptionGC(t *testing.T) {
	// Test that pending GC preempts running goroutines.
	P := 5
	N := 10
	if testing.Short() {
		P = 3
		N = 2
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(P + 1))
	var stop uint32
	for i := 0; i < P; i++ {
		go func() {
			for atomic.LoadUint32(&stop) == 0 {
				preempt()
			}
		}()
	}
	for i := 0; i < N; i++ {
		runtime.Gosched()
		runtime.GC()
	}
	atomic.StoreUint32(&stop, 1)
}

func TestGCFairness(t *testing.T) {
	output := executeTest(t, testGCFairnessSource, nil)
	want := "OK\n"
	if output != want {
		t.Fatalf("want %s, got %s\n", want, output)
	}
}

const testGCFairnessSource = `
package main

import (
	"fmt"
	"os"
	"runtime"
	"time"
)

func main() {
	runtime.GOMAXPROCS(1)
	f, err := os.Open("/dev/null")
	if os.IsNotExist(err) {
		// This test tests what it is intended to test only if writes are fast.
		// If there is no /dev/null, we just don't execute the test.
		fmt.Println("OK")
		return
	}
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
	for i := 0; i < 2; i++ {
		go func() {
			for {
				f.Write([]byte("."))
			}
		}()
	}
	time.Sleep(10 * time.Millisecond)
	fmt.Println("OK")
}
`

func stackGrowthRecursive(i int) {
	var pad [128]uint64
	if i != 0 && pad[0] == 0 {
		stackGrowthRecursive(i - 1)
	}
}

func TestPreemptSplitBig(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping in -short mode")
	}
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(2))
	stop := make(chan int)
	go big(stop)
	for i := 0; i < 3; i++ {
		time.Sleep(10 * time.Microsecond) // let big start running
		runtime.GC()
	}
	close(stop)
}

func big(stop chan int) int {
	n := 0
	for {
		// delay so that gc is sure to have asked for a preemption
		for i := 0; i < 1e9; i++ {
			n++
		}

		// call bigframe, which used to miss the preemption in its prologue.
		bigframe(stop)

		// check if we've been asked to stop.
		select {
		case <-stop:
			return n
		}
	}
}

func bigframe(stop chan int) int {
	// not splitting the stack will overflow.
	// small will notice that it needs a stack split and will
	// catch the overflow.
	var x [8192]byte
	return small(stop, &x)
}

func small(stop chan int, x *[8192]byte) int {
	for i := range x {
		x[i] = byte(i)
	}
	sum := 0
	for i := range x {
		sum += int(x[i])
	}

	// keep small from being a leaf function, which might
	// make it not do any stack check at all.
	nonleaf(stop)

	return sum
}

func nonleaf(stop chan int) bool {
	// do something that won't be inlined:
	select {
	case <-stop:
		return true
	default:
		return false
	}
}

func TestSchedLocalQueue(t *testing.T) {
	runtime.TestSchedLocalQueue1()
}

func TestSchedLocalQueueSteal(t *testing.T) {
	runtime.TestSchedLocalQueueSteal1()
}

func benchmarkStackGrowth(b *testing.B, rec int) {
	const CallsPerSched = 1000
	procs := runtime.GOMAXPROCS(-1)
	N := int32(b.N / CallsPerSched)
	c := make(chan bool, procs)
	for p := 0; p < procs; p++ {
		go func() {
			for atomic.AddInt32(&N, -1) >= 0 {
				runtime.Gosched()
				for g := 0; g < CallsPerSched; g++ {
					stackGrowthRecursive(rec)
				}
			}
			c <- true
		}()
	}
	for p := 0; p < procs; p++ {
		<-c
	}
}

func BenchmarkStackGrowth(b *testing.B) {
	benchmarkStackGrowth(b, 10)
}

func BenchmarkStackGrowthDeep(b *testing.B) {
	benchmarkStackGrowth(b, 1024)
}

func BenchmarkCreateGoroutines(b *testing.B) {
	benchmarkCreateGoroutines(b, 1)
}

func BenchmarkCreateGoroutinesParallel(b *testing.B) {
	benchmarkCreateGoroutines(b, runtime.GOMAXPROCS(-1))
}

func benchmarkCreateGoroutines(b *testing.B, procs int) {
	c := make(chan bool)
	var f func(n int)
	f = func(n int) {
		if n == 0 {
			c <- true
			return
		}
		go f(n - 1)
	}
	for i := 0; i < procs; i++ {
		go f(b.N / procs)
	}
	for i := 0; i < procs; i++ {
		<-c
	}
}

type Matrix [][]float64

func BenchmarkMatmult(b *testing.B) {
	b.StopTimer()
	// matmult is O(N**3) but testing expects O(b.N),
	// so we need to take cube root of b.N
	n := int(math.Cbrt(float64(b.N))) + 1
	A := makeMatrix(n)
	B := makeMatrix(n)
	C := makeMatrix(n)
	b.StartTimer()
	matmult(nil, A, B, C, 0, n, 0, n, 0, n, 8)
}

func makeMatrix(n int) Matrix {
	m := make(Matrix, n)
	for i := 0; i < n; i++ {
		m[i] = make([]float64, n)
		for j := 0; j < n; j++ {
			m[i][j] = float64(i*n + j)
		}
	}
	return m
}

func matmult(done chan<- struct{}, A, B, C Matrix, i0, i1, j0, j1, k0, k1, threshold int) {
	di := i1 - i0
	dj := j1 - j0
	dk := k1 - k0
	if di >= dj && di >= dk && di >= threshold {
		// divide in two by y axis
		mi := i0 + di/2
		done1 := make(chan struct{}, 1)
		go matmult(done1, A, B, C, i0, mi, j0, j1, k0, k1, threshold)
		matmult(nil, A, B, C, mi, i1, j0, j1, k0, k1, threshold)
		<-done1
	} else if dj >= dk && dj >= threshold {
		// divide in two by x axis
		mj := j0 + dj/2
		done1 := make(chan struct{}, 1)
		go matmult(done1, A, B, C, i0, i1, j0, mj, k0, k1, threshold)
		matmult(nil, A, B, C, i0, i1, mj, j1, k0, k1, threshold)
		<-done1
	} else if dk >= threshold {
		// divide in two by "k" axis
		// deliberately not parallel because of data races
		mk := k0 + dk/2
		matmult(nil, A, B, C, i0, i1, j0, j1, k0, mk, threshold)
		matmult(nil, A, B, C, i0, i1, j0, j1, mk, k1, threshold)
	} else {
		// the matrices are small enough, compute directly
		for i := i0; i < i1; i++ {
			for j := j0; j < j1; j++ {
				for k := k0; k < k1; k++ {
					C[i][j] += A[i][k] * B[k][j]
				}
			}
		}
	}
	if done != nil {
		done <- struct{}{}
	}
}
<|endoftext|>


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

package testing

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"time"
)

var matchBenchmarks = flag.String("test.bench", "", "regular expression to select benchmarks to run")
var benchTime = flag.Float64("test.benchtime", 1, "approximate run time for each benchmark, in seconds")

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type InternalBenchmark struct {
	Name string
	F    func(b *B)
}

// B is a type passed to Benchmark functions to manage benchmark
// timing and to specify the number of iterations to run.
type B struct {
	N         int
	benchmark InternalBenchmark
	ns        int64
	bytes     int64
	start     int64
}

// StartTimer starts timing a test.  This function is called automatically
// before a benchmark starts, but it can also used to resume timing after
// a call to StopTimer.
func (b *B) StartTimer() { b.start = time.Nanoseconds() }

// StopTimer stops timing a test.  This can be used to pause the timer
// while performing complex initialization that you don't
// want to measure.
func (b *B) StopTimer() {
	if b.start > 0 {
		b.ns += time.Nanoseconds() - b.start
	}
	b.start = 0
}

// ResetTimer stops the timer and sets the elapsed benchmark time to zero.
func (b *B) ResetTimer() {
	b.start = 0
	b.ns = 0
}

// SetBytes records the number of bytes processed in a single operation.
// If this is called, the benchmark will report ns/op and MB/s.
func (b *B) SetBytes(n int64) { b.bytes = n }

func (b *B) nsPerOp() int64 {
	if b.N <= 0 {
		return 0
	}
	return b.ns / int64(b.N)
}

// runN runs a single benchmark for the specified number of iterations.
func (b *B) runN(n int) {
	// Try to get a comparable environment for each run
	// by clearing garbage from previous runs.
	runtime.GC()
	b.N = n
	b.ResetTimer()
	b.StartTimer()
	b.benchmark.F(b)
	b.StopTimer()
}

func min(x, y int) int {
	if x > y {
		return y
	}
	return x
}

func max(x, y int) int {
	if x < y {
		return y
	}
	return x
}

// roundDown10 rounds a number down to the nearest power of 10.
func roundDown10(n int) int {
	var tens = 0
	// tens = floor(log_10(n))
	for n > 10 {
		n = n / 10
		tens++
	}
	// result = 10^tens
	result := 1
	for i := 0; i < tens; i++ {
		result *= 10
	}
	return result
}

// roundUp rounds x up to a number of the form [1eX, 2eX, 5eX].
func roundUp(n int) int {
	base := roundDown10(n)
	if n < (2 * base) {
		return 2 * base
	}
	if n < (5 * base) {
		return 5 * base
	}
	return 10 * base
}

// run times the benchmark function.  It gradually increases the number
// of benchmark iterations until the benchmark runs for a second in order
// to get a reasonable measurement.  It prints timing information in this form
//		testing.BenchmarkHello	100000		19 ns/op
func (b *B) run() BenchmarkResult {
	// Run the benchmark for a single iteration in case it's expensive.
	n := 1
	b.runN(n)
	// Run the benchmark for at least the specified amount of time.
	time := int64(*benchTime * 1e9)
	for b.ns < time && n < 1e9 {
		last := n
		// Predict iterations/sec.
		if b.nsPerOp() == 0 {
			n = 1e9
		} else {
			n = int(time / b.nsPerOp())
		}
		// Run more iterations than we think we'll need for a second (1.5x).
		// Don't grow too fast in case we had timing errors previously.
		// Be sure to run at least one more than last time.
		n = max(min(n+n/2, 100*last), last+1)
		// Round up to something easy to read.
		n = roundUp(n)
		b.runN(n)
	}
	return BenchmarkResult{b.N, b.ns, b.bytes}
}

// The results of a benchmark run.
type BenchmarkResult struct {
	N     int   // The number of iterations.
	Ns    int64 // The total time taken.
	Bytes int64 // Bytes processed in one iteration.
}

func (r BenchmarkResult) NsPerOp() int64 {
	if r.N <= 0 {
		return 0
	}
	return r.Ns / int64(r.N)
}

func (r BenchmarkResult) mbPerSec() float64 {
	if r.Bytes <= 0 || r.Ns <= 0 || r.N <= 0 {
		return 0
	}
	return float64(r.Bytes) * float64(r.N) / float64(r.Ns) * 1e3
}

func (r BenchmarkResult) String() string {
	mbs := r.mbPerSec()
	mb := ""
	if mbs != 0 {
		mb = fmt.Sprintf("\t%7.2f MB/s", mbs)
	}
	return fmt.Sprintf("%8d\t%10d ns/op%s", r.N, r.NsPerOp(), mb)
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func RunBenchmarks(matchString func(pat, str string) (bool, os.Error), benchmarks []InternalBenchmark) {
	// If no flag was specified, don't run benchmarks.
	if len(*matchBenchmarks) == 0 {
		return
	}
	for _, Benchmark := range benchmarks {
		matched, err := matchString(*matchBenchmarks, Benchmark.Name)
		if err != nil {
			println("invalid regexp for -test.bench:", err.String())
			os.Exit(1)
		}
		if !matched {
			continue
		}
		for _, procs := range cpuList {
			runtime.GOMAXPROCS(procs)
			b := &B{benchmark: Benchmark}
			r := b.run()
			benchName := Benchmark.Name
			if procs != 1 {
				benchName = fmt.Sprintf("%s-%d", Benchmark.Name, procs)
			}
			print(fmt.Sprintf("%s\t%v\n", benchName, r))
			if p := runtime.GOMAXPROCS(-1); p != procs {
				print(fmt.Sprintf("%s left GOMAXPROCS set to %d\n", benchName, p))
			}
		}
	}
}

// Benchmark benchmarks a single function. Useful for creating
// custom benchmarks that do not use gotest.
func Benchmark(f func(b *B)) BenchmarkResult {
	b := &B{benchmark: InternalBenchmark{"", f}}
	return b.run()
}
<commit_msg>testing: scale benchmark precision to 0.01ns if needed<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package testing

import (
	"flag"
	"fmt"
	"os"
	"runtime"
	"time"
)

var matchBenchmarks = flag.String("test.bench", "", "regular expression to select benchmarks to run")
var benchTime = flag.Float64("test.benchtime", 1, "approximate run time for each benchmark, in seconds")

// An internal type but exported because it is cross-package; part of the implementation
// of gotest.
type InternalBenchmark struct {
	Name string
	F    func(b *B)
}

// B is a type passed to Benchmark functions to manage benchmark
// timing and to specify the number of iterations to run.
type B struct {
	N         int
	benchmark InternalBenchmark
	ns        int64
	bytes     int64
	start     int64
}

// StartTimer starts timing a test.  This function is called automatically
// before a benchmark starts, but it can also used to resume timing after
// a call to StopTimer.
func (b *B) StartTimer() { b.start = time.Nanoseconds() }

// StopTimer stops timing a test.  This can be used to pause the timer
// while performing complex initialization that you don't
// want to measure.
func (b *B) StopTimer() {
	if b.start > 0 {
		b.ns += time.Nanoseconds() - b.start
	}
	b.start = 0
}

// ResetTimer stops the timer and sets the elapsed benchmark time to zero.
func (b *B) ResetTimer() {
	b.start = 0
	b.ns = 0
}

// SetBytes records the number of bytes processed in a single operation.
// If this is called, the benchmark will report ns/op and MB/s.
func (b *B) SetBytes(n int64) { b.bytes = n }

func (b *B) nsPerOp() int64 {
	if b.N <= 0 {
		return 0
	}
	return b.ns / int64(b.N)
}

// runN runs a single benchmark for the specified number of iterations.
func (b *B) runN(n int) {
	// Try to get a comparable environment for each run
	// by clearing garbage from previous runs.
	runtime.GC()
	b.N = n
	b.ResetTimer()
	b.StartTimer()
	b.benchmark.F(b)
	b.StopTimer()
}

func min(x, y int) int {
	if x > y {
		return y
	}
	return x
}

func max(x, y int) int {
	if x < y {
		return y
	}
	return x
}

// roundDown10 rounds a number down to the nearest power of 10.
func roundDown10(n int) int {
	var tens = 0
	// tens = floor(log_10(n))
	for n > 10 {
		n = n / 10
		tens++
	}
	// result = 10^tens
	result := 1
	for i := 0; i < tens; i++ {
		result *= 10
	}
	return result
}

// roundUp rounds x up to a number of the form [1eX, 2eX, 5eX].
func roundUp(n int) int {
	base := roundDown10(n)
	if n < (2 * base) {
		return 2 * base
	}
	if n < (5 * base) {
		return 5 * base
	}
	return 10 * base
}

// run times the benchmark function.  It gradually increases the number
// of benchmark iterations until the benchmark runs for a second in order
// to get a reasonable measurement.  It prints timing information in this form
//		testing.BenchmarkHello	100000		19 ns/op
func (b *B) run() BenchmarkResult {
	// Run the benchmark for a single iteration in case it's expensive.
	n := 1
	b.runN(n)
	// Run the benchmark for at least the specified amount of time.
	time := int64(*benchTime * 1e9)
	for b.ns < time && n < 1e9 {
		last := n
		// Predict iterations/sec.
		if b.nsPerOp() == 0 {
			n = 1e9
		} else {
			n = int(time / b.nsPerOp())
		}
		// Run more iterations than we think we'll need for a second (1.5x).
		// Don't grow too fast in case we had timing errors previously.
		// Be sure to run at least one more than last time.
		n = max(min(n+n/2, 100*last), last+1)
		// Round up to something easy to read.
		n = roundUp(n)
		b.runN(n)
	}
	return BenchmarkResult{b.N, b.ns, b.bytes}
}

// The results of a benchmark run.
type BenchmarkResult struct {
	N     int   // The number of iterations.
	Ns    int64 // The total time taken.
	Bytes int64 // Bytes processed in one iteration.
}

func (r BenchmarkResult) NsPerOp() int64 {
	if r.N <= 0 {
		return 0
	}
	return r.Ns / int64(r.N)
}

func (r BenchmarkResult) mbPerSec() float64 {
	if r.Bytes <= 0 || r.Ns <= 0 || r.N <= 0 {
		return 0
	}
	return float64(r.Bytes) * float64(r.N) / float64(r.Ns) * 1e3
}

func (r BenchmarkResult) String() string {
	mbs := r.mbPerSec()
	mb := ""
	if mbs != 0 {
		mb = fmt.Sprintf("\t%7.2f MB/s", mbs)
	}
	nsop := r.NsPerOp()
	ns := fmt.Sprintf("%10d ns/op", nsop)
	if r.N > 0 && nsop < 100 {
		// The format specifiers here make sure that
		// the ones digits line up for all three possible formats.
		if nsop < 10 {
			ns = fmt.Sprintf("%13.2f ns/op", float64(r.Ns)/float64(r.N))
		} else {
			ns = fmt.Sprintf("%12.1f ns/op", float64(r.Ns)/float64(r.N))
		}
	}
	return fmt.Sprintf("%8d\t%s%s", r.N, ns, mb)
}

// An internal function but exported because it is cross-package; part of the implementation
// of gotest.
func RunBenchmarks(matchString func(pat, str string) (bool, os.Error), benchmarks []InternalBenchmark) {
	// If no flag was specified, don't run benchmarks.
	if len(*matchBenchmarks) == 0 {
		return
	}
	for _, Benchmark := range benchmarks {
		matched, err := matchString(*matchBenchmarks, Benchmark.Name)
		if err != nil {
			println("invalid regexp for -test.bench:", err.String())
			os.Exit(1)
		}
		if !matched {
			continue
		}
		for _, procs := range cpuList {
			runtime.GOMAXPROCS(procs)
			b := &B{benchmark: Benchmark}
			r := b.run()
			benchName := Benchmark.Name
			if procs != 1 {
				benchName = fmt.Sprintf("%s-%d", Benchmark.Name, procs)
			}
			print(fmt.Sprintf("%s\t%v\n", benchName, r))
			if p := runtime.GOMAXPROCS(-1); p != procs {
				print(fmt.Sprintf("%s left GOMAXPROCS set to %d\n", benchName, p))
			}
		}
	}
}

// Benchmark benchmarks a single function. Useful for creating
// custom benchmarks that do not use gotest.
func Benchmark(f func(b *B)) BenchmarkResult {
	b := &B{benchmark: InternalBenchmark{"", f}}
	return b.run()
}
<|endoftext|>


<commit_before>package chunkstore

import (
	"log"
	"os"

	. "chunkymonkey/types"
)

type readRequest struct {
	chunkLoc     ChunkXz
	responseChan chan<- ChunkReadResult
}

type IChunkStoreForeground interface {
	ReadChunk(chunkLoc ChunkXz) (reader IChunkReader, err os.Error)
	SupportsWrite() bool
	Writer() IChunkWriter
	WriteChunk(writer IChunkWriter) os.Error
}

// ChunkService adapts an IChunkStoreForeground (which can only be accessed
// from one goroutine) to an IChunkStore.
type ChunkService struct {
	store  IChunkStoreForeground
	reads  chan readRequest
	writes chan IChunkWriter
}

func NewChunkService(store IChunkStoreForeground) (s *ChunkService) {
	return &ChunkService{
		store:  store,
		reads:  make(chan readRequest),
		writes: make(chan IChunkWriter),
	}
}

func (s *ChunkService) Serve() {
	for {
		select {
		case request := <-s.reads:
			reader, err := s.store.ReadChunk(request.chunkLoc)
			request.responseChan <- ChunkReadResult{reader, err}
		case writer := <-s.writes:
			err := s.store.WriteChunk(writer)
			log.Printf("Could not write chunk at %#v: %v", writer.ChunkLoc(), err)
		}
	}
}

func (s *ChunkService) ReadChunk(chunkLoc ChunkXz) <-chan ChunkReadResult {
	responseChan := make(chan ChunkReadResult)

	s.reads <- readRequest{
		chunkLoc:     chunkLoc,
		responseChan: responseChan,
	}

	return responseChan
}

func (s *ChunkService) SupportsWrite() bool {
	return s.store.SupportsWrite()
}

func (s *ChunkService) Writer() IChunkWriter {
	return s.store.Writer()
}

func (s *ChunkService) WriteChunk(writer IChunkWriter) {
	s.writes <- writer
}
<commit_msg>Fixed logging error for chunk write even when no error occurred.<commit_after>package chunkstore

import (
	"log"
	"os"

	. "chunkymonkey/types"
)

type readRequest struct {
	chunkLoc     ChunkXz
	responseChan chan<- ChunkReadResult
}

type IChunkStoreForeground interface {
	ReadChunk(chunkLoc ChunkXz) (reader IChunkReader, err os.Error)
	SupportsWrite() bool
	Writer() IChunkWriter
	WriteChunk(writer IChunkWriter) os.Error
}

// ChunkService adapts an IChunkStoreForeground (which can only be accessed
// from one goroutine) to an IChunkStore.
type ChunkService struct {
	store  IChunkStoreForeground
	reads  chan readRequest
	writes chan IChunkWriter
}

func NewChunkService(store IChunkStoreForeground) (s *ChunkService) {
	return &ChunkService{
		store:  store,
		reads:  make(chan readRequest),
		writes: make(chan IChunkWriter),
	}
}

func (s *ChunkService) Serve() {
	for {
		select {
		case request := <-s.reads:
			reader, err := s.store.ReadChunk(request.chunkLoc)
			request.responseChan <- ChunkReadResult{reader, err}
		case writer := <-s.writes:
			if err := s.store.WriteChunk(writer); err != nil {
				log.Printf("Could not write chunk at %#v: %v", writer.ChunkLoc(), err)
			}
		}
	}
}

func (s *ChunkService) ReadChunk(chunkLoc ChunkXz) <-chan ChunkReadResult {
	responseChan := make(chan ChunkReadResult)

	s.reads <- readRequest{
		chunkLoc:     chunkLoc,
		responseChan: responseChan,
	}

	return responseChan
}

func (s *ChunkService) SupportsWrite() bool {
	return s.store.SupportsWrite()
}

func (s *ChunkService) Writer() IChunkWriter {
	return s.store.Writer()
}

func (s *ChunkService) WriteChunk(writer IChunkWriter) {
	s.writes <- writer
}
<|endoftext|>


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

//go:generate ./COMPILE-PROTOS.sh

package main

import (
	"encoding/json"
	"fmt"
	"github.com/golang/protobuf/proto"
	"io/ioutil"
	"os"
	"sort"

	pb "openapi"
)

func DescribeMap(in interface{}, indent string) string {
	description := ""
	m, ok := in.(map[string]interface{})
	if ok {
		keys := make([]string, 0)
		for k, _ := range m {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			v := m[k]
			description += fmt.Sprintf("%s%s:\n", indent, k)
			description += DescribeMap(v, indent+"  ")
		}
		return description
	}
	a, ok := in.([]interface{})
	if ok {
		for i, v := range a {
			description += fmt.Sprintf("%s%d:\n", indent, i)
			description += DescribeMap(v, indent+"  ")
		}
		return description
	}
	description += fmt.Sprintf("%s%+v\n", indent, in)
	return description
}

func ReadDocumentFromFile(filename string) *pb.Document {
	examplesDir := os.Getenv("GOPATH") + "/src/github.com/googleapis/openapi-compiler/examples"
	file, e := ioutil.ReadFile(examplesDir + "/" + filename)
	if e != nil {
		fmt.Printf("File error: %v\n", e)
		os.Exit(1)
	}
	var raw interface{}
	json.Unmarshal(file, &raw)

	rawDescription := DescribeMap(raw, "")
	ioutil.WriteFile("petstore.txt", []byte(rawDescription), 0644)

	document := buildDocument(raw)
	return document
}

func main() {
	fmt.Printf("Compiling %s\n", version())
	document := ReadDocumentFromFile("petstore.json")
	ioutil.WriteFile("petstore.pb", []byte(proto.MarshalTextString(document)), 0644)
}
<commit_msg>Add a few command-line flags to openapi-compiler<commit_after>// Copyright 2016 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:generate ./COMPILE-PROTOS.sh

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"path"
	"sort"
	"strings"

	"github.com/golang/protobuf/proto"
)

func describeMap(in interface{}, indent string) string {
	description := ""
	m, ok := in.(map[string]interface{})
	if ok {
		keys := make([]string, 0)
		for k, _ := range m {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		for _, k := range keys {
			v := m[k]
			description += fmt.Sprintf("%s%s:\n", indent, k)
			description += describeMap(v, indent+"  ")
		}
		return description
	}
	a, ok := in.([]interface{})
	if ok {
		for i, v := range a {
			description += fmt.Sprintf("%s%d:\n", indent, i)
			description += describeMap(v, indent+"  ")
		}
		return description
	}
	description += fmt.Sprintf("%s%+v\n", indent, in)
	return description
}

func readFile(filename string) interface{} {
	file, e := ioutil.ReadFile(filename)
	if e != nil {
		fmt.Printf("File error: %v\n", e)
		os.Exit(1)
	}
	var raw interface{}
	json.Unmarshal(file, &raw)
	return raw
}

func main() {
	var input = flag.String("input", "", "OpenAPI source file to read")
	var rawInput = flag.Bool("raw", false, "Output the raw json input")
	var textProtobuf = flag.Bool("text", false, "Output a text protobuf representation")
	flag.Parse()

	if *input == "" {
		flag.PrintDefaults()
		return
	}

	fmt.Printf("Compiling %s (%s)\n", *input, version())

	raw := readFile(*input)
	if *rawInput {
		rawDescription := describeMap(raw, "")
		rawFileName := strings.TrimSuffix(path.Base(*input), path.Ext(*input)) + ".raw"
		ioutil.WriteFile(rawFileName, []byte(rawDescription), 0644)
	}

	document := buildDocument(raw)

	if *textProtobuf {
		textProtoFileName := strings.TrimSuffix(path.Base(*input), path.Ext(*input)) + ".text"
		ioutil.WriteFile(textProtoFileName, []byte(proto.MarshalTextString(document)), 0644)
	}
}
<|endoftext|>


<commit_before>package vfs

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/djherbis/times"
	"github.com/ncw/rclone/fstest"
	"github.com/spf13/pflag"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context" // switch to "context" when we stop supporting go1.6
)

// Check CacheMode it satisfies the pflag interface
var _ pflag.Value = (*CacheMode)(nil)

func TestCacheModeString(t *testing.T) {
	assert.Equal(t, "off", CacheModeOff.String())
	assert.Equal(t, "full", CacheModeFull.String())
	assert.Equal(t, "CacheMode(17)", CacheMode(17).String())
}

func TestCacheModeSet(t *testing.T) {
	var m CacheMode

	err := m.Set("full")
	assert.NoError(t, err)
	assert.Equal(t, CacheModeFull, m)

	err = m.Set("potato")
	assert.Error(t, err, "Unknown cache mode level")

	err = m.Set("")
	assert.Error(t, err, "Unknown cache mode level")
}

func TestCacheModeType(t *testing.T) {
	var m CacheMode
	assert.Equal(t, "string", m.Type())
}

// convert c.item to a string
func itemAsString(c *cache) string {
	c.itemMu.Lock()
	defer c.itemMu.Unlock()
	var out []string
	for name, item := range c.item {
		out = append(out, fmt.Sprintf("name=%q isFile=%v opens=%d", name, item.isFile, item.opens))
	}
	sort.Strings(out)
	return strings.Join(out, "\n")
}

func TestCacheNew(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	assert.Contains(t, c.root, "vfs")
	assert.Contains(t, c.f.Root(), filepath.Base(r.Fremote.Root()))
	assert.Equal(t, "", itemAsString(c))

	// mkdir
	p, err := c.mkdir("potato")
	require.NoError(t, err)
	assert.Equal(t, "potato", filepath.Base(p))
	assert.Equal(t, `name="" isFile=false opens=0`, itemAsString(c))

	fi, err := os.Stat(filepath.Dir(p))
	require.NoError(t, err)
	assert.True(t, fi.IsDir())

	// get
	item := c.get("potato")
	item2 := c.get("potato")
	assert.Equal(t, item, item2)
	assert.WithinDuration(t, time.Now(), item.atime, time.Second)

	// updateTime
	//.. before
	t1 := time.Now().Add(-60 * time.Minute)
	c.updateTime("potato", t1)
	item = c.get("potato")
	assert.NotEqual(t, t1, item.atime)
	assert.Equal(t, 0, item.opens)
	//..after
	t2 := time.Now().Add(60 * time.Minute)
	c.updateTime("potato", t2)
	item = c.get("potato")
	assert.Equal(t, t2, item.atime)
	assert.Equal(t, 0, item.opens)

	// open
	assert.Equal(t, `name="" isFile=false opens=0
name="potato" isFile=true opens=0`, itemAsString(c))
	c.open("potato")
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	item = c.get("potato")
	assert.WithinDuration(t, time.Now(), item.atime, time.Second)
	assert.Equal(t, 1, item.opens)

	// write the file
	err = ioutil.WriteFile(p, []byte("hello"), 0600)
	require.NoError(t, err)

	// read its atime
	fi, err = os.Stat(p)
	assert.NoError(t, err)
	atime := times.Get(fi).AccessTime()

	// updateAtimes
	item = c.get("potato")
	item.atime = time.Now().Add(-24 * time.Hour)
	err = c.updateAtimes()
	require.NoError(t, err)
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	item = c.get("potato")
	assert.Equal(t, atime, item.atime)

	// updateAtimes - not in the cache
	oldItem := item
	delete(c.item, "potato") // remove from cache
	err = c.updateAtimes()
	require.NoError(t, err)
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=0`, itemAsString(c))
	item = c.get("potato")
	assert.Equal(t, atime, item.atime)
	c.item["potato"] = oldItem // restore to cache

	// try purging with file open
	c.purgeOld(10 * time.Second)
	_, err = os.Stat(p)
	assert.NoError(t, err)

	// close
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.updateTime("potato", t2)
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.close("potato")
	assert.Equal(t, `name="" isFile=false opens=0
name="potato" isFile=true opens=0`, itemAsString(c))
	item = c.get("potato")
	assert.WithinDuration(t, time.Now(), item.atime, time.Second)
	assert.Equal(t, 0, item.opens)

	// try purging with file closed
	c.purgeOld(10 * time.Second)
	// ...nothing should happend
	_, err = os.Stat(p)
	assert.NoError(t, err)

	//.. purge again with -ve age
	c.purgeOld(-10 * time.Second)
	_, err = os.Stat(p)
	assert.True(t, os.IsNotExist(err))

	// clean - have tested the internals already
	c.clean()

	// cleanup
	err = c.cleanUp()
	require.NoError(t, err)
	_, err = os.Stat(c.root)
	assert.True(t, os.IsNotExist(err))
}

func TestCacheOpens(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	assert.Equal(t, "", itemAsString(c))
	c.open("potato")
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.open("potato")
	assert.Equal(t, `name="" isFile=false opens=2
name="potato" isFile=true opens=2`, itemAsString(c))
	c.close("potato")
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.close("potato")
	assert.Equal(t, `name="" isFile=false opens=0
name="potato" isFile=true opens=0`, itemAsString(c))

	c.open("potato")
	c.open("a/b/c/d/one")
	c.open("a/b/c/d/e/two")
	c.open("a/b/c/d/e/f/three")
	assert.Equal(t, `name="" isFile=false opens=4
name="a" isFile=false opens=3
name="a/b" isFile=false opens=3
name="a/b/c" isFile=false opens=3
name="a/b/c/d" isFile=false opens=3
name="a/b/c/d/e" isFile=false opens=2
name="a/b/c/d/e/f" isFile=false opens=1
name="a/b/c/d/e/f/three" isFile=true opens=1
name="a/b/c/d/e/two" isFile=true opens=1
name="a/b/c/d/one" isFile=true opens=1
name="potato" isFile=true opens=1`, itemAsString(c))

	c.close("potato")
	c.close("a/b/c/d/one")
	c.close("a/b/c/d/e/two")
	c.close("a/b/c/d/e/f/three")
	assert.Equal(t, `name="" isFile=false opens=0
name="a" isFile=false opens=0
name="a/b" isFile=false opens=0
name="a/b/c" isFile=false opens=0
name="a/b/c/d" isFile=false opens=0
name="a/b/c/d/e" isFile=false opens=0
name="a/b/c/d/e/f" isFile=false opens=0
name="a/b/c/d/e/f/three" isFile=true opens=0
name="a/b/c/d/e/two" isFile=true opens=0
name="a/b/c/d/one" isFile=true opens=0
name="potato" isFile=true opens=0`, itemAsString(c))

}

func TestCacheCacheDir(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	assert.Equal(t, "", itemAsString(c))

	c.cacheDir("dir")
	assert.Equal(t, `name="" isFile=false opens=0
name="dir" isFile=false opens=0`, itemAsString(c))

	c.cacheDir("dir/sub")
	assert.Equal(t, `name="" isFile=false opens=0
name="dir" isFile=false opens=0
name="dir/sub" isFile=false opens=0`, itemAsString(c))

	c.cacheDir("dir/sub2/subsub2")
	assert.Equal(t, `name="" isFile=false opens=0
name="dir" isFile=false opens=0
name="dir/sub" isFile=false opens=0
name="dir/sub2" isFile=false opens=0
name="dir/sub2/subsub2" isFile=false opens=0`, itemAsString(c))
}

func TestCachePurgeOld(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	// Test funcs
	var removed []string
	removedDir := true
	removeFile := func(name string) {
		removed = append(removed, name)
	}
	removeDir := func(name string) bool {
		if removedDir {
			removed = append(removed, name+"/")
		}
		return removedDir
	}

	removed = nil
	c._purgeOld(-10*time.Second, removeFile, removeDir)
	assert.Equal(t, []string(nil), removed)

	c.open("sub/dir/potato")
	c.close("sub/dir/potato")

	assert.Equal(t, `name="" isFile=false opens=0
name="sub" isFile=false opens=0
name="sub/dir" isFile=false opens=0
name="sub/dir/potato" isFile=true opens=0`, itemAsString(c))

	removed = nil
	removedDir = false
	c._purgeOld(10*time.Second, removeFile, removeDir)
	assert.Equal(t, []string(nil), removed)

	assert.Equal(t, `name="" isFile=false opens=0
name="sub" isFile=false opens=0
name="sub/dir" isFile=false opens=0
name="sub/dir/potato" isFile=true opens=0`, itemAsString(c))

	removed = nil
	removedDir = true
	c._purgeOld(-10*time.Second, removeFile, removeDir)
	assert.Equal(t, []string{
		"sub/dir/potato",
		"sub/dir/",
		"sub/",
		"/",
	}, removed)

	assert.Equal(t, ``, itemAsString(c))
}
<commit_msg>vfs: fix race in cache tests<commit_after>package vfs

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"testing"
	"time"

	"github.com/djherbis/times"
	"github.com/ncw/rclone/fstest"
	"github.com/spf13/pflag"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"golang.org/x/net/context" // switch to "context" when we stop supporting go1.6
)

// Check CacheMode it satisfies the pflag interface
var _ pflag.Value = (*CacheMode)(nil)

func TestCacheModeString(t *testing.T) {
	assert.Equal(t, "off", CacheModeOff.String())
	assert.Equal(t, "full", CacheModeFull.String())
	assert.Equal(t, "CacheMode(17)", CacheMode(17).String())
}

func TestCacheModeSet(t *testing.T) {
	var m CacheMode

	err := m.Set("full")
	assert.NoError(t, err)
	assert.Equal(t, CacheModeFull, m)

	err = m.Set("potato")
	assert.Error(t, err, "Unknown cache mode level")

	err = m.Set("")
	assert.Error(t, err, "Unknown cache mode level")
}

func TestCacheModeType(t *testing.T) {
	var m CacheMode
	assert.Equal(t, "string", m.Type())
}

// convert c.item to a string
func itemAsString(c *cache) string {
	c.itemMu.Lock()
	defer c.itemMu.Unlock()
	var out []string
	for name, item := range c.item {
		out = append(out, fmt.Sprintf("name=%q isFile=%v opens=%d", name, item.isFile, item.opens))
	}
	sort.Strings(out)
	return strings.Join(out, "\n")
}

func TestCacheNew(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	assert.Contains(t, c.root, "vfs")
	assert.Contains(t, c.f.Root(), filepath.Base(r.Fremote.Root()))
	assert.Equal(t, "", itemAsString(c))

	// mkdir
	p, err := c.mkdir("potato")
	require.NoError(t, err)
	assert.Equal(t, "potato", filepath.Base(p))
	assert.Equal(t, `name="" isFile=false opens=0`, itemAsString(c))

	fi, err := os.Stat(filepath.Dir(p))
	require.NoError(t, err)
	assert.True(t, fi.IsDir())

	// get
	item := c.get("potato")
	item2 := c.get("potato")
	assert.Equal(t, item, item2)
	assert.WithinDuration(t, time.Now(), item.atime, time.Second)

	// updateTime
	//.. before
	t1 := time.Now().Add(-60 * time.Minute)
	c.updateTime("potato", t1)
	item = c.get("potato")
	assert.NotEqual(t, t1, item.atime)
	assert.Equal(t, 0, item.opens)
	//..after
	t2 := time.Now().Add(60 * time.Minute)
	c.updateTime("potato", t2)
	item = c.get("potato")
	assert.Equal(t, t2, item.atime)
	assert.Equal(t, 0, item.opens)

	// open
	assert.Equal(t, `name="" isFile=false opens=0
name="potato" isFile=true opens=0`, itemAsString(c))
	c.open("potato")
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	item = c.get("potato")
	assert.WithinDuration(t, time.Now(), item.atime, time.Second)
	assert.Equal(t, 1, item.opens)

	// write the file
	err = ioutil.WriteFile(p, []byte("hello"), 0600)
	require.NoError(t, err)

	// read its atime
	fi, err = os.Stat(p)
	assert.NoError(t, err)
	atime := times.Get(fi).AccessTime()

	// updateAtimes
	item = c.get("potato")
	item.atime = time.Now().Add(-24 * time.Hour)
	err = c.updateAtimes()
	require.NoError(t, err)
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	item = c.get("potato")
	assert.Equal(t, atime, item.atime)

	// updateAtimes - not in the cache
	oldItem := item
	c.itemMu.Lock()
	delete(c.item, "potato") // remove from cache
	c.itemMu.Unlock()
	err = c.updateAtimes()
	require.NoError(t, err)
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=0`, itemAsString(c))
	item = c.get("potato")
	assert.Equal(t, atime, item.atime)
	c.itemMu.Lock()
	c.item["potato"] = oldItem // restore to cache
	c.itemMu.Unlock()

	// try purging with file open
	c.purgeOld(10 * time.Second)
	_, err = os.Stat(p)
	assert.NoError(t, err)

	// close
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.updateTime("potato", t2)
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.close("potato")
	assert.Equal(t, `name="" isFile=false opens=0
name="potato" isFile=true opens=0`, itemAsString(c))
	item = c.get("potato")
	assert.WithinDuration(t, time.Now(), item.atime, time.Second)
	assert.Equal(t, 0, item.opens)

	// try purging with file closed
	c.purgeOld(10 * time.Second)
	// ...nothing should happend
	_, err = os.Stat(p)
	assert.NoError(t, err)

	//.. purge again with -ve age
	c.purgeOld(-10 * time.Second)
	_, err = os.Stat(p)
	assert.True(t, os.IsNotExist(err))

	// clean - have tested the internals already
	c.clean()

	// cleanup
	err = c.cleanUp()
	require.NoError(t, err)
	_, err = os.Stat(c.root)
	assert.True(t, os.IsNotExist(err))
}

func TestCacheOpens(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	assert.Equal(t, "", itemAsString(c))
	c.open("potato")
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.open("potato")
	assert.Equal(t, `name="" isFile=false opens=2
name="potato" isFile=true opens=2`, itemAsString(c))
	c.close("potato")
	assert.Equal(t, `name="" isFile=false opens=1
name="potato" isFile=true opens=1`, itemAsString(c))
	c.close("potato")
	assert.Equal(t, `name="" isFile=false opens=0
name="potato" isFile=true opens=0`, itemAsString(c))

	c.open("potato")
	c.open("a/b/c/d/one")
	c.open("a/b/c/d/e/two")
	c.open("a/b/c/d/e/f/three")
	assert.Equal(t, `name="" isFile=false opens=4
name="a" isFile=false opens=3
name="a/b" isFile=false opens=3
name="a/b/c" isFile=false opens=3
name="a/b/c/d" isFile=false opens=3
name="a/b/c/d/e" isFile=false opens=2
name="a/b/c/d/e/f" isFile=false opens=1
name="a/b/c/d/e/f/three" isFile=true opens=1
name="a/b/c/d/e/two" isFile=true opens=1
name="a/b/c/d/one" isFile=true opens=1
name="potato" isFile=true opens=1`, itemAsString(c))

	c.close("potato")
	c.close("a/b/c/d/one")
	c.close("a/b/c/d/e/two")
	c.close("a/b/c/d/e/f/three")
	assert.Equal(t, `name="" isFile=false opens=0
name="a" isFile=false opens=0
name="a/b" isFile=false opens=0
name="a/b/c" isFile=false opens=0
name="a/b/c/d" isFile=false opens=0
name="a/b/c/d/e" isFile=false opens=0
name="a/b/c/d/e/f" isFile=false opens=0
name="a/b/c/d/e/f/three" isFile=true opens=0
name="a/b/c/d/e/two" isFile=true opens=0
name="a/b/c/d/one" isFile=true opens=0
name="potato" isFile=true opens=0`, itemAsString(c))

}

func TestCacheCacheDir(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	assert.Equal(t, "", itemAsString(c))

	c.cacheDir("dir")
	assert.Equal(t, `name="" isFile=false opens=0
name="dir" isFile=false opens=0`, itemAsString(c))

	c.cacheDir("dir/sub")
	assert.Equal(t, `name="" isFile=false opens=0
name="dir" isFile=false opens=0
name="dir/sub" isFile=false opens=0`, itemAsString(c))

	c.cacheDir("dir/sub2/subsub2")
	assert.Equal(t, `name="" isFile=false opens=0
name="dir" isFile=false opens=0
name="dir/sub" isFile=false opens=0
name="dir/sub2" isFile=false opens=0
name="dir/sub2/subsub2" isFile=false opens=0`, itemAsString(c))
}

func TestCachePurgeOld(t *testing.T) {
	r := fstest.NewRun(t)
	defer r.Finalise()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	c, err := newCache(ctx, r.Fremote, &DefaultOpt)
	require.NoError(t, err)

	// Test funcs
	var removed []string
	removedDir := true
	removeFile := func(name string) {
		removed = append(removed, name)
	}
	removeDir := func(name string) bool {
		if removedDir {
			removed = append(removed, name+"/")
		}
		return removedDir
	}

	removed = nil
	c._purgeOld(-10*time.Second, removeFile, removeDir)
	assert.Equal(t, []string(nil), removed)

	c.open("sub/dir/potato")
	c.close("sub/dir/potato")

	assert.Equal(t, `name="" isFile=false opens=0
name="sub" isFile=false opens=0
name="sub/dir" isFile=false opens=0
name="sub/dir/potato" isFile=true opens=0`, itemAsString(c))

	removed = nil
	removedDir = false
	c._purgeOld(10*time.Second, removeFile, removeDir)
	assert.Equal(t, []string(nil), removed)

	assert.Equal(t, `name="" isFile=false opens=0
name="sub" isFile=false opens=0
name="sub/dir" isFile=false opens=0
name="sub/dir/potato" isFile=true opens=0`, itemAsString(c))

	removed = nil
	removedDir = true
	c._purgeOld(-10*time.Second, removeFile, removeDir)
	assert.Equal(t, []string{
		"sub/dir/potato",
		"sub/dir/",
		"sub/",
		"/",
	}, removed)

	assert.Equal(t, ``, itemAsString(c))
}
<|endoftext|>


<commit_before>package net

import "math"
import "math/rand"
import "time"

var randr *rand.Rand

func init() {
	t := time.Now()
	s := rand.NewSource(t.Unix() ^ t.UnixNano())
	randr = rand.New(s)
}

// Expresses a backoff and retry specification.
//
// The nil value of this structure results in sensible defaults being used.
type Backoff struct {
	// The maximum number of attempts which may be made.
	// If this is 0, the number of attempts is unlimited.
	MaxTries int

	// The initial delay, in milliseconds. This is the delay used after the first
	// failed attempt.
	InitialDelay time.Duration // ms

	// The maximum delay, in milliseconds. This is the maximum delay between
	// attempts.
	MaxDelay time.Duration // ms

	// Determines when the maximum delay should be reached. If this is 5, the
	// maximum delay will be reached after 5 attempts have been made.
	MaxDelayAfterTries int

	// Positive float expressing the maximum factor by which the delay value may
	// be randomly inflated. e.g. specify 0.05 for a 5% variation. Set to zero to
	// disable jitter.
	Jitter float64

	// The current try. You should not need to set this yourself.
	CurrentTry int
}

// Initialises any nil field in Backoff with sensible defaults. You
// normally do not need to call this method yourself, as it will be called
// automatically.
func (rc *Backoff) InitDefaults() {
	if rc.InitialDelay == 0 {
		rc.InitialDelay = 5 * time.Second
	}
	if rc.MaxDelay == 0 {
		rc.MaxDelay = 120 * time.Second
	}
	if rc.MaxDelayAfterTries == 0 {
		rc.MaxDelayAfterTries = 10
	}
}

// Gets the next delay in milliseconds and increments the internal try counter.
func (rc *Backoff) NextDelay() time.Duration {
	rc.InitDefaults()

	if rc.MaxTries != 0 && rc.CurrentTry >= rc.MaxTries {
		return time.Duration(0)
	}

	initialDelay := float64(rc.InitialDelay)
	maxDelay := float64(rc.MaxDelay)
	maxDelayAfterTries := float64(rc.MaxDelayAfterTries)
	currentTry := float64(rc.CurrentTry)

	// [from backoff.c]
	k := math.Log2(maxDelay/initialDelay) / maxDelayAfterTries
	d := time.Duration(initialDelay * math.Exp2(currentTry*k))
	rc.CurrentTry++

	if d > rc.MaxDelay {
		d = rc.MaxDelay
	}

	if rc.Jitter != 0 {
		f := (randr.Float64() - 0.5) * 2 // random value in range [-1,1)
		d = time.Duration(float64(d) * (1 + rc.Jitter*f))
	}

	return d
}

// Sleep for the duration returned by NextDelay().
func (rc *Backoff) Sleep() {
	time.Sleep(rc.NextDelay())
}

// Sets the internal try counter to zero; the next delay returned will be
// InitialDelay again.
func (rc *Backoff) Reset() {
	rc.CurrentTry = 0
}
<commit_msg>sleep fix<commit_after>package net

import "math"
import "math/rand"
import "time"

var randr *rand.Rand

func init() {
	t := time.Now()
	s := rand.NewSource(t.Unix() ^ t.UnixNano())
	randr = rand.New(s)
}

// Expresses a backoff and retry specification.
//
// The nil value of this structure results in sensible defaults being used.
type Backoff struct {
	// The maximum number of attempts which may be made.
	// If this is 0, the number of attempts is unlimited.
	MaxTries int

	// The initial delay, in milliseconds. This is the delay used after the first
	// failed attempt.
	InitialDelay time.Duration // ms

	// The maximum delay, in milliseconds. This is the maximum delay between
	// attempts.
	MaxDelay time.Duration // ms

	// Determines when the maximum delay should be reached. If this is 5, the
	// maximum delay will be reached after 5 attempts have been made.
	MaxDelayAfterTries int

	// Positive float expressing the maximum factor by which the delay value may
	// be randomly inflated. e.g. specify 0.05 for a 5% variation. Set to zero to
	// disable jitter.
	Jitter float64

	// The current try. You should not need to set this yourself.
	CurrentTry int
}

// Initialises any nil field in Backoff with sensible defaults. You
// normally do not need to call this method yourself, as it will be called
// automatically.
func (rc *Backoff) InitDefaults() {
	if rc.InitialDelay == 0 {
		rc.InitialDelay = 5 * time.Second
	}
	if rc.MaxDelay == 0 {
		rc.MaxDelay = 120 * time.Second
	}
	if rc.MaxDelayAfterTries == 0 {
		rc.MaxDelayAfterTries = 10
	}
}

// Gets the next delay in milliseconds and increments the internal try counter.
func (rc *Backoff) NextDelay() time.Duration {
	rc.InitDefaults()

	if rc.MaxTries != 0 && rc.CurrentTry >= rc.MaxTries {
		return time.Duration(0)
	}

	initialDelay := float64(rc.InitialDelay)
	maxDelay := float64(rc.MaxDelay)
	maxDelayAfterTries := float64(rc.MaxDelayAfterTries)
	currentTry := float64(rc.CurrentTry)

	// [from backoff.c]
	k := math.Log2(maxDelay/initialDelay) / maxDelayAfterTries
	d := time.Duration(initialDelay * math.Exp2(currentTry*k))
	rc.CurrentTry++

	if d > rc.MaxDelay {
		d = rc.MaxDelay
	}

	if rc.Jitter != 0 {
		f := (randr.Float64() - 0.5) * 2 // random value in range [-1,1)
		d = time.Duration(float64(d) * (1 + rc.Jitter*f))
	}

	return d
}

// Sleep for the duration returned by NextDelay().
func (rc *Backoff) Sleep() bool {
	d := rc.NextDelay()
	if d != 0 {
		time.Sleep(d)
		return true
	}
	return false
}

// Sets the internal try counter to zero; the next delay returned will be
// InitialDelay again.
func (rc *Backoff) Reset() {
	rc.CurrentTry = 0
}
<|endoftext|>


<commit_before>package auth

import (
	"errors"
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/memcache"
)

const (
	tokensTable = "tokens"
	cookieName  = "app-cookie"
)

// Token keys
const (
	newTokenHeader       string = "new-auth-token"
	newTokenExpiryHeader string = "new-auth-token-expiry"
)

// Middleware .
type Middleware struct {
	// Reference to the session that is used to set/get account data
	Session Session

	// Whether the request should be allowed to continue if the token has expired, or no token exists.
	// This is useful for pages or endpoints that render/return differently based on whethe the user
	// is authenticated or not
	ContinueWithBadToken bool
}

// AuthenticateCookie authenticates the token with a request cookie
func (m *Middleware) AuthenticateCookie(c context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	var accountKey *datastore.Key

	c, cancel := context.WithCancel(c)
	returnURL := fmt.Sprintf("/signin?returnUrl=%s", r.RequestURI)

	run := func() error {
		cookie, err := r.Cookie(cookieName)
		if err != nil {
			return fmt.Errorf("failed getting cookie: %v", err)
		}

		token, err := m.getToken(c, cookie.Value)
		if err != nil {
			return fmt.Errorf("failed to get token: %v", err)
		}
		if token.isExpired() {
			return errors.New("expired token")
		}

		accountKey = token.Key.Parent()

		// if the token's expiry less than a week away, get new token and kill the current
		if token.willExpireIn(time.Hour * 24 * 7) {
			newToken, err := m.getNewToken(c, token)
			if err != nil {
				return fmt.Errorf("failed to get new token: %v", err)
			}

			http.SetCookie(w, &http.Cookie{
				Name:     cookieName,
				Expires:  time.Now().Add(time.Hour * 24 * 14),
				HttpOnly: true,
				Secure:   !appengine.IsDevAppServer(),
				Value:    newToken.Value(),
			})
		}

		// add accountKey to context
		c = m.Session.SetAccountKey(c, accountKey)

		return nil
	}

	if err := run(); err != nil && !m.ContinueWithBadToken {
		http.Redirect(w, r, returnURL, http.StatusTemporaryRedirect)
		cancel()
	}

	return c
}

// AuthenticateToken authenticates the Authorization request header token
func (m *Middleware) AuthenticateToken(c context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	// let option requests through
	if r.Method == http.MethodOptions {
		return c
	}

	c, cancel := context.WithCancel(c)

	run := func() error {
		authHeader := r.Header.Get("Authorization")
		if len(authHeader) <= len("token=") {
			return errors.New("missing token header")
		}

		// prevent token caching with blank string value
		rawToken := authHeader[len("token="):]
		if len(rawToken) == 0 {
			return errors.New("missing token value")
		}

		token, err := m.getToken(c, rawToken)
		if err != nil {
			return fmt.Errorf("failed to get token details: %v", err)
		}

		// if token has expired return 401
		if token.isExpired() {
			return errors.New("expired token")
		}

		accountKey := token.Key.Parent()

		// if the token's expiry less than a week away, get new token
		if token.willExpireIn(time.Hour * 24 * 7) {
			newToken, err := m.getNewToken(c, token)
			if err != nil {
				return fmt.Errorf("failed to create new token: %v", err)
			}

			// send back the new token values
			w.Header().Add(newTokenHeader, newToken.Value())
			w.Header().Add(newTokenExpiryHeader, newToken.Expiry.Format(time.RFC3339))
		}

		// add accountKey to context
		c = m.Session.SetAccountKey(c, accountKey)

		return nil
	}

	if err := run(); err != nil && !m.ContinueWithBadToken {
		w.WriteHeader(http.StatusUnauthorized)
		cancel()
	}

	return c
}

// Gets the token for the rawToken value
func (m *Middleware) getToken(c context.Context, rawToken string) (*Token, error) {
	var err error

	var token Token
	_, err = memcache.Gob.Get(c, rawToken, &token)
	if err != nil && err != memcache.ErrCacheMiss {
		return nil, err
	}

	if err == memcache.ErrCacheMiss {
		tokenKey, err := datastore.DecodeKey(rawToken)
		if err != nil {
			return nil, fmt.Errorf("decoding token key: %v", err)
		}

		var store = NewTokenStore()
		token.Key, err = store.Get(c, tokenKey, &token)
		if err != nil {
			return nil, err
		}

		// add the token to memcache
		err = memcache.Gob.Set(c, &memcache.Item{
			Key:        token.Value(),
			Object:     token,
			Expiration: -1 * time.Since(token.Expiry),
		})

		return &token, err
	}
	return &token, nil
}

// Creates a new token and links it to the account for the old token
func (m *Middleware) getNewToken(c context.Context, accountKey *datastore.Key) (*Token, error) {
	if accountKey == nil {
		return nil, errors.New("account key is required to create a token")
	}

	store := NewTokenStore()
	return store.Create(c, accountKey)
}
<commit_msg>add detail to returned error<commit_after>package auth

import (
	"errors"
	"fmt"
	"net/http"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/appengine"
	"google.golang.org/appengine/datastore"
	"google.golang.org/appengine/memcache"
)

const (
	tokensTable = "tokens"
	cookieName  = "app-cookie"
)

// Token keys
const (
	newTokenHeader       string = "new-auth-token"
	newTokenExpiryHeader string = "new-auth-token-expiry"
)

// Middleware .
type Middleware struct {
	// Reference to the session that is used to set/get account data
	Session Session

	// Whether the request should be allowed to continue if the token has expired, or no token exists.
	// This is useful for pages or endpoints that render/return differently based on whethe the user
	// is authenticated or not
	ContinueWithBadToken bool
}

// AuthenticateCookie authenticates the token with a request cookie
func (m *Middleware) AuthenticateCookie(c context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	var accountKey *datastore.Key

	c, cancel := context.WithCancel(c)
	returnURL := fmt.Sprintf("/signin?returnUrl=%s", r.RequestURI)

	run := func() error {
		cookie, err := r.Cookie(cookieName)
		if err != nil {
			return fmt.Errorf("failed getting cookie: %v", err)
		}

		token, err := m.getToken(c, cookie.Value)
		if err != nil {
			return fmt.Errorf("failed to get token: %v", err)
		}
		if token.isExpired() {
			return errors.New("expired token")
		}

		accountKey = token.Key.Parent()

		// if the token's expiry less than a week away, get new token and kill the current
		if token.willExpireIn(time.Hour * 24 * 7) {
			newToken, err := m.getNewToken(c, token)
			if err != nil {
				return fmt.Errorf("failed to get new token: %v", err)
			}

			http.SetCookie(w, &http.Cookie{
				Name:     cookieName,
				Expires:  time.Now().Add(time.Hour * 24 * 14),
				HttpOnly: true,
				Secure:   !appengine.IsDevAppServer(),
				Value:    newToken.Value(),
			})
		}

		// add accountKey to context
		c = m.Session.SetAccountKey(c, accountKey)

		return nil
	}

	if err := run(); err != nil && !m.ContinueWithBadToken {
		http.Redirect(w, r, returnURL, http.StatusTemporaryRedirect)
		cancel()
	}

	return c
}

// AuthenticateToken authenticates the Authorization request header token
func (m *Middleware) AuthenticateToken(c context.Context, w http.ResponseWriter, r *http.Request) context.Context {
	// let option requests through
	if r.Method == http.MethodOptions {
		return c
	}

	c, cancel := context.WithCancel(c)

	run := func() error {
		authHeader := r.Header.Get("Authorization")
		if len(authHeader) <= len("token=") {
			return errors.New("missing token header")
		}

		// prevent token caching with blank string value
		rawToken := authHeader[len("token="):]
		if len(rawToken) == 0 {
			return errors.New("missing token value")
		}

		token, err := m.getToken(c, rawToken)
		if err != nil {
			return fmt.Errorf("failed to get token details: %v", err)
		}

		// if token has expired return 401
		if token.isExpired() {
			return errors.New("expired token")
		}

		accountKey := token.Key.Parent()

		// if the token's expiry less than a week away, get new token
		if token.willExpireIn(time.Hour * 24 * 7) {
			newToken, err := m.getNewToken(c, token)
			if err != nil {
				return fmt.Errorf("failed to create new token: %v", err)
			}

			// send back the new token values
			w.Header().Add(newTokenHeader, newToken.Value())
			w.Header().Add(newTokenExpiryHeader, newToken.Expiry.Format(time.RFC3339))
		}

		// add accountKey to context
		c = m.Session.SetAccountKey(c, accountKey)

		return nil
	}

	if err := run(); err != nil && !m.ContinueWithBadToken {
		w.WriteHeader(http.StatusUnauthorized)
		cancel()
	}

	return c
}

// Gets the token for the rawToken value
func (m *Middleware) getToken(c context.Context, rawToken string) (*Token, error) {
	var err error

	var token Token
	_, err = memcache.Gob.Get(c, rawToken, &token)
	if err != nil && err != memcache.ErrCacheMiss {
		return nil, fmt.Errorf("faild to get token from memcache: %v", err)
	}

	if err == memcache.ErrCacheMiss {
		tokenKey, err := datastore.DecodeKey(rawToken)
		if err != nil {
			return nil, fmt.Errorf("decoding token key: %v", err)
		}

		var store = NewTokenStore()
		token.Key, err = store.Get(c, tokenKey, &token)
		if err != nil {
			return nil, err
		}

		// add the token to memcache
		err = memcache.Gob.Set(c, &memcache.Item{
			Key:        token.Value(),
			Object:     token,
			Expiration: -1 * time.Since(token.Expiry),
		})

		return &token, err
	}
	return &token, nil
}

// Creates a new token and links it to the account for the old token
func (m *Middleware) getNewToken(c context.Context, accountKey *datastore.Key) (*Token, error) {
	if accountKey == nil {
		return nil, errors.New("account key is required to create a token")
	}

	store := NewTokenStore()
	return store.Create(c, accountKey)
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"io/ioutil"
	"log"
	"os"
	"regexp"
	"strings"

	"path/filepath"

	"code.google.com/p/go.crypto/ssh"
	yaml "gopkg.in/yaml.v1"
)

type authorizedKeys struct {
	Users  map[string]map[string][]string
	Access map[string]map[string][]string
}

type configData struct {
	Log  string
	Data string
}

// Allow only lowercase letters, dot, underscore or hyphen
func safestring(s string) bool {
	r, err := regexp.Compile("[-.a-z_]+")
	if err != nil {
		log.Fatalln(err)
	}

	return r.MatchString(s)
}

// Error during program init
// print error + exit with code 2
func initError(s string) {
	os.Stderr.WriteString(s + "\n")
	os.Exit(2)
}

func main() {
	if len(os.Args) != 2 {
		initError("Need exactly one argument - username")
	}

	// Load config file
	const configfile = "/usr/local/etc/authorized-keys.conf"
	configsource, err := ioutil.ReadFile(configfile)
	config := configData{}
	if err != nil {
		os.Stderr.WriteString("Error opening config file " + configfile + ", using defaults\n")
		programDir := filepath.Dir(os.Args[0])
		config.Data = programDir + "/authorized-keys.yaml"
		config.Log = programDir + "/authorized-keys.log"
	} else {
		err = yaml.Unmarshal(configsource, &config)
		if err != nil {
			initError(err.Error())
		}
	}

	// Open log file
	l, err := os.OpenFile(config.Log, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
	if err != nil {
		initError("error opening log file: '" + config.Log + "'" + err.Error())
	}
	defer l.Close()
	log.SetOutput(io.MultiWriter(l, os.Stderr))

	// Username
	username := os.Args[1]
	if !safestring(username) {
		log.Fatalln("User '" + username + " is not safe string")
	}

	// Hostname (server name)
	hostname, _ := os.Hostname()
	if !safestring(hostname) {
		log.Fatalln("Host '" + hostname + "' is not safe string")
	}

	// Data file
	source, err := ioutil.ReadFile(config.Data)
	if err != nil {
		initError("Can't read data file " + config.Data)
	}
	authorizedKeys := authorizedKeys{}
	err = yaml.Unmarshal(source, &authorizedKeys)
	if err != nil {
		initError("Cannot parse yaml file " + err.Error())
	}

	// Debug
	log.Println("Requested authorized_keys for server '" + hostname + "' and user '" + username + "'")

	// Find aliases for server/user
	users, ok := authorizedKeys.Access[hostname][username]
	if !ok {
		log.Fatalln("No users found")
	}

	// Find all keys for aliases
	var keys []string
	for _, user := range users {
		k, ok := authorizedKeys.Users[user]["keys"]
		if !ok {
			log.Fatalln("User '" + user + "' doesn't have any keys defined")
		}
		// Validate and append each key
		for index, key := range k {
			_, comment, _, _, err := ssh.ParseAuthorizedKey([]byte(key))

			if err == nil {
				keys = append(keys, key)
				log.Printf("Using key %q for user %q", comment, user)
			} else {
				log.Printf("Skipping key %d for user %q", index, user)
			}
		}
	}

	// Output for sshd
	os.Stdout.WriteString(strings.Join(keys, "\n"))
}
<commit_msg>Command line options, noref<commit_after>// Authorized keys SSHD helper
package main

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

	"path/filepath"

	"crypto/md5"

	"code.google.com/p/go.crypto/ssh"
	docopt "github.com/docopt/docopt-go"
	yaml "gopkg.in/yaml.v1"
)

type authorizedKeys struct {
	Users  map[string]map[string][]string
	Access map[string]map[string][]string
}

type configData struct {
	Log  string
	Data string
}

// Allow only lowercase letters, dot, underscore or hyphen
func safestring(s string) bool {
	r, err := regexp.Compile("[-.a-z_]+")
	if err != nil {
		log.Fatalln(err)
	}

	return r.MatchString(s)
}

// Error during program init
// print error + exit with code 2
func initError(s string, code int) {
	os.Stderr.WriteString(s + "\n")
	os.Exit(code)
}

var builddate string

func main() {
	usage := `Authorized Keys.

Usage:
  authorized-keys <user> [--force-server=<server>] [--test]
  authorized-keys -h | --help
  authorized-keys -v | --version

Options:
  -h --help                Show this screen.
  -v --version             Show version.
  --force-server=<server>  Force server name.
  --test                  Test mode - no logging, just print users and key fingerprints.
`
	arguments, err := docopt.Parse(usage, nil, true, "Authorized Keys build "+builddate, false)

	if err != nil {
		initError(usage, 2)
	}

	var test = false
	if arguments["--test"] == true {
		test = true
	}

	if arguments["<user>"] == nil {
		initError("No user defined", 3)
	}

	// Load config file
	const configfile = "/usr/local/etc/authorized-keys.conf"
	configsource, err := ioutil.ReadFile(configfile)
	config := configData{}
	if err != nil {
		os.Stderr.WriteString("Error opening config file " + configfile + ", using defaults\n")
		programDir := filepath.Dir(os.Args[0])
		config.Data = programDir + "/authorized-keys.yaml"
		config.Log = programDir + "/authorized-keys.log"
	} else {
		err = yaml.Unmarshal(configsource, &config)
		if err != nil {
			initError(err.Error(), 4)
		}
	}

	// Open log file
	if !test {
		l, err := os.OpenFile(config.Log, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0660)
		if err != nil {
			initError("error opening log file: '"+config.Log+"'"+err.Error(), 5)
		}
		defer l.Close()
		log.SetOutput(io.MultiWriter(l, os.Stderr))
	} else {
		log.SetOutput(os.Stderr)
	}

	// Username
	username := fmt.Sprintf("%s", arguments["<user>"])
	if !safestring(username) {
		log.Fatalln("User '" + username + "' is not safe string")
	}

	// Hostname (server name)
	var hostname string
	if arguments["--force-server"] == nil {
		hostname, _ = os.Hostname()
	} else {
		hostname = fmt.Sprintf("%s", arguments["--force-server"])
	}

	if !safestring(hostname) {
		log.Fatalln("Host '" + hostname + "' is not safe string")
	}

	// Data file
	source, err := ioutil.ReadFile(config.Data)
	if err != nil {
		initError("Can't read data file "+config.Data, 5)
	}
	authorizedKeys := authorizedKeys{}
	err = yaml.Unmarshal(source, &authorizedKeys)
	if err != nil {
		initError("Cannot parse yaml file "+err.Error(), 6)
	}

	// Debug
	log.Println("Requested authorized_keys for server '" + hostname + "' and user '" + username + "'")

	// Find server/user
	_, ok := authorizedKeys.Access[hostname][username]
	if !ok {
		log.Fatalln("Server '" + hostname + "' not found")
	}

	// Find aliases for server/user
	users, ok := authorizedKeys.Access[hostname][username]
	if !ok {
		log.Fatalln("No users found")
	}

	// Find all keys for aliases
	var keys []string
	for _, user := range users {
		k, ok := authorizedKeys.Users[user]["keys"]
		if !ok {
			log.Fatalln("User '" + user + "' doesn't have any keys defined")
		}
		// Validate and append each key
		for index, key := range k {
			publickey, comment, _, _, err := ssh.ParseAuthorizedKey([]byte(key))

			if err == nil {
				keys = append(keys, key)
				if test {
					fingerprint := fmt.Sprintf("%x", md5.Sum(publickey.Marshal()))
					log.Printf("Using key %q for user %q. Fingerprint: %s", comment, user, formatFingerPrint(fingerprint))
				} else {
					log.Printf("Using key %q for user %q", comment, user)
				}

			} else {
				log.Printf("Skipping key %d for user %q", index, user)
			}
		}
	}

	// Output for sshd
	if !test {
		os.Stdout.WriteString(strings.Join(keys, "\n"))
	}
}

func formatFingerPrint(f string) string {
	var fprint []string
	for i := 0; i < len(f); i = i + 2 {
		fprint = append(fprint, f[i:i+2])
	}
	return strings.Join(fprint, ":")
}
<|endoftext|>


<commit_before>package agent

import (
	"errors"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/swarm-v2/api"
	"github.com/docker/swarm-v2/log"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

var (
	errSessionClosed = errors.New("agent: session closed")
)

// session encapsulates one round of registration with the manager. session
// starts the registration and heartbeat control cycle. Any failure will result
// in a complete shutdown of the session and it must be reestablished.
//
// All communication with the master is done through session.  Changes that
// flow into the agent, such as task assignment, are called back into the
// agent through errs, messages and tasks.
type session struct {
	agent     *Agent
	sessionID string
	errs      chan error
	messages  chan *api.SessionMessage
	tasks     chan *api.TasksMessage

	registered chan struct{} // closed registration
	closed     chan struct{}
}

func newSession(ctx context.Context, agent *Agent, delay time.Duration) *session {
	s := &session{
		agent:      agent,
		errs:       make(chan error),
		messages:   make(chan *api.SessionMessage),
		tasks:      make(chan *api.TasksMessage),
		registered: make(chan struct{}),
		closed:     make(chan struct{}),
	}

	go s.run(ctx, delay)
	return s
}

func (s *session) run(ctx context.Context, delay time.Duration) {
	time.Sleep(delay) // delay before registering.

	sessionID, err := s.register(ctx)
	if err != nil {
		select {
		case s.errs <- err:
		case <-s.closed:
		case <-ctx.Done():
		}
		return
	}

	ctx = log.WithLogger(ctx, log.G(ctx).WithField("session.id", sessionID))
	s.sessionID = sessionID
	close(s.registered)

	go runctx(ctx, s.heartbeat, s.closed, s.errs)
	go runctx(ctx, s.watch, s.closed, s.errs)
	go runctx(ctx, s.listen, s.closed, s.errs)
}

func (s *session) register(ctx context.Context) (string, error) {
	log.G(ctx).Debugf("(*session).register")
	client := api.NewDispatcherClient(s.agent.conn)

	description, err := s.agent.config.Executor.Describe(ctx)
	if err != nil {
		log.G(ctx).WithError(err).WithField("executor", s.agent.config.Executor).
			Errorf("node description unavailable")
		return "", err
	}

	resp, err := client.Register(ctx, &api.RegisterRequest{
		NodeID:      s.agent.config.ID,
		Description: description,
	})
	if err != nil {
		if grpc.Code(err) == codes.NotFound {
			return "", errNodeNotRegistered
		}

		return "", err
	}

	return resp.SessionID, nil
}

func (s *session) heartbeat(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).heartbeat")
	client := api.NewDispatcherClient(s.agent.conn)
	heartbeat := time.NewTimer(1) // send out a heartbeat right away
	defer heartbeat.Stop()

	for {
		select {
		case <-heartbeat.C:
			start := time.Now()
			resp, err := client.Heartbeat(ctx, &api.HeartbeatRequest{
				NodeID:    s.agent.config.ID,
				SessionID: s.sessionID,
			})
			if err != nil {
				if grpc.Code(err) == codes.NotFound {
					err = errNodeNotRegistered
				}

				return err
			}

			heartbeat.Reset(resp.Period)
			log.G(ctx).WithFields(
				logrus.Fields{
					"period":        resp.Period,
					"grpc.duration": time.Since(start),
				}).Debugf("heartbeat")
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *session) listen(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).listen")
	client := api.NewDispatcherClient(s.agent.conn)
	session, err := client.Session(ctx, &api.SessionRequest{
		NodeID:    s.agent.config.ID,
		SessionID: s.sessionID,
	})
	if err != nil {
		return err
	}

	for {
		resp, err := session.Recv()
		if err != nil {
			return err
		}

		select {
		case s.messages <- resp:
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *session) watch(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).watch")
	client := api.NewDispatcherClient(s.agent.conn)
	watch, err := client.Tasks(ctx, &api.TasksRequest{
		NodeID:    s.agent.config.ID,
		SessionID: s.sessionID})
	if err != nil {
		return err
	}

	for {
		resp, err := watch.Recv()
		if err != nil {
			return err
		}

		select {
		case s.tasks <- resp:
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// sendTaskStatus uses the current session to send the status of a single task.
func (s *session) sendTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {
	select {
	case <-s.registered:
		select {
		case <-s.closed:
			return errSessionClosed
		default:
		}
	case <-s.closed:
		return errSessionClosed
	case <-ctx.Done():
		return ctx.Err()
	}

	client := api.NewDispatcherClient(s.agent.conn)
	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{
		NodeID:    s.agent.config.ID,
		SessionID: s.sessionID,
		Updates: []*api.UpdateTaskStatusRequest_TaskStatusUpdate{
			{
				TaskID: taskID,
				Status: status,
			},
		},
	}); err != nil {
		return err
	}

	return nil
}

func (s *session) close() error {
	select {
	case <-s.closed:
		return errSessionClosed
	default:
		close(s.closed)
		return nil
	}
}

// runctx blocks until the function exits, closed is closed, or the context is
// cancelled. Call as part os go statement.
func runctx(ctx context.Context, fn func(ctx context.Context) error, closed chan struct{}, errs chan error) {
	select {
	case errs <- fn(ctx):
	case <-closed:
	case <-ctx.Done():
	}
}
<commit_msg>agent: Fix hostname overriding.<commit_after>package agent

import (
	"errors"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/swarm-v2/api"
	"github.com/docker/swarm-v2/log"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
)

var (
	errSessionClosed = errors.New("agent: session closed")
)

// session encapsulates one round of registration with the manager. session
// starts the registration and heartbeat control cycle. Any failure will result
// in a complete shutdown of the session and it must be reestablished.
//
// All communication with the master is done through session.  Changes that
// flow into the agent, such as task assignment, are called back into the
// agent through errs, messages and tasks.
type session struct {
	agent     *Agent
	sessionID string
	errs      chan error
	messages  chan *api.SessionMessage
	tasks     chan *api.TasksMessage

	registered chan struct{} // closed registration
	closed     chan struct{}
}

func newSession(ctx context.Context, agent *Agent, delay time.Duration) *session {
	s := &session{
		agent:      agent,
		errs:       make(chan error),
		messages:   make(chan *api.SessionMessage),
		tasks:      make(chan *api.TasksMessage),
		registered: make(chan struct{}),
		closed:     make(chan struct{}),
	}

	go s.run(ctx, delay)
	return s
}

func (s *session) run(ctx context.Context, delay time.Duration) {
	time.Sleep(delay) // delay before registering.

	sessionID, err := s.register(ctx)
	if err != nil {
		select {
		case s.errs <- err:
		case <-s.closed:
		case <-ctx.Done():
		}
		return
	}

	ctx = log.WithLogger(ctx, log.G(ctx).WithField("session.id", sessionID))
	s.sessionID = sessionID
	close(s.registered)

	go runctx(ctx, s.heartbeat, s.closed, s.errs)
	go runctx(ctx, s.watch, s.closed, s.errs)
	go runctx(ctx, s.listen, s.closed, s.errs)
}

func (s *session) register(ctx context.Context) (string, error) {
	log.G(ctx).Debugf("(*session).register")
	client := api.NewDispatcherClient(s.agent.conn)

	description, err := s.agent.config.Executor.Describe(ctx)
	if err != nil {
		log.G(ctx).WithError(err).WithField("executor", s.agent.config.Executor).
			Errorf("node description unavailable")
		return "", err
	}
	// Override hostname
	if s.agent.config.Hostname != "" {
		description.Hostname = s.agent.config.Hostname
	}

	resp, err := client.Register(ctx, &api.RegisterRequest{
		NodeID:      s.agent.config.ID,
		Description: description,
	})
	if err != nil {
		if grpc.Code(err) == codes.NotFound {
			return "", errNodeNotRegistered
		}

		return "", err
	}

	return resp.SessionID, nil
}

func (s *session) heartbeat(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).heartbeat")
	client := api.NewDispatcherClient(s.agent.conn)
	heartbeat := time.NewTimer(1) // send out a heartbeat right away
	defer heartbeat.Stop()

	for {
		select {
		case <-heartbeat.C:
			start := time.Now()
			resp, err := client.Heartbeat(ctx, &api.HeartbeatRequest{
				NodeID:    s.agent.config.ID,
				SessionID: s.sessionID,
			})
			if err != nil {
				if grpc.Code(err) == codes.NotFound {
					err = errNodeNotRegistered
				}

				return err
			}

			heartbeat.Reset(resp.Period)
			log.G(ctx).WithFields(
				logrus.Fields{
					"period":        resp.Period,
					"grpc.duration": time.Since(start),
				}).Debugf("heartbeat")
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *session) listen(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).listen")
	client := api.NewDispatcherClient(s.agent.conn)
	session, err := client.Session(ctx, &api.SessionRequest{
		NodeID:    s.agent.config.ID,
		SessionID: s.sessionID,
	})
	if err != nil {
		return err
	}

	for {
		resp, err := session.Recv()
		if err != nil {
			return err
		}

		select {
		case s.messages <- resp:
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

func (s *session) watch(ctx context.Context) error {
	log.G(ctx).Debugf("(*session).watch")
	client := api.NewDispatcherClient(s.agent.conn)
	watch, err := client.Tasks(ctx, &api.TasksRequest{
		NodeID:    s.agent.config.ID,
		SessionID: s.sessionID})
	if err != nil {
		return err
	}

	for {
		resp, err := watch.Recv()
		if err != nil {
			return err
		}

		select {
		case s.tasks <- resp:
		case <-s.closed:
			return errSessionClosed
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// sendTaskStatus uses the current session to send the status of a single task.
func (s *session) sendTaskStatus(ctx context.Context, taskID string, status *api.TaskStatus) error {
	select {
	case <-s.registered:
		select {
		case <-s.closed:
			return errSessionClosed
		default:
		}
	case <-s.closed:
		return errSessionClosed
	case <-ctx.Done():
		return ctx.Err()
	}

	client := api.NewDispatcherClient(s.agent.conn)
	if _, err := client.UpdateTaskStatus(ctx, &api.UpdateTaskStatusRequest{
		NodeID:    s.agent.config.ID,
		SessionID: s.sessionID,
		Updates: []*api.UpdateTaskStatusRequest_TaskStatusUpdate{
			{
				TaskID: taskID,
				Status: status,
			},
		},
	}); err != nil {
		return err
	}

	return nil
}

func (s *session) close() error {
	select {
	case <-s.closed:
		return errSessionClosed
	default:
		close(s.closed)
		return nil
	}
}

// runctx blocks until the function exits, closed is closed, or the context is
// cancelled. Call as part os go statement.
func runctx(ctx context.Context, fn func(ctx context.Context) error, closed chan struct{}, errs chan error) {
	select {
	case errs <- fn(ctx):
	case <-closed:
	case <-ctx.Done():
	}
}
<|endoftext|>