<commit_before>// Copyright ©2015 The gonum 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 path

import (
	"math"
	"testing"

	"github.com/gonum/graph"
	"github.com/gonum/graph/graphs/gen"
	"github.com/gonum/graph/simple"
)

var (
	gnpUndirected_10_tenth   = gnpUndirected(10, 0.1)
	gnpUndirected_100_tenth  = gnpUndirected(100, 0.1)
	gnpUndirected_1000_tenth = gnpUndirected(1000, 0.1)
	gnpUndirected_10_half    = gnpUndirected(10, 0.5)
	gnpUndirected_100_half   = gnpUndirected(100, 0.5)
	gnpUndirected_1000_half  = gnpUndirected(1000, 0.5)
)

func gnpUndirected(n int, p float64) graph.Undirected {
	g := simple.NewUndirectedGraph(0, math.Inf(1))
	gen.Gnp(g, n, p, nil)
	return g
}

func benchmarkAStarNilHeuristic(b *testing.B, g graph.Undirected) {
	var expanded int
	for i := 0; i < b.N; i++ {
		_, expanded = AStar(simple.Node(0), simple.Node(1), g, nil)
	}
	if expanded == 0 {
		b.Fatal("unexpected number of expanded nodes")
	}
}

func BenchmarkAStarGnp_10_tenth(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_10_tenth)
}
func BenchmarkAStarGnp_100_tenth(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_100_tenth)
}
func BenchmarkAStarGnp_1000_tenth(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_1000_tenth)
}
func BenchmarkAStarGnp_10_half(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_10_half)
}
func BenchmarkAStarGnp_100_half(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_100_half)
}
func BenchmarkAStarGnp_1000_half(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_1000_half)
}

var (
	nswUndirected_10_2_2_2   = navigableSmallWorldUndirected(10, 2, 2, 2)
	nswUndirected_10_2_5_2   = navigableSmallWorldUndirected(10, 2, 5, 2)
	nswUndirected_100_5_10_2 = navigableSmallWorldUndirected(100, 5, 10, 2)
	nswUndirected_100_5_20_2 = navigableSmallWorldUndirected(100, 5, 20, 2)
)

func navigableSmallWorldUndirected(n, p, q int, r float64) graph.Undirected {
	g := simple.NewUndirectedGraph(0, math.Inf(1))
	gen.NavigableSmallWorld(g, []int{n, n}, p, q, r, nil)
	return g
}

func coordinatesForID(n graph.Node, c, r int) [2]int {
	id := n.ID()
	if id >= c*r {
		panic("out of range")
	}
	return [2]int{id / r, id % r}
}

// manhattanBetween returns the Manhattan distance between a and b.
func manhattanBetween(a, b [2]int) float64 {
	var d int
	for i, v := range a {
		d += abs(v - b[i])
	}
	return float64(d)
}

func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

func benchmarkAStarHeuristic(b *testing.B, g graph.Undirected, h Heuristic) {
	var expanded int
	for i := 0; i < b.N; i++ {
		_, expanded = AStar(simple.Node(0), simple.Node(1), g, h)
	}
	if expanded == 0 {
		b.Fatal("unexpected number of expanded nodes")
	}
}

func BenchmarkAStarUndirectedmallWorld_10_2_2_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_10_2_2_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_10_2_2_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 10, 10), coordinatesForID(y, 10, 10))
	}
	benchmarkAStarHeuristic(b, nswUndirected_10_2_2_2, h)
}
func BenchmarkAStarUndirectedmallWorld_10_2_5_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_10_2_2_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_10_2_5_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 10, 10), coordinatesForID(y, 10, 10))
	}
	benchmarkAStarHeuristic(b, nswUndirected_10_2_2_2, h)
}
func BenchmarkAStarUndirectedmallWorld_100_5_10_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_100_5_10_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_100_5_10_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 100, 100), coordinatesForID(y, 100, 100))
	}
	benchmarkAStarHeuristic(b, nswUndirected_100_5_10_2, h)
}
func BenchmarkAStarUndirectedmallWorld_100_5_20_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_100_5_20_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_100_5_20_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 100, 100), coordinatesForID(y, 100, 100))
	}
	benchmarkAStarHeuristic(b, nswUndirected_100_5_20_2, h)
}
<commit_msg>path: fix benchmark call<commit_after>// Copyright ©2015 The gonum 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 path

import (
	"math"
	"testing"

	"github.com/gonum/graph"
	"github.com/gonum/graph/graphs/gen"
	"github.com/gonum/graph/simple"
)

var (
	gnpUndirected_10_tenth   = gnpUndirected(10, 0.1)
	gnpUndirected_100_tenth  = gnpUndirected(100, 0.1)
	gnpUndirected_1000_tenth = gnpUndirected(1000, 0.1)
	gnpUndirected_10_half    = gnpUndirected(10, 0.5)
	gnpUndirected_100_half   = gnpUndirected(100, 0.5)
	gnpUndirected_1000_half  = gnpUndirected(1000, 0.5)
)

func gnpUndirected(n int, p float64) graph.Undirected {
	g := simple.NewUndirectedGraph(0, math.Inf(1))
	gen.Gnp(g, n, p, nil)
	return g
}

func benchmarkAStarNilHeuristic(b *testing.B, g graph.Undirected) {
	var expanded int
	for i := 0; i < b.N; i++ {
		_, expanded = AStar(simple.Node(0), simple.Node(1), g, nil)
	}
	if expanded == 0 {
		b.Fatal("unexpected number of expanded nodes")
	}
}

func BenchmarkAStarGnp_10_tenth(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_10_tenth)
}
func BenchmarkAStarGnp_100_tenth(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_100_tenth)
}
func BenchmarkAStarGnp_1000_tenth(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_1000_tenth)
}
func BenchmarkAStarGnp_10_half(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_10_half)
}
func BenchmarkAStarGnp_100_half(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_100_half)
}
func BenchmarkAStarGnp_1000_half(b *testing.B) {
	benchmarkAStarNilHeuristic(b, gnpUndirected_1000_half)
}

var (
	nswUndirected_10_2_2_2   = navigableSmallWorldUndirected(10, 2, 2, 2)
	nswUndirected_10_2_5_2   = navigableSmallWorldUndirected(10, 2, 5, 2)
	nswUndirected_100_5_10_2 = navigableSmallWorldUndirected(100, 5, 10, 2)
	nswUndirected_100_5_20_2 = navigableSmallWorldUndirected(100, 5, 20, 2)
)

func navigableSmallWorldUndirected(n, p, q int, r float64) graph.Undirected {
	g := simple.NewUndirectedGraph(0, math.Inf(1))
	gen.NavigableSmallWorld(g, []int{n, n}, p, q, r, nil)
	return g
}

func coordinatesForID(n graph.Node, c, r int) [2]int {
	id := n.ID()
	if id >= c*r {
		panic("out of range")
	}
	return [2]int{id / r, id % r}
}

// manhattanBetween returns the Manhattan distance between a and b.
func manhattanBetween(a, b [2]int) float64 {
	var d int
	for i, v := range a {
		d += abs(v - b[i])
	}
	return float64(d)
}

func abs(a int) int {
	if a < 0 {
		return -a
	}
	return a
}

func benchmarkAStarHeuristic(b *testing.B, g graph.Undirected, h Heuristic) {
	var expanded int
	for i := 0; i < b.N; i++ {
		_, expanded = AStar(simple.Node(0), simple.Node(1), g, h)
	}
	if expanded == 0 {
		b.Fatal("unexpected number of expanded nodes")
	}
}

func BenchmarkAStarUndirectedmallWorld_10_2_2_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_10_2_2_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_10_2_2_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 10, 10), coordinatesForID(y, 10, 10))
	}
	benchmarkAStarHeuristic(b, nswUndirected_10_2_2_2, h)
}
func BenchmarkAStarUndirectedmallWorld_10_2_5_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_10_2_5_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_10_2_5_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 10, 10), coordinatesForID(y, 10, 10))
	}
	benchmarkAStarHeuristic(b, nswUndirected_10_2_5_2, h)
}
func BenchmarkAStarUndirectedmallWorld_100_5_10_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_100_5_10_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_100_5_10_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 100, 100), coordinatesForID(y, 100, 100))
	}
	benchmarkAStarHeuristic(b, nswUndirected_100_5_10_2, h)
}
func BenchmarkAStarUndirectedmallWorld_100_5_20_2(b *testing.B) {
	benchmarkAStarHeuristic(b, nswUndirected_100_5_20_2, nil)
}
func BenchmarkAStarUndirectedmallWorld_100_5_20_2_Heur(b *testing.B) {
	h := func(x, y graph.Node) float64 {
		return manhattanBetween(coordinatesForID(x, 100, 100), coordinatesForID(y, 100, 100))
	}
	benchmarkAStarHeuristic(b, nswUndirected_100_5_20_2, h)
}
<|endoftext|>


<commit_before>/*
Copyright (c) 2015, Emir Pasic
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package examples

import (
	"fmt"
	rbt "github.com/emirpasic/gods/trees/redblacktree"
)

type RedBlackTreeExtended struct {
	*rbt.Tree
}

func (tree *RedBlackTreeExtended) GetMin() (value interface{}, found bool) {
	node, found := tree.getMinFromNode(tree.Root)
	if node != nil {
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) GetMax() (value interface{}, found bool) {
	node, found := tree.getMaxFromNode(tree.Root)
	if node != nil {
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) RemoveMin() (value interface{}, deleted bool) {
	node, found := tree.getMinFromNode(tree.Root)
	if found {
		tree.Remove(node.Key)
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) RemoveMax() (value interface{}, deleted bool) {
	node, found := tree.getMaxFromNode(tree.Root)
	if found {
		tree.Remove(node.Key)
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) getMinFromNode(node *rbt.Node) (foundNode *rbt.Node, found bool) {
	if node == nil {
		return nil, false
	}
	if node.Left == nil {
		return node, true
	} else {
		return tree.getMinFromNode(node.Left)
	}
}

func (tree *RedBlackTreeExtended) getMaxFromNode(node *rbt.Node) (foundNode *rbt.Node, found bool) {
	if node == nil {
		return nil, false
	}
	if node.Right == nil {
		return node, true
	} else {
		return tree.getMaxFromNode(node.Right)
	}
}

func print(tree *RedBlackTreeExtended) {
	max, _ := tree.GetMax()
	min, _ := tree.GetMin()
	fmt.Printf("Value for max key: %v \n", max)
	fmt.Printf("Value for min key: %v \n", min)
	fmt.Println(tree)
}

// Find ceiling node of the input key, return its key and value or nil if no ceiling is found.
// Third return parameter is true if ceiling was found, otherwise false.
// Ceiling node is defined as the smallest node that is larger than or equal to the given node.
// A ceiling node may not be found, either because the tree is empty, or because
// all nodes in the tree is smaller than the given node.
// Key should adhere to the comparator's type assertion, otherwise method panics.
func (tree *RedBlackTreeExtended) Ceiling(key interface{}) (ceilingKey interface{}, value interface{}, found bool) {
	var ceiling *rbt.Node
	found = false
	comparator := tree.Comparator()

	node := tree.Root
	for node != nil {
		compare := comparator(key, node.Key)
		switch {
		case compare == 0:
			return node.Key, node.Value, true
		case compare < 0:
			ceiling, found = node, true
			node = node.Left
		case compare > 0:
			node = node.Right
		}
	}
	if found {
		return ceiling.Key, ceiling.Value, true
	}
	return nil, nil, false
}

// Find floor node of the input key, return its key and value or nil if no ceiling is found.
// Third return parameter is true if floor was found, otherwise false.
// Floor node is defined as the largest node that is smaller than or equal to the given node.
// A floor node may not be found, either because the tree is empty, or because
// all nodes in the tree is larger than the given node.
// Key should adhere to the comparator's type assertion, otherwise method panics.
func (tree *RedBlackTreeExtended) Floor(key interface{}) (floorKey interface{}, value interface{}, found bool) {
	var floor *rbt.Node
	found = false
	comparator := tree.Comparator()

	node := tree.Root
	for node != nil {
		compare := comparator(key, node.Key)
		switch {
		case compare == 0:
			return node.Key, node.Value, true
		case compare < 0:
			node = node.Left
		case compare > 0:
			floor, found = node, true
			node = node.Right
		}
	}
	if found {
		return floor.Key, floor.Value, true
	}
	return nil, nil, false
}

func RedBlackTreeExtendedExample() {
	tree := RedBlackTreeExtended{rbt.NewWithIntComparator()}

	tree.Put(1, "a") // 1->x (in order)
	tree.Put(2, "b") // 1->x, 2->b (in order)
	tree.Put(3, "c") // 1->x, 2->b, 3->c (in order)
	tree.Put(4, "d") // 1->x, 2->b, 3->c, 4->d (in order)
	tree.Put(5, "e") // 1->x, 2->b, 3->c, 4->d, 5->e (in order)

	print(&tree)
	// Value for max key: e
	// Value for min key: a
	// RedBlackTree
	// │       ┌── 5
	// │   ┌── 4
	// │   │   └── 3
	// └── 2
	//     └── 1

	tree.RemoveMin() // 2->b, 3->c, 4->d, 5->e (in order)
	tree.RemoveMax() // 2->b, 3->c, 4->d (in order)
	tree.RemoveMin() // 3->c, 4->d (in order)
	tree.RemoveMax() // 3->c (in order)

	print(&tree)
	// Value for max key: c
	// Value for min key: c
	// RedBlackTree
	// └── 3

	/*
	 * Ceiling and Floor functions
	 */
	tree = RedBlackTreeExtended{rbt.NewWithIntComparator()}
	tree.Put(1, "a")
	tree.Put(2, "b")
	tree.Put(4, "d")
	tree.Put(6, "f")
	tree.Put(7, "g")

	//index, ceiling, floor
	testValues := [][]interface{}{
		{0, 1, nil},
		{1, 1, 1},
		{2, 2, 2},
		{3, 4, 2},
		{4, 4, 4},
		{5, 6, 4},
		{6, 6, 6},
		{7, 7, 7},
		{8, nil, 7},
	}
	for _, tt := range testValues {
		actualCeiling, _, _ := tree.Ceiling(tt[0])
		actualFloor, _, _ := tree.Floor(tt[0])
		fmt.Printf("test key %d, expected (%d, %d), actual (%d, %d)\n", tt[0], tt[1], tt[2], actualCeiling, actualFloor)
	}
}
<commit_msg>Ceiling and Floor function have two return parameter instead of three<commit_after>/*
Copyright (c) 2015, Emir Pasic
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

package examples

import (
	"fmt"
	rbt "github.com/emirpasic/gods/trees/redblacktree"
)

type RedBlackTreeExtended struct {
	*rbt.Tree
}

func (tree *RedBlackTreeExtended) GetMin() (value interface{}, found bool) {
	node, found := tree.getMinFromNode(tree.Root)
	if node != nil {
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) GetMax() (value interface{}, found bool) {
	node, found := tree.getMaxFromNode(tree.Root)
	if node != nil {
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) RemoveMin() (value interface{}, deleted bool) {
	node, found := tree.getMinFromNode(tree.Root)
	if found {
		tree.Remove(node.Key)
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) RemoveMax() (value interface{}, deleted bool) {
	node, found := tree.getMaxFromNode(tree.Root)
	if found {
		tree.Remove(node.Key)
		return node.Value, found
	} else {
		return nil, false
	}
}

func (tree *RedBlackTreeExtended) getMinFromNode(node *rbt.Node) (foundNode *rbt.Node, found bool) {
	if node == nil {
		return nil, false
	}
	if node.Left == nil {
		return node, true
	} else {
		return tree.getMinFromNode(node.Left)
	}
}

func (tree *RedBlackTreeExtended) getMaxFromNode(node *rbt.Node) (foundNode *rbt.Node, found bool) {
	if node == nil {
		return nil, false
	}
	if node.Right == nil {
		return node, true
	} else {
		return tree.getMaxFromNode(node.Right)
	}
}

func print(tree *RedBlackTreeExtended) {
	max, _ := tree.GetMax()
	min, _ := tree.GetMin()
	fmt.Printf("Value for max key: %v \n", max)
	fmt.Printf("Value for min key: %v \n", min)
	fmt.Println(tree)
}

// Find ceiling node of the input key, return the ceiling node or nil if no ceiling is found.
// Second return parameter is true if ceiling was found, otherwise false.
//
// Ceiling node is defined as the smallest node that is larger than or equal to the given node.
// A ceiling node may not be found, either because the tree is empty, or because
// all nodes in the tree is smaller than the given node.
//
// Key should adhere to the comparator's type assertion, otherwise method panics.
func (tree *RedBlackTreeExtended) Ceiling(key interface{}) (ceiling *rbt.Node, found bool) {
	found = false
	comparator := tree.Comparator()

	node := tree.Root
	for node != nil {
		compare := comparator(key, node.Key)
		switch {
		case compare == 0:
			return node, true
		case compare < 0:
			ceiling, found = node, true
			node = node.Left
		case compare > 0:
			node = node.Right
		}
	}
	if found {
		return ceiling, true
	}
	return nil, false
}

// Find floor node of the input key, return the floor node or nil if no ceiling is found.
// Second return parameter is true if floor was found, otherwise false.
//
// Floor node is defined as the largest node that is smaller than or equal to the given node.
// A floor node may not be found, either because the tree is empty, or because
// all nodes in the tree is larger than the given node.
//
// Key should adhere to the comparator's type assertion, otherwise method panics.
func (tree *RedBlackTreeExtended) Floor(key interface{}) (floor *rbt.Node, found bool) {
	found = false
	comparator := tree.Comparator()

	node := tree.Root
	for node != nil {
		compare := comparator(key, node.Key)
		switch {
		case compare == 0:
			return node, true
		case compare < 0:
			node = node.Left
		case compare > 0:
			floor, found = node, true
			node = node.Right
		}
	}
	if found {
		return floor, true
	}
	return nil, false
}

func RedBlackTreeExtendedExample() {
	tree := RedBlackTreeExtended{rbt.NewWithIntComparator()}

	tree.Put(1, "a") // 1->x (in order)
	tree.Put(2, "b") // 1->x, 2->b (in order)
	tree.Put(3, "c") // 1->x, 2->b, 3->c (in order)
	tree.Put(4, "d") // 1->x, 2->b, 3->c, 4->d (in order)
	tree.Put(5, "e") // 1->x, 2->b, 3->c, 4->d, 5->e (in order)

	print(&tree)
	// Value for max key: e
	// Value for min key: a
	// RedBlackTree
	// │       ┌── 5
	// │   ┌── 4
	// │   │   └── 3
	// └── 2
	//     └── 1

	tree.RemoveMin() // 2->b, 3->c, 4->d, 5->e (in order)
	tree.RemoveMax() // 2->b, 3->c, 4->d (in order)
	tree.RemoveMin() // 3->c, 4->d (in order)
	tree.RemoveMax() // 3->c (in order)

	print(&tree)
	// Value for max key: c
	// Value for min key: c
	// RedBlackTree
	// └── 3

	/*
	 * Ceiling and Floor functions
	 */
	tree = RedBlackTreeExtended{rbt.NewWithIntComparator()}
	tree.Put(1, "a")
	tree.Put(2, "b")
	tree.Put(4, "d")
	tree.Put(6, "f")
	tree.Put(7, "g")

	//index, ceiling, floor
	testValues := [][]interface{}{
		{0, 1, nil},
		{1, 1, 1},
		{2, 2, 2},
		{3, 4, 2},
		{4, 4, 4},
		{5, 6, 4},
		{6, 6, 6},
		{7, 7, 7},
		{8, nil, 7},
	}
	for _, tt := range testValues {
		actualCeiling, _ := tree.Ceiling(tt[0])
		actualFloor, _ := tree.Floor(tt[0])
		fmt.Printf("test key %d, expected (%d, %d), actual (%d, %d)\n", tt[0], tt[1], tt[2], actualCeiling.Key, actualFloor.Key)
	}
}
<|endoftext|>


<commit_before>package core

import (
	"fmt"

	"github.com/satori/go.uuid"
)

// UserManager is the process that listens for new session connections and spins of the session go-routine
func UserManager(sessionListener <-chan UserSession, jobNotify <-chan Job, dataStore Database, userConn UserConnector, logger Logger) {

	//structure for our lookup table
	type sessionLookup struct {
		jobLink chan Job
	}

	userMap := make(map[uuid.UUID]sessionLookup)
	clientClosed := make(chan uuid.UUID, 64)

	//create notification hub

	//listen for user sessions and hook them in
	go func() {
		for {
			select {
			case s := <-sessionListener:
				//new user session
				userMap[s.GetGUID()] = sessionLookup{make(chan Job, 32)}
				logger.Info("User %v Connected", s.GetGUID().String())
				go userSession(s, userMap[s.GetGUID()].jobLink, dataStore, userConn, logger)
			case id := <-clientClosed:
				//user session has disconnected
				logger.Info("User %v Disconnected", id.String())
				delete(userMap, id)
			case j := <-jobNotify:
				//a job has been updated, is this user connected?
				session, ok := userMap[j.User]
				if ok {
					//tell session coroutine to inform user
					session.jobLink <- j
				}
			}
		}
	}()

}

func userSession(session UserSession, jobLink <-chan Job, dataStore Database, userConn UserConnector, logger Logger) {

	clientMsg := make(chan []byte, 32)
	clientClosed := make(chan bool)

	go session.Read(clientMsg, clientClosed)

	for {
		select {
		case j := <-jobLink:
			session.SendJob(0, j)
		case msg := <-clientMsg:
			//message from client
			m := userRequest{}
			m.load(msg)
			switch m.MessageType {
			case "IarUpload":
				userID, password, err := m.readPassword()
				if err != nil {
					logger.Error("Error reading iar upload request")
					continue
				}
				logger.Info("Iar upload request from %v:%v", userID, password)
				isValid, err := userConn.ValidatePassword(userID, password)
				if err != nil {
					session.SignalError(m.MessageID, err.Error())
				} else {
					if isValid {
						//password is valid, create the upload job
						job, err := dataStore.CreateLoadIarJob(userID, "/")
						if err != nil {
							session.SignalError(m.MessageID, err.Error())
						} else {
							session.SendJob(m.MessageID, job)
							session.SignalSuccess(m.MessageID, fmt.Sprintf("%v", job.ID))
						}
					} else {
						session.SignalError(m.MessageID, "Invalid Password")
					}
				}
			case "SetPassword":
				userID, password, err := m.readPassword()
				if err != nil {
					logger.Error("Error reading password request")
					continue
				}
				logger.Info("Setting password for %v to %v", userID, password)
				if userID != session.GetGUID() && session.GetAccessLevel() < 250 {
					session.SignalError(m.MessageID, "Permission Denied")
				} else {
					if password == "" {
						session.SignalError(m.MessageID, "Password Cannot be blank")
					} else {
						err = userConn.SetPassword(session.GetGUID(), password)
						if err != nil {
							session.SignalError(m.MessageID, err.Error())
						} else {
							session.SignalSuccess(m.MessageID, "Password Set Successfully")
						}
					}
				}
			case "GetDefaultConfig":
				if session.GetAccessLevel() > 249 {
					logger.Info("Serving Default Region Configs.  Request: %v", m.MessageID)
					cfgs, err := dataStore.GetDefaultConfigs()
					if err != nil {
						logger.Error("Error getting default configs: ", err)
					} else {
						for _, cfg := range cfgs {
							session.SendConfig(m.MessageID, cfg)
						}
						session.SignalSuccess(m.MessageID, "Default Config Retrieved")
					}
				}
			case "GetConfig":
				if session.GetAccessLevel() > 249 {
					logger.Info("Serving Region Configs.  Request: %v", m.MessageID)
					rid, err := m.readRegionID()
					if err != nil {
						logger.Error("Error reading region id for configs: ", err)
					} else {
						logger.Info("Serving Region Configs for %v.", rid)
						cfgs, err := dataStore.GetConfigs(rid)
						if err != nil {
							logger.Error("Error getting configs: ", err)
						} else {
							for _, cfg := range cfgs {
								session.SendConfig(m.MessageID, cfg)
							}
							session.SignalSuccess(m.MessageID, "Config Retrieved")
						}
					}
				}
			case "GetState":
				logger.Info("Service state request")
				users, err := userConn.GetUsers()
				if err != nil {
					logger.Error("Error lookin up activeuser account: ", err)
				}
				for _, user := range users {
					if user.Suspended && session.GetAccessLevel() < 250 {
						continue
					}
					session.SendUser(m.MessageID, user)
				}
				users = nil

				jobs, err := dataStore.GetJobsForUser(session.GetGUID())
				if err != nil {
					logger.Error("Error lookin up tasks: ", err)
				}
				for _, job := range jobs {
					session.SendJob(m.MessageID, job)
				}
				jobs = nil

				pendingUsers, err := dataStore.GetPendingUsers()
				if err != nil {
					logger.Error("Error lookin up pending user account: ", err)
				}
				for _, user := range pendingUsers {
					session.SendPendingUser(m.MessageID, user)
				}
				pendingUsers = nil

				//send regions this user may control
				regions, err := dataStore.GetRegions()
				if err != nil {
					logger.Error("Error lookin up user regions: ", err)
				}
				for _, r := range regions {
					session.SendRegion(0, r)
				}
				regions = nil

				//send Estate, Group, and Host dataManager
				estates, err := dataStore.GetEstates()
				if err != nil {
					logger.Error("Error lookin up estates: ", err)
				}
				for _, e := range estates {
					session.SendEstate(m.MessageID, e)
				}
				estates = nil
				groups, err := userConn.GetGroups()
				if err != nil {
					logger.Error("Error lookin up groups: ", err)
				}
				for _, g := range groups {
					session.SendGroup(m.MessageID, g)
				}
				groups = nil
				//only administrative users need host access
				if session.GetAccessLevel() > 249 {
					hosts, err := dataStore.GetHosts()
					if err != nil {
						logger.Error("Error lookin up hosts: ", err)
					}
					for _, h := range hosts {
						session.SendHost(m.MessageID, h)
					}
				}

				//signal to the client that we have completed initial state sync
				session.SignalSuccess(m.MessageID, "State Sync Complete")
				logger.Info("Sync Complete")

			default:
				logger.Error("Error on message from client: ", m.MessageType)

			}
		case <-clientClosed:
			//the client connection has closed
			logger.Info("Client went away")
			return
		}

	}
}
<commit_msg>adding logging to user requests<commit_after>package core

import (
	"fmt"

	"github.com/satori/go.uuid"
)

// UserManager is the process that listens for new session connections and spins of the session go-routine
func UserManager(sessionListener <-chan UserSession, jobNotify <-chan Job, dataStore Database, userConn UserConnector, logger Logger) {

	//structure for our lookup table
	type sessionLookup struct {
		jobLink chan Job
	}

	userMap := make(map[uuid.UUID]sessionLookup)
	clientClosed := make(chan uuid.UUID, 64)

	//create notification hub

	//listen for user sessions and hook them in
	go func() {
		for {
			select {
			case s := <-sessionListener:
				//new user session
				userMap[s.GetGUID()] = sessionLookup{make(chan Job, 32)}
				logger.Info("User %v Connected", s.GetGUID().String())
				go userSession(s, userMap[s.GetGUID()].jobLink, clientClosed, dataStore, userConn, logger)
			case id := <-clientClosed:
				//user session has disconnected
				logger.Info("User %v Disconnected", id.String())
				delete(userMap, id)
			case j := <-jobNotify:
				//a job has been updated, is this user connected?
				session, ok := userMap[j.User]
				if ok {
					//tell session coroutine to inform user
					session.jobLink <- j
				}
			}
		}
	}()

}

func userSession(session UserSession, jobLink <-chan Job, exitLink chan<- uuid.UUID, dataStore Database, userConn UserConnector, logger Logger) {

	clientMsg := make(chan []byte, 32)
	clientClosed := make(chan bool)

	go session.Read(clientMsg, clientClosed)

	for {
		select {
		case j := <-jobLink:
			session.SendJob(0, j)
		case msg := <-clientMsg:
			//message from client
			m := userRequest{}
			m.load(msg)
			switch m.MessageType {
			case "IarUpload":
				logger.Info("User %v requesting iar upload", session.GetGUID())
				userID, password, err := m.readPassword()
				if err != nil {
					logger.Error("Error reading iar upload request")
					continue
				}
				logger.Info("Iar upload request from %v:%v", userID, password)
				isValid, err := userConn.ValidatePassword(userID, password)
				if err != nil {
					session.SignalError(m.MessageID, err.Error())
				} else {
					if isValid {
						//password is valid, create the upload job
						job, err := dataStore.CreateLoadIarJob(userID, "/")
						if err != nil {
							session.SignalError(m.MessageID, err.Error())
						} else {
							session.SendJob(m.MessageID, job)
							session.SignalSuccess(m.MessageID, fmt.Sprintf("%v", job.ID))
						}
					} else {
						session.SignalError(m.MessageID, "Invalid Password")
					}
				}
			case "SetPassword":
				logger.Info("User %v requesting password change", session.GetGUID())
				userID, password, err := m.readPassword()
				if err != nil {
					logger.Error("Error reading password request")
					continue
				}
				logger.Info("Setting password for %v to %v", userID, password)
				if userID != session.GetGUID() && session.GetAccessLevel() < 250 {
					session.SignalError(m.MessageID, "Permission Denied")
				} else {
					if password == "" {
						session.SignalError(m.MessageID, "Password Cannot be blank")
					} else {
						err = userConn.SetPassword(session.GetGUID(), password)
						if err != nil {
							session.SignalError(m.MessageID, err.Error())
						} else {
							session.SignalSuccess(m.MessageID, "Password Set Successfully")
						}
					}
				}
			case "GetDefaultConfig":
				logger.Info("User %v requesting default configuration", session.GetGUID())
				if session.GetAccessLevel() > 249 {
					logger.Info("Serving Default Region Configs.  Request: %v", m.MessageID)
					cfgs, err := dataStore.GetDefaultConfigs()
					if err != nil {
						logger.Error("Error getting default configs: ", err)
					} else {
						for _, cfg := range cfgs {
							session.SendConfig(m.MessageID, cfg)
						}
						session.SignalSuccess(m.MessageID, "Default Config Retrieved")
					}
				}
			case "GetConfig":
				logger.Info("User %v requesting configuration", session.GetGUID())
				if session.GetAccessLevel() > 249 {
					logger.Info("Serving Region Configs.  Request: %v", m.MessageID)
					rid, err := m.readRegionID()
					if err != nil {
						logger.Error("Error reading region id for configs: ", err)
					} else {
						logger.Info("Serving Region Configs for %v.", rid)
						cfgs, err := dataStore.GetConfigs(rid)
						if err != nil {
							logger.Error("Error getting configs: ", err)
						} else {
							for _, cfg := range cfgs {
								session.SendConfig(m.MessageID, cfg)
							}
							session.SignalSuccess(m.MessageID, "Config Retrieved")
						}
					}
				}
			case "GetState":
				logger.Info("User %v requesting state sync", session.GetGUID())
				users, err := userConn.GetUsers()
				if err != nil {
					logger.Error("Error lookin up activeuser account: ", err)
				}
				for _, user := range users {
					if user.Suspended && session.GetAccessLevel() < 250 {
						continue
					}
					session.SendUser(m.MessageID, user)
				}
				users = nil

				jobs, err := dataStore.GetJobsForUser(session.GetGUID())
				if err != nil {
					logger.Error("Error lookin up tasks: ", err)
				}
				for _, job := range jobs {
					session.SendJob(m.MessageID, job)
				}
				jobs = nil

				pendingUsers, err := dataStore.GetPendingUsers()
				if err != nil {
					logger.Error("Error lookin up pending user account: ", err)
				}
				for _, user := range pendingUsers {
					session.SendPendingUser(m.MessageID, user)
				}
				pendingUsers = nil

				//send regions this user may control
				regions, err := dataStore.GetRegions()
				if err != nil {
					logger.Error("Error lookin up user regions: ", err)
				}
				for _, r := range regions {
					session.SendRegion(0, r)
				}
				regions = nil

				//send Estate, Group, and Host dataManager
				estates, err := dataStore.GetEstates()
				if err != nil {
					logger.Error("Error lookin up estates: ", err)
				}
				for _, e := range estates {
					session.SendEstate(m.MessageID, e)
				}
				estates = nil
				groups, err := userConn.GetGroups()
				if err != nil {
					logger.Error("Error lookin up groups: ", err)
				}
				for _, g := range groups {
					session.SendGroup(m.MessageID, g)
				}
				groups = nil
				//only administrative users need host access
				if session.GetAccessLevel() > 249 {
					hosts, err := dataStore.GetHosts()
					if err != nil {
						logger.Error("Error lookin up hosts: ", err)
					}
					for _, h := range hosts {
						session.SendHost(m.MessageID, h)
					}
				}

				//signal to the client that we have completed initial state sync
				session.SignalSuccess(m.MessageID, "State Sync Complete")

			default:
				logger.Error("Error on message from client: ", m.MessageType)

			}
		case <-clientClosed:
			//the client connection has closed
			exitLink <- session.GetGUID()
			return
		}

	}
}
<|endoftext|>


<commit_before>package goyaec_test

import (
	"fmt"
	"os"

	"github.com/utrack/goyaec"
)

// This example shows the basic usage.
//
// It shows how to parse basic structs and nested structs.
func ExampleProcess() {

	type BarType struct {
		C uint `env:"baz"`
	}

	type FooType struct {
		A string
		B int

		Bar BarType

		BazField string
	}

	// Wireup variables
	// You can use prefixes for structs
	os.Setenv("PREF_A", "VarA")
	os.Setenv("PREF_B", "2")

	// Use tag "env" to change the envvar name.
	// Previous prefixes are applied, so for BarType's C it becomes
	// pref (common prefix) -> bar (field name in FooType) -> baz (field name from tag).
	os.Setenv("PREF_BAR_BAZ", "2")

	os.Setenv("PREF_BAZFIELD","BazField")

	conf := FooType{}

	// Remember to pass the pointer!
	err := goyaec.Process("pref", &conf)

	if err != nil {
		panic(err)
	}

	fmt.Println(conf)
	// Output: {VarA 2 {2} BazField}
}

// This example shows the usage with map[string]string included
// in the struct.
func ExampleProcess_map() {
	type T struct {
		A bool `env:"somebool"`
		B map[string]string `env:"dict"`
	}

	// Wire the vars up
	os.Setenv("SOMEBOOL","true")
	// Case of the map keys is saved, capitalize as you want.
	os.Setenv("DICT_Foo","bar")
	os.Setenv("DICT_BaR","baz")
	os.Setenv("DICT_qux","whatever")

	t := T{}
	// Init the map and load some values in it - existing values
	// will be saved!
	t.B = map[string]string{"Foo":"prevFoo","SomeVar":"wow"}

	err := goyaec.Process("",&t)
	if err != nil {
		panic(err)
	}

	// Existing values are saved
	fmt.Printf("Key SomeVar: %v\n",t.B["SomeVar"])

	// ... but replaced if there's an envvar matching it
	fmt.Printf("Key Foo: %v\n",t.B["Foo"])

	// New values are loaded as they should.
	fmt.Printf("Key BaR: %v\n",t.B["BaR"])
	fmt.Printf("Key qux: %v\n",t.B["qux"])

	fmt.Printf("Member A: %v\n",t.A)
	// Output: Key SomeVar: wow
	// Key Foo: bar
	// Key BaR: baz
	// Key qux: whatever
	// Member A: true
}
<commit_msg>Made the basic example more basic<commit_after>package goyaec_test

import (
	"fmt"
	"os"

	"github.com/utrack/goyaec"
)

// This example shows the basic usage.
//
// It shows how to parse basic structs and use the tag overrides.
func ExampleProcess() {
	type Config struct {
		URI string
		Type string
		Active bool `env:"enabled"`
	}

	// Wireup variables
	// You can use prefixes for structs.
	os.Setenv("APP_URI", "blank://")
	os.Setenv("APP_TYPE", "Your Type")

	// Use tag "env" to change the envvar name.
	os.Setenv("APP_ENABLED", "true")

	conf := Config{}

	// Remember to pass the pointer!
	err := goyaec.Process("app", &conf)

	if err != nil {
		panic(err)
	}

	fmt.Println(conf)
	// Output: {blank:// Your Type true}
}

// This example shows the advanced usage.
//
// It shows how to parse basic structs and nested structs with env tags.
func ExampleProcess_advanced() {

	type BarType struct {
		C uint `env:"baz"`
	}

	type FooType struct {
		A string
		B int

		Bar BarType

		BazField string
	}

	// Wireup variables
	// You can use prefixes for structs
	os.Setenv("PREF_A", "VarA")
	os.Setenv("PREF_B", "2")

	// Use tag "env" to change the envvar name.
	// Previous prefixes are applied, so for BarType's C it becomes
	// pref (common prefix) -> bar (field name in FooType) -> baz (field name from tag).
	os.Setenv("PREF_BAR_BAZ", "2")

	os.Setenv("PREF_BAZFIELD","BazField")

	conf := FooType{}

	// Remember to pass the pointer!
	err := goyaec.Process("pref", &conf)

	if err != nil {
		panic(err)
	}

	fmt.Println(conf)
	// Output: {VarA 2 {2} BazField}
}

// This example shows the usage with map[string]string included
// in the struct.
func ExampleProcess_map() {
	type T struct {
		A bool `env:"somebool"`
		B map[string]string `env:"dict"`
	}

	// Wire the vars up
	os.Setenv("SOMEBOOL","true")
	// Case of the map keys is saved, capitalize as you want.
	os.Setenv("DICT_Foo","bar")
	os.Setenv("DICT_BaR","baz")
	os.Setenv("DICT_qux","whatever")

	t := T{}
	// Init the map and load some values in it - existing values
	// will be saved!
	t.B = map[string]string{"Foo":"prevFoo","SomeVar":"wow"}

	err := goyaec.Process("",&t)
	if err != nil {
		panic(err)
	}

	// Existing values are saved
	fmt.Printf("Key SomeVar: %v\n",t.B["SomeVar"])

	// ... but replaced if there's an envvar matching it
	fmt.Printf("Key Foo: %v\n",t.B["Foo"])

	// New values are loaded as they should.
	fmt.Printf("Key BaR: %v\n",t.B["BaR"])
	fmt.Printf("Key qux: %v\n",t.B["qux"])

	fmt.Printf("Member A: %v\n",t.A)
	// Output: Key SomeVar: wow
	// Key Foo: bar
	// Key BaR: baz
	// Key qux: whatever
	// Member A: true
}
<|endoftext|>


<commit_before>package cross_validation

import (
	"fmt"
	mat "github.com/gonum/matrix/mat64"
	"math/rand"
	"sync"
	"time"
)

func shuffleMatrix(returnDatasets []*mat.Dense, dataset mat.Matrix, testSize int, seed int64, wg *sync.WaitGroup) {
	numGen := rand.New(rand.NewSource(seed))

	// We don't want to alter the original dataset.
	shuffledSet := mat.DenseCopyOf(dataset)
	rowCount, colCount := shuffledSet.Dims()
	temp := make([]float64, colCount)

	// Fisher–Yates shuffle
	for i := 0; i < rowCount; i++ {
		j := numGen.Intn(i+1)
		if j != i {
			// Make a "hard" copy to avoid pointer craziness.
			copy(temp, shuffledSet.RowView(i))
			shuffledSet.SetRow(i, shuffledSet.RowView(j))
			shuffledSet.SetRow(j, temp)
		}
	}
	trainSize := rowCount - testSize
	returnDatasets[0] = mat.NewDense(trainSize, colCount, shuffledSet.RawMatrix().Data[:trainSize*colCount])
	returnDatasets[1] = mat.NewDense(testSize, colCount, shuffledSet.RawMatrix().Data[trainSize*colCount:])

	wg.Done()
}

// TrainTestSplit splits input DenseMatrix into subsets for testing.
// The function expects a test size number (int) or percentage (float64), and a random state or nil to get "random" shuffle.
// It returns a list containing the train-test split and an error status.
func TrainTestSplit(size interface{}, randomState interface{}, datasets ...*mat.Dense) ([]*mat.Dense, error) {
	// Get number of instances (rows).
	instanceCount, _ := datasets[0].Dims()

	// Input should be one or two matrices.
	dataCount := len(datasets)
	if dataCount > 2 {
		return nil, fmt.Errorf("Expected 1 or 2 datasets, got %d\n", dataCount)
	}

	if dataCount == 2 {
		// Test for consistency.
		labelCount, labelFeatures := datasets[1].Dims()
		if labelCount != instanceCount {
			return nil, fmt.Errorf("Data and labels must have the same number of instances")
		} else if labelFeatures != 1 {
			return nil, fmt.Errorf("Label matrix must have single feature")
		}
	}

	var testSize int
	switch size := size.(type) {
		// If size is an integer, treat it as the test data instance count.
	case int:
		testSize = size
	case float64:
		// If size is a float, treat it as a percentage of the instances to be allocated to the test set.
		testSize = int(float64(instanceCount)*size + 0.5)
	default:
		return nil, fmt.Errorf("Expected a test instance count (int) or percentage (float64)")
	}

	var randSeed int64
	// Create a deterministic shuffle, or a "random" one based on current time.
	if seed, ok := randomState.(int); ok {
		randSeed = int64(seed)
	} else {
		// Use seconds since epoch as seed
		randSeed = time.Now().Unix()
	}

	// Wait group for goroutine syncronization.
	wg := new(sync.WaitGroup)
	wg.Add(dataCount)

	// Return slice will hold training and test data and optional labels matrix.
	returnDatasets := make([]*mat.Dense, 2*dataCount)

	for i, dataset := range datasets {
		// Send proper returnDataset slice.
		// This is needed so goroutine doesn't mess up the expected return order.
		// Perhaps returning a map is a better solution...
		go shuffleMatrix(returnDatasets[i:i+2], dataset, testSize, randSeed, wg)
	}
	wg.Wait()

	return returnDatasets, nil
}
<commit_msg>error strings should not be capitalized<commit_after>package cross_validation

import (
	"fmt"
	mat "github.com/gonum/matrix/mat64"
	"math/rand"
	"sync"
	"time"
)

func shuffleMatrix(returnDatasets []*mat.Dense, dataset mat.Matrix, testSize int, seed int64, wg *sync.WaitGroup) {
	numGen := rand.New(rand.NewSource(seed))

	// We don't want to alter the original dataset.
	shuffledSet := mat.DenseCopyOf(dataset)
	rowCount, colCount := shuffledSet.Dims()
	temp := make([]float64, colCount)

	// Fisher–Yates shuffle
	for i := 0; i < rowCount; i++ {
		j := numGen.Intn(i + 1)
		if j != i {
			// Make a "hard" copy to avoid pointer craziness.
			copy(temp, shuffledSet.RowView(i))
			shuffledSet.SetRow(i, shuffledSet.RowView(j))
			shuffledSet.SetRow(j, temp)
		}
	}
	trainSize := rowCount - testSize
	returnDatasets[0] = mat.NewDense(trainSize, colCount, shuffledSet.RawMatrix().Data[:trainSize*colCount])
	returnDatasets[1] = mat.NewDense(testSize, colCount, shuffledSet.RawMatrix().Data[trainSize*colCount:])

	wg.Done()
}

// TrainTestSplit splits input DenseMatrix into subsets for testing.
// The function expects a test size number (int) or percentage (float64), and a random state or nil to get "random" shuffle.
// It returns a list containing the train-test split and an error status.
func TrainTestSplit(size interface{}, randomState interface{}, datasets ...*mat.Dense) ([]*mat.Dense, error) {
	// Get number of instances (rows).
	instanceCount, _ := datasets[0].Dims()

	// Input should be one or two matrices.
	dataCount := len(datasets)
	if dataCount > 2 {
		return nil, fmt.Errorf("expected 1 or 2 datasets, got %d\n", dataCount)
	}

	if dataCount == 2 {
		// Test for consistency.
		labelCount, labelFeatures := datasets[1].Dims()
		if labelCount != instanceCount {
			return nil, fmt.Errorf("data and labels must have the same number of instances")
		} else if labelFeatures != 1 {
			return nil, fmt.Errorf("label matrix must have single feature")
		}
	}

	var testSize int
	switch size := size.(type) {
	// If size is an integer, treat it as the test data instance count.
	case int:
		testSize = size
	case float64:
		// If size is a float, treat it as a percentage of the instances to be allocated to the test set.
		testSize = int(float64(instanceCount)*size + 0.5)
	default:
		return nil, fmt.Errorf("expected a test instance count (int) or percentage (float64)")
	}

	var randSeed int64
	// Create a deterministic shuffle, or a "random" one based on current time.
	if seed, ok := randomState.(int); ok {
		randSeed = int64(seed)
	} else {
		// Use seconds since epoch as seed
		randSeed = time.Now().Unix()
	}

	// Wait group for goroutine syncronization.
	wg := new(sync.WaitGroup)
	wg.Add(dataCount)

	// Return slice will hold training and test data and optional labels matrix.
	returnDatasets := make([]*mat.Dense, 2*dataCount)

	for i, dataset := range datasets {
		// Send proper returnDataset slice.
		// This is needed so goroutine doesn't mess up the expected return order.
		// Perhaps returning a map is a better solution...
		go shuffleMatrix(returnDatasets[i:i+2], dataset, testSize, randSeed, wg)
	}
	wg.Wait()

	return returnDatasets, nil
}
<|endoftext|>


<commit_before>package config

// EqualClusterReference returns whether two contexts appear to point to the
// same underlying kubernetes cluster
func (c *Context) EqualClusterReference(other *Context) bool {
	if other == nil {
		return false
	}
	if c.ClusterName != other.ClusterName {
		return false
	}
	if c.AuthInfo != other.AuthInfo {
		return false
	}
	if c.Namespace != other.Namespace {
		return false
	}
	return true
}
<commit_msg>Compare pachd address in contexts, as otherwise prior deployments with direct connect enabled will cause subsequent deployments without direct connect to fail<commit_after>package config

// EqualClusterReference returns whether two contexts appear to point to the
// same underlying kubernetes cluster
func (c *Context) EqualClusterReference(other *Context) bool {
	if other == nil {
		return false
	}
	if c.ClusterName != other.ClusterName {
		return false
	}
	if c.AuthInfo != other.AuthInfo {
		return false
	}
	if c.Namespace != other.Namespace {
		return false
	}
	if c.PachdAddress != other.PachdAddress {
		return false
	}
	return true
}
<|endoftext|>


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

package mailer

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/mail"
	"net/smtp"
	"os"
	"strings"
	"errors"
	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/setting"
)

type loginAuth struct {
  username, password string
}

//SMTP AUTH LOGIN Auth Handler
func LoginAuth(username, password string) smtp.Auth {
	return &loginAuth{username, password}
}

func (a *loginAuth) Start(server *smtp.ServerInfo) (string, []byte, error) {
	return "LOGIN", []byte{}, nil
}

func (a *loginAuth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		switch string(fromServer) {
		case "Username:":
			return []byte(a.username), nil
		case "Password:":
			return []byte(a.password), nil
		default:
			return nil, errors.New("Unkown fromServer")
		}
	}
	return nil, nil
}

type Message struct {
	To      []string
	From    string
	Subject string
	Body    string
	Type    string
	Massive bool
	Info    string
}

// create mail content
func (m Message) Content() string {
	// set mail type
	contentType := "text/plain; charset=UTF-8"
	if m.Type == "html" {
		contentType = "text/html; charset=UTF-8"
	}

	// create mail content
	content := "From: " + m.From + "\r\nSubject: " + m.Subject + "\r\nContent-Type: " + contentType + "\r\n\r\n" + m.Body
	return content
}

var mailQueue chan *Message

func NewMailerContext() {
	mailQueue = make(chan *Message, setting.Cfg.Section("mailer").Key("SEND_BUFFER_LEN").MustInt(10))
	go processMailQueue()
}

func processMailQueue() {
	for {
		select {
		case msg := <-mailQueue:
			num, err := Send(msg)
			tos := strings.Join(msg.To, "; ")
			info := ""
			if err != nil {
				if len(msg.Info) > 0 {
					info = ", info: " + msg.Info
				}
				log.Error(4, fmt.Sprintf("Async sent email %d succeed, not send emails: %s%s err: %s", num, tos, info, err))
			} else {
				log.Trace(fmt.Sprintf("Async sent email %d succeed, sent emails: %s%s", num, tos, info))
			}
		}
	}
}

// sendMail allows mail with self-signed certificates.
func sendMail(settings *setting.Mailer, recipients []string, msgContent []byte) error {
	host, port, err := net.SplitHostPort(settings.Host)
	if err != nil {
		return err
	}

	tlsconfig := &tls.Config{
		InsecureSkipVerify: settings.SkipVerify,
		ServerName:         host,
	}

	if settings.UseCertificate {
		cert, err := tls.LoadX509KeyPair(settings.CertFile, settings.KeyFile)
		if err != nil {
			return err
		}
		tlsconfig.Certificates = []tls.Certificate{cert}
	}

	conn, err := net.Dial("tcp", net.JoinHostPort(host, port))
	if err != nil {
		return err
	}
	defer conn.Close()

	isSecureConn := false
	// Start TLS directly if the port ends with 465 (SMTPS protocol)
	if strings.HasSuffix(port, "465") {
		conn = tls.Client(conn, tlsconfig)
		isSecureConn = true
	}

	client, err := smtp.NewClient(conn, host)
	if err != nil {
		return err
	}

	if !setting.MailService.DisableHelo {
		hostname := setting.MailService.HeloHostname
		if len(hostname) == 0 {
			hostname, err = os.Hostname()
			if err != nil {
				return err
			}
		}

		if err = client.Hello(hostname); err != nil {
			return err
		}
	}

	// If not using SMTPS, alway use STARTTLS if available
	hasStartTLS, _ := client.Extension("STARTTLS")
	if !isSecureConn && hasStartTLS {
		if err = client.StartTLS(tlsconfig); err != nil {
			return err
		}
	}

	canAuth, options := client.Extension("AUTH")

	if canAuth && len(settings.User) > 0 {
		var auth smtp.Auth

		if strings.Contains(options, "CRAM-MD5") {
			auth = smtp.CRAMMD5Auth(settings.User, settings.Passwd)
		} else if strings.Contains(options, "PLAIN") {
			auth = smtp.PlainAuth("", settings.User, settings.Passwd, host)
		//Patch for AUTH LOGIN
		} else if strings.Contains(options, "LOGIN") {
			auth = LoginAuth(settings.User, settings.Passwd)
		}

		if auth != nil {
			if err = client.Auth(auth); err != nil {
				return err
			}
		}
	}

	if fromAddress, err := mail.ParseAddress(settings.From); err != nil {
		return err
	} else {
		if err = client.Mail(fromAddress.Address); err != nil {
			return err
		}
	}

	for _, rec := range recipients {
		if err = client.Rcpt(rec); err != nil {
			return err
		}
	}

	w, err := client.Data()
	if err != nil {
		return err
	}
	if _, err = w.Write([]byte(msgContent)); err != nil {
		return err
	}

	if err = w.Close(); err != nil {
		return err
	}

	return client.Quit()
}

// Direct Send mail message
func Send(msg *Message) (int, error) {
	log.Trace("Sending mails to: %s", strings.Join(msg.To, "; "))

	// get message body
	content := msg.Content()

	if len(msg.To) == 0 {
		return 0, fmt.Errorf("empty receive emails")
	} else if len(msg.Body) == 0 {
		return 0, fmt.Errorf("empty email body")
	}

	if msg.Massive {
		// send mail to multiple emails one by one
		num := 0
		for _, to := range msg.To {
			body := []byte("To: " + to + "\r\n" + content)
			err := sendMail(setting.MailService, []string{to}, body)
			if err != nil {
				return num, err
			}
			num++
		}
		return num, nil
	} else {
		body := []byte("To: " + strings.Join(msg.To, ";") + "\r\n" + content)

		// send to multiple emails in one message
		err := sendMail(setting.MailService, msg.To, body)
		if err != nil {
			return 0, err
		} else {
			return 1, nil
		}
	}
}

// Async Send mail message
func SendAsync(msg *Message) {
	go func() {
		mailQueue <- msg
	}()
}

// Create html mail message
func NewHtmlMessage(To []string, From, Subject, Body string) Message {
	return Message{
		To:      To,
		From:    From,
		Subject: Subject,
		Body:    Body,
		Type:    "html",
	}
}
<commit_msg>minor fix on #1517<commit_after>// Copyright 2014 The Gogs Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package mailer

import (
	"crypto/tls"
	"fmt"
	"net"
	"net/mail"
	"net/smtp"
	"os"
	"strings"

	"github.com/gogits/gogs/modules/log"
	"github.com/gogits/gogs/modules/setting"
)

type loginAuth struct {
	username, password string
}

// SMTP AUTH LOGIN Auth Handler
func LoginAuth(username, password string) smtp.Auth {
	return &loginAuth{username, password}
}

func (a *loginAuth) Start(server *smtp.ServerInfo) (string, []byte, error) {
	return "LOGIN", []byte{}, nil
}

func (a *loginAuth) Next(fromServer []byte, more bool) ([]byte, error) {
	if more {
		switch string(fromServer) {
		case "Username:":
			return []byte(a.username), nil
		case "Password:":
			return []byte(a.password), nil
		default:
			return nil, fmt.Errorf("unknwon fromServer: %s", string(fromServer))
		}
	}
	return nil, nil
}

type Message struct {
	To      []string
	From    string
	Subject string
	Body    string
	Type    string
	Massive bool
	Info    string
}

// create mail content
func (m Message) Content() string {
	// set mail type
	contentType := "text/plain; charset=UTF-8"
	if m.Type == "html" {
		contentType = "text/html; charset=UTF-8"
	}

	// create mail content
	content := "From: " + m.From + "\r\nSubject: " + m.Subject + "\r\nContent-Type: " + contentType + "\r\n\r\n" + m.Body
	return content
}

var mailQueue chan *Message

func NewMailerContext() {
	mailQueue = make(chan *Message, setting.Cfg.Section("mailer").Key("SEND_BUFFER_LEN").MustInt(10))
	go processMailQueue()
}

func processMailQueue() {
	for {
		select {
		case msg := <-mailQueue:
			num, err := Send(msg)
			tos := strings.Join(msg.To, "; ")
			info := ""
			if err != nil {
				if len(msg.Info) > 0 {
					info = ", info: " + msg.Info
				}
				log.Error(4, fmt.Sprintf("Async sent email %d succeed, not send emails: %s%s err: %s", num, tos, info, err))
			} else {
				log.Trace(fmt.Sprintf("Async sent email %d succeed, sent emails: %s%s", num, tos, info))
			}
		}
	}
}

// sendMail allows mail with self-signed certificates.
func sendMail(settings *setting.Mailer, recipients []string, msgContent []byte) error {
	host, port, err := net.SplitHostPort(settings.Host)
	if err != nil {
		return err
	}

	tlsconfig := &tls.Config{
		InsecureSkipVerify: settings.SkipVerify,
		ServerName:         host,
	}

	if settings.UseCertificate {
		cert, err := tls.LoadX509KeyPair(settings.CertFile, settings.KeyFile)
		if err != nil {
			return err
		}
		tlsconfig.Certificates = []tls.Certificate{cert}
	}

	conn, err := net.Dial("tcp", net.JoinHostPort(host, port))
	if err != nil {
		return err
	}
	defer conn.Close()

	isSecureConn := false
	// Start TLS directly if the port ends with 465 (SMTPS protocol)
	if strings.HasSuffix(port, "465") {
		conn = tls.Client(conn, tlsconfig)
		isSecureConn = true
	}

	client, err := smtp.NewClient(conn, host)
	if err != nil {
		return err
	}

	if !setting.MailService.DisableHelo {
		hostname := setting.MailService.HeloHostname
		if len(hostname) == 0 {
			hostname, err = os.Hostname()
			if err != nil {
				return err
			}
		}

		if err = client.Hello(hostname); err != nil {
			return err
		}
	}

	// If not using SMTPS, alway use STARTTLS if available
	hasStartTLS, _ := client.Extension("STARTTLS")
	if !isSecureConn && hasStartTLS {
		if err = client.StartTLS(tlsconfig); err != nil {
			return err
		}
	}

	canAuth, options := client.Extension("AUTH")

	if canAuth && len(settings.User) > 0 {
		var auth smtp.Auth

		if strings.Contains(options, "CRAM-MD5") {
			auth = smtp.CRAMMD5Auth(settings.User, settings.Passwd)
		} else if strings.Contains(options, "PLAIN") {
			auth = smtp.PlainAuth("", settings.User, settings.Passwd, host)
		} else if strings.Contains(options, "LOGIN") {
			// Patch for AUTH LOGIN
			auth = LoginAuth(settings.User, settings.Passwd)
		}

		if auth != nil {
			if err = client.Auth(auth); err != nil {
				return err
			}
		}
	}

	if fromAddress, err := mail.ParseAddress(settings.From); err != nil {
		return err
	} else {
		if err = client.Mail(fromAddress.Address); err != nil {
			return err
		}
	}

	for _, rec := range recipients {
		if err = client.Rcpt(rec); err != nil {
			return err
		}
	}

	w, err := client.Data()
	if err != nil {
		return err
	}
	if _, err = w.Write([]byte(msgContent)); err != nil {
		return err
	}

	if err = w.Close(); err != nil {
		return err
	}

	return client.Quit()
}

// Direct Send mail message
func Send(msg *Message) (int, error) {
	log.Trace("Sending mails to: %s", strings.Join(msg.To, "; "))

	// get message body
	content := msg.Content()

	if len(msg.To) == 0 {
		return 0, fmt.Errorf("empty receive emails")
	} else if len(msg.Body) == 0 {
		return 0, fmt.Errorf("empty email body")
	}

	if msg.Massive {
		// send mail to multiple emails one by one
		num := 0
		for _, to := range msg.To {
			body := []byte("To: " + to + "\r\n" + content)
			err := sendMail(setting.MailService, []string{to}, body)
			if err != nil {
				return num, err
			}
			num++
		}
		return num, nil
	} else {
		body := []byte("To: " + strings.Join(msg.To, ";") + "\r\n" + content)

		// send to multiple emails in one message
		err := sendMail(setting.MailService, msg.To, body)
		if err != nil {
			return 0, err
		} else {
			return 1, nil
		}
	}
}

// Async Send mail message
func SendAsync(msg *Message) {
	go func() {
		mailQueue <- msg
	}()
}

// Create html mail message
func NewHtmlMessage(To []string, From, Subject, Body string) Message {
	return Message{
		To:      To,
		From:    From,
		Subject: Subject,
		Body:    Body,
		Type:    "html",
	}
}
<|endoftext|>


<commit_before>package renter

import (
	"crypto/rand"
	"errors"
	"os"
	"path/filepath"
	"time"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	maxUploadAttempts = 3
	parallelUploads   = 3
)

var (
	errUploadFailed = errors.New("failed to upload to the desired host")

	redundancy = 12
)

// checkWalletBalance looks at an upload and determines if there is enough
// money in the wallet to support such an upload. An error is returned if it is
// determined that there is not enough money.
func (r *Renter) checkWalletBalance(up modules.FileUploadParams) error {
	// Get the size of the file.
	fileInfo, err := os.Stat(up.Filename)
	if err != nil {
		return err
	}
	curSize := types.NewCurrency64(uint64(fileInfo.Size()))

	var averagePrice types.Currency
	sampleSize := redundancy * 2
	hosts := r.hostDB.RandomHosts(sampleSize)
	for _, host := range hosts {
		averagePrice = averagePrice.Add(host.Price)
	}
	averagePrice = averagePrice.Div(types.NewCurrency64(uint64(sampleSize)))
	estimatedCost := averagePrice.Mul(types.NewCurrency64(uint64(up.Duration))).Mul(curSize)
	bufferedCost := estimatedCost.Mul(types.NewCurrency64(2))

	if bufferedCost.Cmp(r.wallet.Balance(false)) > 0 {
		return errors.New("insufficient balance for upload")
	}
	return nil
}

// threadedUploadPiece will upload the piece of a file to a randomly chosen
// host. If the wallet has insufficient balance to support uploading,
// uploadPiece will give up. The file uploading can be continued using a repair
// tool. Upon completion, the memory containg the piece's information is
// updated.
func (r *Renter) threadedUploadPiece(host modules.HostSettings, up modules.FileUploadParams, piece *filePiece) error {
	// Set 'Repairing' for the piece to true.
	lockID := r.mu.Lock()
	piece.Repairing = true
	r.mu.Unlock(lockID)

	// Try 'maxUploadAttempts' hosts before giving up.
	for attempts := 0; attempts < maxUploadAttempts; attempts++ {
		// Negotiate the contract with the host. If the negotiation is
		// unsuccessful, we need to try again with a new host.
		err := r.negotiateContract(host, up, piece)
		if err == nil {
			lockID := r.mu.Lock()
			piece.Repairing = false
			r.mu.Unlock(lockID)
			return nil
		}

		// The previous attempt didn't work. We will try again after
		// sleeping for a randomized amount of time to increase our chances
		// of success. This will help spread things out if there are
		// problems with network congestion or other randomized issues.
		randSource := make([]byte, 1)
		rand.Read(randSource)
		time.Sleep(100 * time.Millisecond * time.Duration(randSource[0]))
	}

	// All attempts failed.
	return errors.New("failed to upload filePiece")
}

// Upload takes an upload parameters, which contain a file to upload, and then
// creates a redundant copy of the file on the Sia network.
func (r *Renter) Upload(up modules.FileUploadParams) error {
	// TODO: This type of restriction is something that should be handled by
	// the frontend, not the backend.
	if filepath.Ext(up.Filename) != filepath.Ext(up.Nickname) {
		return errors.New("nickname and file name must have the same extension")
	}

	err := r.checkWalletBalance(up)
	if err != nil {
		return err
	}

	// Check for a nickname conflict.
	lockID := r.mu.RLock()
	_, exists := r.files[up.Nickname]
	r.mu.RUnlock(lockID)
	if exists {
		return errors.New("file with that nickname already exists")
	}

	// Check that the file exists and is less than 500 MiB.
	fileInfo, err := os.Stat(up.Filename)
	if err != nil {
		return err
	}
	// NOTE: The upload max of 500 MiB is temporary and therefore does not have a
	// constant. This should be removed once micropayments + upload resuming
	// are in place. 500 MiB is chosen to prevent confusion - on anybody's
	// machine any file appearing to be under 500 MB will be below the hard
	// limit.
	if fileInfo.Size() > 500*1024*1024 {
		return errors.New("cannot upload a file larger than 500 MB")
	}

	// Check that the hostdb is sufficiently large to support an upload. Right
	// now that value is set to 3, but in the future the logic will be a bit
	// more complex; once there is erasure coding we'll want to hit the minimum
	// number of pieces plus some buffer before we decide that an upload is
	// okay.
	if len(r.hostDB.ActiveHosts()) < 1 {
		return errors.New("not enough hosts on the network to upload a file")
	}

	// Create file object.
	f := &file{
		Name: up.Nickname,

		PiecesRequired: 1,
		Pieces:         make([]filePiece, up.Pieces),
		UploadParams:   up,
		renter:         r,
	}
	for i := range f.Pieces {
		f.Pieces[i].Repairing = true
		f.Pieces[i].PieceSize = uint64(fileInfo.Size())
	}

	// Add file to renter.
	lockID = r.mu.Lock()
	r.files[up.Nickname] = f
	r.save()
	r.mu.Unlock(lockID)

	// Upload to hosts in parallel. To facilitate this, we create channels of
	// hosts and file pieces, and spawn goroutines that attempt to match each
	// piece to a host.
	hostPool := make(chan modules.HostSettings, 3*redundancy)
	for _, host := range r.hostDB.RandomHosts(3 * redundancy) {
		hostPool <- host
	}
	piecePool := make(chan *filePiece, len(f.Pieces))
	for i := range f.Pieces {
		piecePool <- &f.Pieces[i]
	}
	errChan := make(chan error, len(f.Pieces))
	for i := 0; i < parallelUploads; i++ {
		go func() {
			for piece := range piecePool {
				var err error
				for host := range hostPool {
					err = r.threadedUploadPiece(host, up, piece)
					if err == nil {
						break
					}
				}
				errChan <- err
			}
		}()
	}

	// Wait for success or failure. Since we are (currently) using full
	// replication, success means "one piece was uploaded," while failure
	// means "zero pieces were uploaded."
	reqPieces := f.PiecesRequired
	for i := 0; i < up.Pieces; i++ {
		if <-errChan == nil {
			reqPieces--
			if reqPieces <= 0 {
				return nil
			}
		}
	}

	// All uploads failed. Remove the file object.
	lockID = r.mu.Lock()
	delete(r.files, up.Nickname)
	r.save()
	r.mu.Unlock(lockID)

	return errors.New("failed to upload any file pieces")
}
<commit_msg>remove redundant piece update<commit_after>package renter

import (
	"crypto/rand"
	"errors"
	"os"
	"path/filepath"
	"time"

	"github.com/NebulousLabs/Sia/modules"
	"github.com/NebulousLabs/Sia/types"
)

const (
	maxUploadAttempts = 3
	parallelUploads   = 3
)

var (
	errUploadFailed = errors.New("failed to upload to the desired host")

	redundancy = 12
)

// checkWalletBalance looks at an upload and determines if there is enough
// money in the wallet to support such an upload. An error is returned if it is
// determined that there is not enough money.
func (r *Renter) checkWalletBalance(up modules.FileUploadParams) error {
	// Get the size of the file.
	fileInfo, err := os.Stat(up.Filename)
	if err != nil {
		return err
	}
	curSize := types.NewCurrency64(uint64(fileInfo.Size()))

	var averagePrice types.Currency
	sampleSize := redundancy * 2
	hosts := r.hostDB.RandomHosts(sampleSize)
	for _, host := range hosts {
		averagePrice = averagePrice.Add(host.Price)
	}
	averagePrice = averagePrice.Div(types.NewCurrency64(uint64(sampleSize)))
	estimatedCost := averagePrice.Mul(types.NewCurrency64(uint64(up.Duration))).Mul(curSize)
	bufferedCost := estimatedCost.Mul(types.NewCurrency64(2))

	if bufferedCost.Cmp(r.wallet.Balance(false)) > 0 {
		return errors.New("insufficient balance for upload")
	}
	return nil
}

// threadedUploadPiece will upload the piece of a file to a randomly chosen
// host. If the wallet has insufficient balance to support uploading,
// uploadPiece will give up. The file uploading can be continued using a repair
// tool. Upon completion, the memory containg the piece's information is
// updated.
func (r *Renter) threadedUploadPiece(host modules.HostSettings, up modules.FileUploadParams, piece *filePiece) error {
	// Set 'Repairing' for the piece to true.
	lockID := r.mu.Lock()
	piece.Repairing = true
	r.mu.Unlock(lockID)

	// Try 'maxUploadAttempts' hosts before giving up.
	for attempts := 0; attempts < maxUploadAttempts; attempts++ {
		// Negotiate the contract with the host. If the negotiation is
		// unsuccessful, we need to try again with a new host.
		err := r.negotiateContract(host, up, piece)
		if err == nil {
			return nil
		}

		// The previous attempt didn't work. We will try again after
		// sleeping for a randomized amount of time to increase our chances
		// of success. This will help spread things out if there are
		// problems with network congestion or other randomized issues.
		randSource := make([]byte, 1)
		rand.Read(randSource)
		time.Sleep(100 * time.Millisecond * time.Duration(randSource[0]))
	}

	// All attempts failed.
	return errors.New("failed to upload filePiece")
}

// Upload takes an upload parameters, which contain a file to upload, and then
// creates a redundant copy of the file on the Sia network.
func (r *Renter) Upload(up modules.FileUploadParams) error {
	// TODO: This type of restriction is something that should be handled by
	// the frontend, not the backend.
	if filepath.Ext(up.Filename) != filepath.Ext(up.Nickname) {
		return errors.New("nickname and file name must have the same extension")
	}

	err := r.checkWalletBalance(up)
	if err != nil {
		return err
	}

	// Check for a nickname conflict.
	lockID := r.mu.RLock()
	_, exists := r.files[up.Nickname]
	r.mu.RUnlock(lockID)
	if exists {
		return errors.New("file with that nickname already exists")
	}

	// Check that the file exists and is less than 500 MiB.
	fileInfo, err := os.Stat(up.Filename)
	if err != nil {
		return err
	}
	// NOTE: The upload max of 500 MiB is temporary and therefore does not have a
	// constant. This should be removed once micropayments + upload resuming
	// are in place. 500 MiB is chosen to prevent confusion - on anybody's
	// machine any file appearing to be under 500 MB will be below the hard
	// limit.
	if fileInfo.Size() > 500*1024*1024 {
		return errors.New("cannot upload a file larger than 500 MB")
	}

	// Check that the hostdb is sufficiently large to support an upload. Right
	// now that value is set to 3, but in the future the logic will be a bit
	// more complex; once there is erasure coding we'll want to hit the minimum
	// number of pieces plus some buffer before we decide that an upload is
	// okay.
	if len(r.hostDB.ActiveHosts()) < 1 {
		return errors.New("not enough hosts on the network to upload a file")
	}

	// Create file object.
	f := &file{
		Name: up.Nickname,

		PiecesRequired: 1,
		Pieces:         make([]filePiece, up.Pieces),
		UploadParams:   up,
		renter:         r,
	}
	for i := range f.Pieces {
		f.Pieces[i].Repairing = true
		f.Pieces[i].PieceSize = uint64(fileInfo.Size())
	}

	// Add file to renter.
	lockID = r.mu.Lock()
	r.files[up.Nickname] = f
	r.save()
	r.mu.Unlock(lockID)

	// Upload to hosts in parallel. To facilitate this, we create channels of
	// hosts and file pieces, and spawn goroutines that attempt to match each
	// piece to a host.
	hostPool := make(chan modules.HostSettings, 3*redundancy)
	for _, host := range r.hostDB.RandomHosts(3 * redundancy) {
		hostPool <- host
	}
	piecePool := make(chan *filePiece, len(f.Pieces))
	for i := range f.Pieces {
		piecePool <- &f.Pieces[i]
	}
	errChan := make(chan error, len(f.Pieces))
	for i := 0; i < parallelUploads; i++ {
		go func() {
			for piece := range piecePool {
				var err error
				for host := range hostPool {
					err = r.threadedUploadPiece(host, up, piece)
					if err == nil {
						break
					}
				}
				errChan <- err
			}
		}()
	}

	// Wait for success or failure. Since we are (currently) using full
	// replication, success means "one piece was uploaded," while failure
	// means "zero pieces were uploaded."
	reqPieces := f.PiecesRequired
	for i := 0; i < up.Pieces; i++ {
		if <-errChan == nil {
			reqPieces--
			if reqPieces <= 0 {
				return nil
			}
		}
	}

	// All uploads failed. Remove the file object.
	lockID = r.mu.Lock()
	delete(r.files, up.Nickname)
	r.save()
	r.mu.Unlock(lockID)

	return errors.New("failed to upload any file pieces")
}
<|endoftext|>


<commit_before>package bitswap

import (
	"bytes"
	"sync"
	"testing"
	"time"

	detectrace "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-detect-race"
	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
	travis "github.com/ipfs/go-ipfs/util/testutil/ci/travis"

	blocks "github.com/ipfs/go-ipfs/blocks"
	blocksutil "github.com/ipfs/go-ipfs/blocks/blocksutil"
	key "github.com/ipfs/go-ipfs/blocks/key"
	tn "github.com/ipfs/go-ipfs/exchange/bitswap/testnet"
	p2ptestutil "github.com/ipfs/go-ipfs/p2p/test/util"
	mockrouting "github.com/ipfs/go-ipfs/routing/mock"
	delay "github.com/ipfs/go-ipfs/thirdparty/delay"
)

// FIXME the tests are really sensitive to the network delay. fix them to work
// well under varying conditions
const kNetworkDelay = 0 * time.Millisecond

func TestClose(t *testing.T) {
	vnet := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sesgen := NewTestSessionGenerator(vnet)
	defer sesgen.Close()
	bgen := blocksutil.NewBlockGenerator()

	block := bgen.Next()
	bitswap := sesgen.Next()

	bitswap.Exchange.Close()
	bitswap.Exchange.GetBlock(context.Background(), block.Key())
}

func TestProviderForKeyButNetworkCannotFind(t *testing.T) { // TODO revisit this

	rs := mockrouting.NewServer()
	net := tn.VirtualNetwork(rs, delay.Fixed(kNetworkDelay))
	g := NewTestSessionGenerator(net)
	defer g.Close()

	block := blocks.NewBlock([]byte("block"))
	pinfo := p2ptestutil.RandTestBogusIdentityOrFatal(t)
	rs.Client(pinfo).Provide(context.Background(), block.Key()) // but not on network

	solo := g.Next()
	defer solo.Exchange.Close()

	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
	defer cancel()
	_, err := solo.Exchange.GetBlock(ctx, block.Key())

	if err != context.DeadlineExceeded {
		t.Fatal("Expected DeadlineExceeded error")
	}
}

func TestGetBlockFromPeerAfterPeerAnnounces(t *testing.T) {

	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	block := blocks.NewBlock([]byte("block"))
	g := NewTestSessionGenerator(net)
	defer g.Close()

	peers := g.Instances(2)
	hasBlock := peers[0]
	defer hasBlock.Exchange.Close()

	if err := hasBlock.Exchange.HasBlock(block); err != nil {
		t.Fatal(err)
	}

	wantsBlock := peers[1]
	defer wantsBlock.Exchange.Close()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	received, err := wantsBlock.Exchange.GetBlock(ctx, block.Key())
	if err != nil {
		t.Log(err)
		t.Fatal("Expected to succeed")
	}

	if !bytes.Equal(block.Data, received.Data) {
		t.Fatal("Data doesn't match")
	}
}

func TestLargeSwarm(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	numInstances := 100
	numBlocks := 2
	if detectrace.WithRace() {
		// when running with the race detector, 500 instances launches
		// well over 8k goroutines. This hits a race detector limit.
		numInstances = 100
	} else if travis.IsRunning() {
		numInstances = 200
	} else {
		t.Parallel()
	}
	PerformDistributionTest(t, numInstances, numBlocks)
}

func TestLargeFile(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	if !travis.IsRunning() {
		t.Parallel()
	}

	numInstances := 10
	numBlocks := 100
	PerformDistributionTest(t, numInstances, numBlocks)
}

func TestLargeFileNoRebroadcast(t *testing.T) {
	rbd := rebroadcastDelay.Get()
	rebroadcastDelay.Set(time.Hour * 24 * 365 * 10) // ten years should be long enough
	if testing.Short() {
		t.SkipNow()
	}
	numInstances := 10
	numBlocks := 100
	PerformDistributionTest(t, numInstances, numBlocks)
	rebroadcastDelay.Set(rbd)
}

func TestLargeFileTwoPeers(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	numInstances := 2
	numBlocks := 100
	PerformDistributionTest(t, numInstances, numBlocks)
}

func PerformDistributionTest(t *testing.T, numInstances, numBlocks int) {
	ctx := context.Background()
	if testing.Short() {
		t.SkipNow()
	}
	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sg := NewTestSessionGenerator(net)
	defer sg.Close()
	bg := blocksutil.NewBlockGenerator()

	instances := sg.Instances(numInstances)
	blocks := bg.Blocks(numBlocks)

	t.Log("Give the blocks to the first instance")

	var blkeys []key.Key
	first := instances[0]
	for _, b := range blocks {
		blkeys = append(blkeys, b.Key())
		first.Exchange.HasBlock(b)
	}

	t.Log("Distribute!")

	wg := sync.WaitGroup{}
	for _, inst := range instances[1:] {
		wg.Add(1)
		go func(inst Instance) {
			defer wg.Done()
			outch, err := inst.Exchange.GetBlocks(ctx, blkeys)
			if err != nil {
				t.Fatal(err)
			}
			for _ = range outch {
			}
		}(inst)
	}
	wg.Wait()

	t.Log("Verify!")

	for _, inst := range instances {
		for _, b := range blocks {
			if _, err := inst.Blockstore().Get(b.Key()); err != nil {
				t.Fatal(err)
			}
		}
	}
}

func getOrFail(bitswap Instance, b *blocks.Block, t *testing.T, wg *sync.WaitGroup) {
	if _, err := bitswap.Blockstore().Get(b.Key()); err != nil {
		_, err := bitswap.Exchange.GetBlock(context.Background(), b.Key())
		if err != nil {
			t.Fatal(err)
		}
	}
	wg.Done()
}

// TODO simplify this test. get to the _essence_!
func TestSendToWantingPeer(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sg := NewTestSessionGenerator(net)
	defer sg.Close()
	bg := blocksutil.NewBlockGenerator()

	prev := rebroadcastDelay.Set(time.Second / 2)
	defer func() { rebroadcastDelay.Set(prev) }()

	peers := sg.Instances(2)
	peerA := peers[0]
	peerB := peers[1]

	t.Logf("Session %v\n", peerA.Peer)
	t.Logf("Session %v\n", peerB.Peer)

	waitTime := time.Second * 5

	alpha := bg.Next()
	// peerA requests and waits for block alpha
	ctx, cancel := context.WithTimeout(context.Background(), waitTime)
	defer cancel()
	alphaPromise, err := peerA.Exchange.GetBlocks(ctx, []key.Key{alpha.Key()})
	if err != nil {
		t.Fatal(err)
	}

	// peerB announces to the network that he has block alpha
	err = peerB.Exchange.HasBlock(alpha)
	if err != nil {
		t.Fatal(err)
	}

	// At some point, peerA should get alpha (or timeout)
	blkrecvd, ok := <-alphaPromise
	if !ok {
		t.Fatal("context timed out and broke promise channel!")
	}

	if blkrecvd.Key() != alpha.Key() {
		t.Fatal("Wrong block!")
	}

}

func TestBasicBitswap(t *testing.T) {
	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sg := NewTestSessionGenerator(net)
	defer sg.Close()
	bg := blocksutil.NewBlockGenerator()

	t.Log("Test a one node trying to get one block from another")

	instances := sg.Instances(2)
	blocks := bg.Blocks(1)
	err := instances[0].Exchange.HasBlock(blocks[0])
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	blk, err := instances[1].Exchange.GetBlock(ctx, blocks[0].Key())
	if err != nil {
		t.Fatal(err)
	}

	t.Log(blk)
	for _, inst := range instances {
		err := inst.Exchange.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}
<commit_msg>exchange/bitswap/bitswap_test: fix t.Fatal in a goroutine<commit_after>package bitswap

import (
	"bytes"
	"sync"
	"testing"
	"time"

	detectrace "github.com/ipfs/go-ipfs/Godeps/_workspace/src/github.com/jbenet/go-detect-race"
	context "github.com/ipfs/go-ipfs/Godeps/_workspace/src/golang.org/x/net/context"
	travis "github.com/ipfs/go-ipfs/util/testutil/ci/travis"

	blocks "github.com/ipfs/go-ipfs/blocks"
	blocksutil "github.com/ipfs/go-ipfs/blocks/blocksutil"
	key "github.com/ipfs/go-ipfs/blocks/key"
	tn "github.com/ipfs/go-ipfs/exchange/bitswap/testnet"
	p2ptestutil "github.com/ipfs/go-ipfs/p2p/test/util"
	mockrouting "github.com/ipfs/go-ipfs/routing/mock"
	delay "github.com/ipfs/go-ipfs/thirdparty/delay"
)

// FIXME the tests are really sensitive to the network delay. fix them to work
// well under varying conditions
const kNetworkDelay = 0 * time.Millisecond

func TestClose(t *testing.T) {
	vnet := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sesgen := NewTestSessionGenerator(vnet)
	defer sesgen.Close()
	bgen := blocksutil.NewBlockGenerator()

	block := bgen.Next()
	bitswap := sesgen.Next()

	bitswap.Exchange.Close()
	bitswap.Exchange.GetBlock(context.Background(), block.Key())
}

func TestProviderForKeyButNetworkCannotFind(t *testing.T) { // TODO revisit this

	rs := mockrouting.NewServer()
	net := tn.VirtualNetwork(rs, delay.Fixed(kNetworkDelay))
	g := NewTestSessionGenerator(net)
	defer g.Close()

	block := blocks.NewBlock([]byte("block"))
	pinfo := p2ptestutil.RandTestBogusIdentityOrFatal(t)
	rs.Client(pinfo).Provide(context.Background(), block.Key()) // but not on network

	solo := g.Next()
	defer solo.Exchange.Close()

	ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
	defer cancel()
	_, err := solo.Exchange.GetBlock(ctx, block.Key())

	if err != context.DeadlineExceeded {
		t.Fatal("Expected DeadlineExceeded error")
	}
}

func TestGetBlockFromPeerAfterPeerAnnounces(t *testing.T) {

	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	block := blocks.NewBlock([]byte("block"))
	g := NewTestSessionGenerator(net)
	defer g.Close()

	peers := g.Instances(2)
	hasBlock := peers[0]
	defer hasBlock.Exchange.Close()

	if err := hasBlock.Exchange.HasBlock(block); err != nil {
		t.Fatal(err)
	}

	wantsBlock := peers[1]
	defer wantsBlock.Exchange.Close()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	received, err := wantsBlock.Exchange.GetBlock(ctx, block.Key())
	if err != nil {
		t.Log(err)
		t.Fatal("Expected to succeed")
	}

	if !bytes.Equal(block.Data, received.Data) {
		t.Fatal("Data doesn't match")
	}
}

func TestLargeSwarm(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	numInstances := 100
	numBlocks := 2
	if detectrace.WithRace() {
		// when running with the race detector, 500 instances launches
		// well over 8k goroutines. This hits a race detector limit.
		numInstances = 100
	} else if travis.IsRunning() {
		numInstances = 200
	} else {
		t.Parallel()
	}
	PerformDistributionTest(t, numInstances, numBlocks)
}

func TestLargeFile(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	if !travis.IsRunning() {
		t.Parallel()
	}

	numInstances := 10
	numBlocks := 100
	PerformDistributionTest(t, numInstances, numBlocks)
}

func TestLargeFileNoRebroadcast(t *testing.T) {
	rbd := rebroadcastDelay.Get()
	rebroadcastDelay.Set(time.Hour * 24 * 365 * 10) // ten years should be long enough
	if testing.Short() {
		t.SkipNow()
	}
	numInstances := 10
	numBlocks := 100
	PerformDistributionTest(t, numInstances, numBlocks)
	rebroadcastDelay.Set(rbd)
}

func TestLargeFileTwoPeers(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}
	numInstances := 2
	numBlocks := 100
	PerformDistributionTest(t, numInstances, numBlocks)
}

func PerformDistributionTest(t *testing.T, numInstances, numBlocks int) {
	ctx := context.Background()
	if testing.Short() {
		t.SkipNow()
	}
	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sg := NewTestSessionGenerator(net)
	defer sg.Close()
	bg := blocksutil.NewBlockGenerator()

	instances := sg.Instances(numInstances)
	blocks := bg.Blocks(numBlocks)

	t.Log("Give the blocks to the first instance")

	var blkeys []key.Key
	first := instances[0]
	for _, b := range blocks {
		blkeys = append(blkeys, b.Key())
		first.Exchange.HasBlock(b)
	}

	t.Log("Distribute!")

	wg := sync.WaitGroup{}
	errs := make(chan error)

	for _, inst := range instances[1:] {
		wg.Add(1)
		go func(inst Instance) {
			defer wg.Done()
			outch, err := inst.Exchange.GetBlocks(ctx, blkeys)
			if err != nil {
				errs <- err
			}
			for _ = range outch {
			}
		}(inst)
	}

	go func() {
		wg.Wait()
		close(errs)
	}()

	for err := range errs {
		if err != nil {
			t.Fatal(err)
		}
	}

	t.Log("Verify!")

	for _, inst := range instances {
		for _, b := range blocks {
			if _, err := inst.Blockstore().Get(b.Key()); err != nil {
				t.Fatal(err)
			}
		}
	}
}

func getOrFail(bitswap Instance, b *blocks.Block, t *testing.T, wg *sync.WaitGroup) {
	if _, err := bitswap.Blockstore().Get(b.Key()); err != nil {
		_, err := bitswap.Exchange.GetBlock(context.Background(), b.Key())
		if err != nil {
			t.Fatal(err)
		}
	}
	wg.Done()
}

// TODO simplify this test. get to the _essence_!
func TestSendToWantingPeer(t *testing.T) {
	if testing.Short() {
		t.SkipNow()
	}

	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sg := NewTestSessionGenerator(net)
	defer sg.Close()
	bg := blocksutil.NewBlockGenerator()

	prev := rebroadcastDelay.Set(time.Second / 2)
	defer func() { rebroadcastDelay.Set(prev) }()

	peers := sg.Instances(2)
	peerA := peers[0]
	peerB := peers[1]

	t.Logf("Session %v\n", peerA.Peer)
	t.Logf("Session %v\n", peerB.Peer)

	waitTime := time.Second * 5

	alpha := bg.Next()
	// peerA requests and waits for block alpha
	ctx, cancel := context.WithTimeout(context.Background(), waitTime)
	defer cancel()
	alphaPromise, err := peerA.Exchange.GetBlocks(ctx, []key.Key{alpha.Key()})
	if err != nil {
		t.Fatal(err)
	}

	// peerB announces to the network that he has block alpha
	err = peerB.Exchange.HasBlock(alpha)
	if err != nil {
		t.Fatal(err)
	}

	// At some point, peerA should get alpha (or timeout)
	blkrecvd, ok := <-alphaPromise
	if !ok {
		t.Fatal("context timed out and broke promise channel!")
	}

	if blkrecvd.Key() != alpha.Key() {
		t.Fatal("Wrong block!")
	}

}

func TestBasicBitswap(t *testing.T) {
	net := tn.VirtualNetwork(mockrouting.NewServer(), delay.Fixed(kNetworkDelay))
	sg := NewTestSessionGenerator(net)
	defer sg.Close()
	bg := blocksutil.NewBlockGenerator()

	t.Log("Test a one node trying to get one block from another")

	instances := sg.Instances(2)
	blocks := bg.Blocks(1)
	err := instances[0].Exchange.HasBlock(blocks[0])
	if err != nil {
		t.Fatal(err)
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	blk, err := instances[1].Exchange.GetBlock(ctx, blocks[0].Key())
	if err != nil {
		t.Fatal(err)
	}

	t.Log(blk)
	for _, inst := range instances {
		err := inst.Exchange.Close()
		if err != nil {
			t.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>package mysql

import (
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"fmt"
	"io/ioutil"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	log "github.com/mgutz/logxi/v1"

	"github.com/armon/go-metrics"
	mysql "github.com/go-sql-driver/mysql"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/vault/helper/strutil"
	"github.com/hashicorp/vault/physical"
)

// Unreserved tls key
// Reserved values are "true", "false", "skip-verify"
const mysqlTLSKey = "default"

// MySQLBackend is a physical backend that stores data
// within MySQL database.
type MySQLBackend struct {
	dbTable    string
	client     *sql.DB
	statements map[string]*sql.Stmt
	logger     log.Logger
	permitPool *physical.PermitPool
}

// NewMySQLBackend constructs a MySQL backend using the given API client and
// server address and credential for accessing mysql database.
func NewMySQLBackend(conf map[string]string, logger log.Logger) (physical.Backend, error) {
	var err error

	// Get the MySQL credentials to perform read/write operations.
	username, ok := conf["username"]
	if !ok || username == "" {
		return nil, fmt.Errorf("missing username")
	}
	password, ok := conf["password"]
	if !ok || username == "" {
		return nil, fmt.Errorf("missing password")
	}

	// Get or set MySQL server address. Defaults to localhost and default port(3306)
	address, ok := conf["address"]
	if !ok {
		address = "127.0.0.1:3306"
	}

	// Get the MySQL database and table details.
	database, ok := conf["database"]
	if !ok {
		database = "vault"
	}
	table, ok := conf["table"]
	if !ok {
		table = "vault"
	}
	dbTable := database + "." + table

	maxParStr, ok := conf["max_parallel"]
	var maxParInt int
	if ok {
		maxParInt, err = strconv.Atoi(maxParStr)
		if err != nil {
			return nil, errwrap.Wrapf("failed parsing max_parallel parameter: {{err}}", err)
		}
		if logger.IsDebug() {
			logger.Debug("mysql: max_parallel set", "max_parallel", maxParInt)
		}
	} else {
		maxParInt = physical.DefaultParallelOperations
	}

	dsnParams := url.Values{}
	tlsCaFile, ok := conf["tls_ca_file"]
	if ok {
		if err := setupMySQLTLSConfig(tlsCaFile); err != nil {
			return nil, fmt.Errorf("failed register TLS config: %v", err)
		}

		dsnParams.Add("tls", mysqlTLSKey)
	}

	// Create MySQL handle for the database.
	dsn := username + ":" + password + "@tcp(" + address + ")/?" + dsnParams.Encode()
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to mysql: %v", err)
	}

	db.SetMaxOpenConns(maxParInt)

	// Create the required database if it doesn't exists.
	if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS " + database); err != nil {
		return nil, fmt.Errorf("failed to create mysql database: %v", err)
	}

	// Create the required table if it doesn't exists.
	create_query := "CREATE TABLE IF NOT EXISTS " + dbTable +
		" (vault_key varbinary(512), vault_value mediumblob, PRIMARY KEY (vault_key))"
	if _, err := db.Exec(create_query); err != nil {
		return nil, fmt.Errorf("failed to create mysql table: %v", err)
	}

	// Setup the backend.
	m := &MySQLBackend{
		dbTable:    dbTable,
		client:     db,
		statements: make(map[string]*sql.Stmt),
		logger:     logger,
		permitPool: physical.NewPermitPool(maxParInt),
	}

	// Prepare all the statements required
	statements := map[string]string{
		"put": "INSERT INTO " + dbTable +
			" VALUES( ?, ? ) ON DUPLICATE KEY UPDATE vault_value=VALUES(vault_value)",
		"get":    "SELECT vault_value FROM " + dbTable + " WHERE vault_key = ?",
		"delete": "DELETE FROM " + dbTable + " WHERE vault_key = ?",
		"list":   "SELECT vault_key FROM " + dbTable + " WHERE vault_key LIKE ?",
	}
	for name, query := range statements {
		if err := m.prepare(name, query); err != nil {
			return nil, err
		}
	}

	return m, nil
}

// prepare is a helper to prepare a query for future execution
func (m *MySQLBackend) prepare(name, query string) error {
	stmt, err := m.client.Prepare(query)
	if err != nil {
		return fmt.Errorf("failed to prepare '%s': %v", name, err)
	}
	m.statements[name] = stmt
	return nil
}

// Put is used to insert or update an entry.
func (m *MySQLBackend) Put(entry *physical.Entry) error {
	defer metrics.MeasureSince([]string{"mysql", "put"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	_, err := m.statements["put"].Exec(entry.Key, entry.Value)
	if err != nil {
		return err
	}
	return nil
}

// Get is used to fetch and entry.
func (m *MySQLBackend) Get(key string) (*physical.Entry, error) {
	defer metrics.MeasureSince([]string{"mysql", "get"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	var result []byte
	err := m.statements["get"].QueryRow(key).Scan(&result)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	ent := &physical.Entry{
		Key:   key,
		Value: result,
	}
	return ent, nil
}

// Delete is used to permanently delete an entry
func (m *MySQLBackend) Delete(key string) error {
	defer metrics.MeasureSince([]string{"mysql", "delete"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	_, err := m.statements["delete"].Exec(key)
	if err != nil {
		return err
	}
	return nil
}

// List is used to list all the keys under a given
// prefix, up to the next prefix.
func (m *MySQLBackend) List(prefix string) ([]string, error) {
	defer metrics.MeasureSince([]string{"mysql", "list"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	// Add the % wildcard to the prefix to do the prefix search
	likePrefix := prefix + "%"
	rows, err := m.statements["list"].Query(likePrefix)
	if err != nil {
		return nil, fmt.Errorf("failed to execute statement: %v", err)
	}

	var keys []string
	for rows.Next() {
		var key string
		err = rows.Scan(&key)
		if err != nil {
			return nil, fmt.Errorf("failed to scan rows: %v", err)
		}

		key = strings.TrimPrefix(key, prefix)
		if i := strings.Index(key, "/"); i == -1 {
			// Add objects only from the current 'folder'
			keys = append(keys, key)
		} else if i != -1 {
			// Add truncated 'folder' paths
			keys = strutil.AppendIfMissing(keys, string(key[:i+1]))
		}
	}

	sort.Strings(keys)
	return keys, nil
}

// Establish a TLS connection with a given CA certificate
// Register a tsl.Config associted with the same key as the dns param from sql.Open
// foo:bar@tcp(127.0.0.1:3306)/dbname?tls=default
func setupMySQLTLSConfig(tlsCaFile string) error {
	rootCertPool := x509.NewCertPool()

	pem, err := ioutil.ReadFile(tlsCaFile)
	if err != nil {
		return err
	}

	if ok := rootCertPool.AppendCertsFromPEM(pem); !ok {
		return err
	}

	err = mysql.RegisterTLSConfig(mysqlTLSKey, &tls.Config{
		RootCAs: rootCertPool,
	})
	if err != nil {
		return err
	}

	return nil
}
<commit_msg>check schema and table before create it (#3716)<commit_after>package mysql

import (
	"crypto/tls"
	"crypto/x509"
	"database/sql"
	"fmt"
	"io/ioutil"
	"net/url"
	"sort"
	"strconv"
	"strings"
	"time"

	log "github.com/mgutz/logxi/v1"

	"github.com/armon/go-metrics"
	mysql "github.com/go-sql-driver/mysql"
	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/vault/helper/strutil"
	"github.com/hashicorp/vault/physical"
)

// Unreserved tls key
// Reserved values are "true", "false", "skip-verify"
const mysqlTLSKey = "default"

// MySQLBackend is a physical backend that stores data
// within MySQL database.
type MySQLBackend struct {
	dbTable    string
	client     *sql.DB
	statements map[string]*sql.Stmt
	logger     log.Logger
	permitPool *physical.PermitPool
}

// NewMySQLBackend constructs a MySQL backend using the given API client and
// server address and credential for accessing mysql database.
func NewMySQLBackend(conf map[string]string, logger log.Logger) (physical.Backend, error) {
	var err error

	// Get the MySQL credentials to perform read/write operations.
	username, ok := conf["username"]
	if !ok || username == "" {
		return nil, fmt.Errorf("missing username")
	}
	password, ok := conf["password"]
	if !ok || username == "" {
		return nil, fmt.Errorf("missing password")
	}

	// Get or set MySQL server address. Defaults to localhost and default port(3306)
	address, ok := conf["address"]
	if !ok {
		address = "127.0.0.1:3306"
	}

	// Get the MySQL database and table details.
	database, ok := conf["database"]
	if !ok {
		database = "vault"
	}
	table, ok := conf["table"]
	if !ok {
		table = "vault"
	}
	dbTable := database + "." + table

	maxParStr, ok := conf["max_parallel"]
	var maxParInt int
	if ok {
		maxParInt, err = strconv.Atoi(maxParStr)
		if err != nil {
			return nil, errwrap.Wrapf("failed parsing max_parallel parameter: {{err}}", err)
		}
		if logger.IsDebug() {
			logger.Debug("mysql: max_parallel set", "max_parallel", maxParInt)
		}
	} else {
		maxParInt = physical.DefaultParallelOperations
	}

	dsnParams := url.Values{}
	tlsCaFile, ok := conf["tls_ca_file"]
	if ok {
		if err := setupMySQLTLSConfig(tlsCaFile); err != nil {
			return nil, fmt.Errorf("failed register TLS config: %v", err)
		}

		dsnParams.Add("tls", mysqlTLSKey)
	}

	// Create MySQL handle for the database.
	dsn := username + ":" + password + "@tcp(" + address + ")/?" + dsnParams.Encode()
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to mysql: %v", err)
	}

	db.SetMaxOpenConns(maxParInt)

	// Check schema exists
	var schemaExist bool
	schemaRows, err := db.Query("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = ?", database)
	if err != nil {
		return nil, fmt.Errorf("failed to check mysql schema exist: %v", err)
	}
	defer schemaRows.Close()
	schemaExist = schemaRows.Next()

	// Check table exists
	var tableExist bool
	tableRows, err := db.Query("SELECT TABLE_NAME FROM information_schema.TABLES WHERE TABLE_NAME = ? AND TABLE_SCHEMA = ?", table, database)

	if err != nil {
		return nil, fmt.Errorf("failed to check mysql table exist: %v", err)
	}
	defer tableRows.Close()
	tableExist = tableRows.Next()

	// Create the required database if it doesn't exists.
	if !schemaExist {
		if _, err := db.Exec("CREATE DATABASE IF NOT EXISTS " + database); err != nil {
			return nil, fmt.Errorf("failed to create mysql database: %v", err)
		}
	}

	// Create the required table if it doesn't exists.
	if !tableExist {
		create_query := "CREATE TABLE IF NOT EXISTS " + dbTable +
			" (vault_key varbinary(512), vault_value mediumblob, PRIMARY KEY (vault_key))"
		if _, err := db.Exec(create_query); err != nil {
			return nil, fmt.Errorf("failed to create mysql table: %v", err)
		}
	}

	// Setup the backend.
	m := &MySQLBackend{
		dbTable:    dbTable,
		client:     db,
		statements: make(map[string]*sql.Stmt),
		logger:     logger,
		permitPool: physical.NewPermitPool(maxParInt),
	}

	// Prepare all the statements required
	statements := map[string]string{
		"put": "INSERT INTO " + dbTable +
			" VALUES( ?, ? ) ON DUPLICATE KEY UPDATE vault_value=VALUES(vault_value)",
		"get":    "SELECT vault_value FROM " + dbTable + " WHERE vault_key = ?",
		"delete": "DELETE FROM " + dbTable + " WHERE vault_key = ?",
		"list":   "SELECT vault_key FROM " + dbTable + " WHERE vault_key LIKE ?",
	}
	for name, query := range statements {
		if err := m.prepare(name, query); err != nil {
			return nil, err
		}
	}

	return m, nil
}

// prepare is a helper to prepare a query for future execution
func (m *MySQLBackend) prepare(name, query string) error {
	stmt, err := m.client.Prepare(query)
	if err != nil {
		return fmt.Errorf("failed to prepare '%s': %v", name, err)
	}
	m.statements[name] = stmt
	return nil
}

// Put is used to insert or update an entry.
func (m *MySQLBackend) Put(entry *physical.Entry) error {
	defer metrics.MeasureSince([]string{"mysql", "put"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	_, err := m.statements["put"].Exec(entry.Key, entry.Value)
	if err != nil {
		return err
	}
	return nil
}

// Get is used to fetch and entry.
func (m *MySQLBackend) Get(key string) (*physical.Entry, error) {
	defer metrics.MeasureSince([]string{"mysql", "get"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	var result []byte
	err := m.statements["get"].QueryRow(key).Scan(&result)
	if err == sql.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	ent := &physical.Entry{
		Key:   key,
		Value: result,
	}
	return ent, nil
}

// Delete is used to permanently delete an entry
func (m *MySQLBackend) Delete(key string) error {
	defer metrics.MeasureSince([]string{"mysql", "delete"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	_, err := m.statements["delete"].Exec(key)
	if err != nil {
		return err
	}
	return nil
}

// List is used to list all the keys under a given
// prefix, up to the next prefix.
func (m *MySQLBackend) List(prefix string) ([]string, error) {
	defer metrics.MeasureSince([]string{"mysql", "list"}, time.Now())

	m.permitPool.Acquire()
	defer m.permitPool.Release()

	// Add the % wildcard to the prefix to do the prefix search
	likePrefix := prefix + "%"
	rows, err := m.statements["list"].Query(likePrefix)
	if err != nil {
		return nil, fmt.Errorf("failed to execute statement: %v", err)
	}

	var keys []string
	for rows.Next() {
		var key string
		err = rows.Scan(&key)
		if err != nil {
			return nil, fmt.Errorf("failed to scan rows: %v", err)
		}

		key = strings.TrimPrefix(key, prefix)
		if i := strings.Index(key, "/"); i == -1 {
			// Add objects only from the current 'folder'
			keys = append(keys, key)
		} else if i != -1 {
			// Add truncated 'folder' paths
			keys = strutil.AppendIfMissing(keys, string(key[:i+1]))
		}
	}

	sort.Strings(keys)
	return keys, nil
}

// Establish a TLS connection with a given CA certificate
// Register a tsl.Config associted with the same key as the dns param from sql.Open
// foo:bar@tcp(127.0.0.1:3306)/dbname?tls=default
func setupMySQLTLSConfig(tlsCaFile string) error {
	rootCertPool := x509.NewCertPool()

	pem, err := ioutil.ReadFile(tlsCaFile)
	if err != nil {
		return err
	}

	if ok := rootCertPool.AppendCertsFromPEM(pem); !ok {
		return err
	}

	err = mysql.RegisterTLSConfig(mysqlTLSKey, &tls.Config{
		RootCAs: rootCertPool,
	})
	if err != nil {
		return err
	}

	return nil
}
<|endoftext|>