<commit_before>package main_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"testing"
)

func TestMain(m *testing.M) {
	os.MkdirAll("tmp", os.ModeDir|os.ModePerm)
	err := exec.Command("go", "build", "-o", "tmp/sqlfmt").Run()
	if err != nil {
		fmt.Println("Failed to build sqlfmt binary:", err)
		os.Exit(1)
	}

	os.Exit(m.Run())
}

func sqlfmt(t *testing.T, sql string, args ...string) string {
	cmd := exec.Command("tmp/sqlfmt", args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		t.Fatalf("cmd.StdinPipe failed: %v", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		t.Fatalf("cmd.StdoutPipe failed: %v", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		t.Fatalf("cmd.StderrPipe failed: %v", err)
	}

	err = cmd.Start()
	if err != nil {
		t.Fatalf("cmd.Start failed: %v", err)
	}

	_, err = fmt.Fprint(stdin, sql)
	if err != nil {
		t.Fatalf("fmt.Fprint failed: %v", err)
	}

	err = stdin.Close()
	if err != nil {
		t.Fatalf("stdin.Close failed: %v", err)
	}

	output, err := ioutil.ReadAll(stdout)
	if err != nil {
		t.Fatalf("ioutil.ReadAll(stdout) failed: %v", err)
	}

	errout, err := ioutil.ReadAll(stderr)
	if err != nil {
		t.Fatalf("ioutil.ReadAll(stderr) failed: %v", err)
	}

	err = cmd.Wait()
	if err != nil {
		t.Fatalf("cmd.Wait failed: %v\n%s", err, string(errout))
	}

	return string(output)
}

func TestSqlFmt(t *testing.T) {
	tests := []struct {
		inputFile          string
		expectedOutputFile string
	}{
		{
			inputFile:          "simple_select_without_from.sql",
			expectedOutputFile: "simple_select_without_from.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_from.sql",
			expectedOutputFile: "simple_select_with_from.fmt.sql",
		},
		{
			inputFile:          "select_from_aliased.sql",
			expectedOutputFile: "select_from_aliased.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_selection_alias.sql",
			expectedOutputFile: "simple_select_with_selection_alias.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_selection_alias_no_as.sql",
			expectedOutputFile: "simple_select_with_selection_alias.fmt.sql",
		},
		{
			inputFile:          "select_table_dot_column.sql",
			expectedOutputFile: "select_table_dot_column.fmt.sql",
		},
		{
			inputFile:          "simple_select_literal_integer.sql",
			expectedOutputFile: "simple_select_literal_integer.fmt.sql",
		},
		{
			inputFile:          "simple_select_literal_text.sql",
			expectedOutputFile: "simple_select_literal_text.fmt.sql",
		},
		{
			inputFile:          "arithmetic_expression.sql",
			expectedOutputFile: "arithmetic_expression.fmt.sql",
		},
		{
			inputFile:          "paren_expression.sql",
			expectedOutputFile: "paren_expression.fmt.sql",
		},
		{
			inputFile:          "subselect_expression.sql",
			expectedOutputFile: "subselect_expression.fmt.sql",
		},
		{
			inputFile:          "comparison_expression.sql",
			expectedOutputFile: "comparison_expression.fmt.sql",
		},
		{
			inputFile:          "select_from_comma_join.sql",
			expectedOutputFile: "select_from_comma_join.fmt.sql",
		},
		{
			inputFile:          "select_from_cross_join.sql",
			expectedOutputFile: "select_from_cross_join.fmt.sql",
		},
		{
			inputFile:          "select_from_natural_join.sql",
			expectedOutputFile: "select_from_natural_join.fmt.sql",
		},
		{
			inputFile:          "select_from_join_using.sql",
			expectedOutputFile: "select_from_join_using.fmt.sql",
		},
		{
			inputFile:          "select_from_join_using_multiple.sql",
			expectedOutputFile: "select_from_join_using_multiple.fmt.sql",
		},
		{
			inputFile:          "select_from_join_on.sql",
			expectedOutputFile: "select_from_join_on.fmt.sql",
		},
		{
			inputFile:          "quoted_identifier.sql",
			expectedOutputFile: "quoted_identifier.fmt.sql",
		},
		{
			inputFile:          "boolean_binary_op.sql",
			expectedOutputFile: "boolean_binary_op.fmt.sql",
		},
		{
			inputFile:          "boolean_not.sql",
			expectedOutputFile: "boolean_not.fmt.sql",
		},
		{
			inputFile:          "select_where.sql",
			expectedOutputFile: "select_where.fmt.sql",
		},
		{
			inputFile:          "order.sql",
			expectedOutputFile: "order.fmt.sql",
		},
		{
			inputFile:          "order_column_num.sql",
			expectedOutputFile: "order_column_num.fmt.sql",
		},
		{
			inputFile:          "order_desc.sql",
			expectedOutputFile: "order_desc.fmt.sql",
		},
		{
			inputFile:          "order_multiple.sql",
			expectedOutputFile: "order_multiple.fmt.sql",
		},
	}

	for i, tt := range tests {
		input, err := ioutil.ReadFile(path.Join("testdata", tt.inputFile))
		if err != nil {
			t.Fatal(err)
		}

		expected, err := ioutil.ReadFile(path.Join("testdata", tt.expectedOutputFile))
		if err != nil {
			t.Fatal(err)
		}

		output := sqlfmt(t, string(input))

		if output != string(expected) {
			actualFileName := path.Join("tmp", fmt.Sprintf("%d.sql", i))
			err = ioutil.WriteFile(actualFileName, []byte(output), os.ModePerm)
			if err != nil {
				t.Fatal(err)
			}

			t.Errorf("%d. Given %s, did not receive %s. Unexpected output written to %s", i, tt.inputFile, tt.expectedOutputFile, actualFileName)
		}
	}
}
<commit_msg>Continue testing even if sqlfmt execution fails<commit_after>package main_test

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"testing"
)

func TestMain(m *testing.M) {
	os.MkdirAll("tmp", os.ModeDir|os.ModePerm)
	err := exec.Command("go", "build", "-o", "tmp/sqlfmt").Run()
	if err != nil {
		fmt.Println("Failed to build sqlfmt binary:", err)
		os.Exit(1)
	}

	os.Exit(m.Run())
}

func sqlfmt(sql string, args ...string) (string, error) {
	cmd := exec.Command("tmp/sqlfmt", args...)
	stdin, err := cmd.StdinPipe()
	if err != nil {
		return "", fmt.Errorf("cmd.StdinPipe failed: %v", err)
	}
	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return "", fmt.Errorf("cmd.StdoutPipe failed: %v", err)
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return "", fmt.Errorf("cmd.StderrPipe failed: %v", err)
	}

	err = cmd.Start()
	if err != nil {
		return "", fmt.Errorf("cmd.Start failed: %v", err)
	}

	_, err = fmt.Fprint(stdin, sql)
	if err != nil {
		return "", fmt.Errorf("fmt.Fprint failed: %v", err)
	}

	err = stdin.Close()
	if err != nil {
		return "", fmt.Errorf("stdin.Close failed: %v", err)
	}

	output, err := ioutil.ReadAll(stdout)
	if err != nil {
		return "", fmt.Errorf("ioutil.ReadAll(stdout) failed: %v", err)
	}

	errout, err := ioutil.ReadAll(stderr)
	if err != nil {
		return "", fmt.Errorf("ioutil.ReadAll(stderr) failed: %v", err)
	}

	err = cmd.Wait()
	if err != nil {
		return "", fmt.Errorf("cmd.Wait failed: %v\n%s", err, string(errout))
	}

	return string(output), nil
}

func TestSqlFmt(t *testing.T) {
	tests := []struct {
		inputFile          string
		expectedOutputFile string
	}{
		{
			inputFile:          "simple_select_without_from.sql",
			expectedOutputFile: "simple_select_without_from.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_from.sql",
			expectedOutputFile: "simple_select_with_from.fmt.sql",
		},
		{
			inputFile:          "select_from_aliased.sql",
			expectedOutputFile: "select_from_aliased.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_selection_alias.sql",
			expectedOutputFile: "simple_select_with_selection_alias.fmt.sql",
		},
		{
			inputFile:          "simple_select_with_selection_alias_no_as.sql",
			expectedOutputFile: "simple_select_with_selection_alias.fmt.sql",
		},
		{
			inputFile:          "select_table_dot_column.sql",
			expectedOutputFile: "select_table_dot_column.fmt.sql",
		},
		{
			inputFile:          "simple_select_literal_integer.sql",
			expectedOutputFile: "simple_select_literal_integer.fmt.sql",
		},
		{
			inputFile:          "simple_select_literal_text.sql",
			expectedOutputFile: "simple_select_literal_text.fmt.sql",
		},
		{
			inputFile:          "arithmetic_expression.sql",
			expectedOutputFile: "arithmetic_expression.fmt.sql",
		},
		{
			inputFile:          "paren_expression.sql",
			expectedOutputFile: "paren_expression.fmt.sql",
		},
		{
			inputFile:          "subselect_expression.sql",
			expectedOutputFile: "subselect_expression.fmt.sql",
		},
		{
			inputFile:          "comparison_expression.sql",
			expectedOutputFile: "comparison_expression.fmt.sql",
		},
		{
			inputFile:          "select_from_comma_join.sql",
			expectedOutputFile: "select_from_comma_join.fmt.sql",
		},
		{
			inputFile:          "select_from_cross_join.sql",
			expectedOutputFile: "select_from_cross_join.fmt.sql",
		},
		{
			inputFile:          "select_from_natural_join.sql",
			expectedOutputFile: "select_from_natural_join.fmt.sql",
		},
		{
			inputFile:          "select_from_join_using.sql",
			expectedOutputFile: "select_from_join_using.fmt.sql",
		},
		{
			inputFile:          "select_from_join_using_multiple.sql",
			expectedOutputFile: "select_from_join_using_multiple.fmt.sql",
		},
		{
			inputFile:          "select_from_join_on.sql",
			expectedOutputFile: "select_from_join_on.fmt.sql",
		},
		{
			inputFile:          "quoted_identifier.sql",
			expectedOutputFile: "quoted_identifier.fmt.sql",
		},
		{
			inputFile:          "boolean_binary_op.sql",
			expectedOutputFile: "boolean_binary_op.fmt.sql",
		},
		{
			inputFile:          "boolean_not.sql",
			expectedOutputFile: "boolean_not.fmt.sql",
		},
		{
			inputFile:          "select_where.sql",
			expectedOutputFile: "select_where.fmt.sql",
		},
		{
			inputFile:          "order.sql",
			expectedOutputFile: "order.fmt.sql",
		},
		{
			inputFile:          "order_column_num.sql",
			expectedOutputFile: "order_column_num.fmt.sql",
		},
		{
			inputFile:          "order_desc.sql",
			expectedOutputFile: "order_desc.fmt.sql",
		},
		{
			inputFile:          "order_multiple.sql",
			expectedOutputFile: "order_multiple.fmt.sql",
		},
	}

	for i, tt := range tests {
		input, err := ioutil.ReadFile(path.Join("testdata", tt.inputFile))
		if err != nil {
			t.Errorf("%d. %v", i, err)
			continue
		}

		expected, err := ioutil.ReadFile(path.Join("testdata", tt.expectedOutputFile))
		if err != nil {
			t.Errorf("%d. %v", i, err)
			continue
		}

		output, err := sqlfmt(string(input))
		if err != nil {
			t.Errorf("%d. sqfmt failed: %v", i, err)
			continue
		}

		if output != string(expected) {
			actualFileName := path.Join("tmp", fmt.Sprintf("%d.sql", i))
			err = ioutil.WriteFile(actualFileName, []byte(output), os.ModePerm)
			if err != nil {
				t.Fatal(err)
			}

			t.Errorf("%d. Given %s, did not receive %s. Unexpected output written to %s", i, tt.inputFile, tt.expectedOutputFile, actualFileName)
		}
	}
}
<|endoftext|>


<commit_before>package elkrem

import (
	"testing"

	"github.com/btcsuite/btcd/wire"
)

// TestElkremBig tries 10K hashes
func TestElkremBig(t *testing.T) {
	sndr := NewElkremSender(wire.DoubleSha256SH([]byte("elktest")))
	var rcv ElkremReceiver
	//	SenderSerdesTest(t, sndr)
	for n := uint64(0); n < 10000; n++ {
		sha, err := sndr.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		err = rcv.AddNext(sha)
		if err != nil {
			t.Fatal(err)
		}
		if n%1000 == 999 {
			t.Logf("stack with %d received hashes\n", n+1)
			for i, n := range rcv.s {
				t.Logf("Stack element %d: index %d height %d %s\n",
					i, n.i, n.h, n.sha.String())
			}
		}
	}
	//	SenderSerdesTest(t, sndr)
	ReceiverSerdesTest(t, &rcv)
	for n := uint64(0); n < 10000; n += 500 {
		sha, err := rcv.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("Retreived index %d %s\n", n, sha.String())
	}
}

// TestElkremLess tries 10K hashes
func TestElkremLess(t *testing.T) {
	sndr := NewElkremSender(wire.DoubleSha256SH([]byte("elktest2")))
	var rcv ElkremReceiver
	for n := uint64(0); n < 5000; n++ {
		sha, err := sndr.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		err = rcv.AddNext(sha)
		if err != nil {
			t.Fatal(err)
		}
		if n%1000 == 999 {
			t.Logf("stack with %d received hashes\n", n+1)
			for i, n := range rcv.s {
				t.Logf("Stack element %d: index %d height %d %s\n",
					i, n.i, n.h, n.sha.String())
			}
		}
	}
	for n := uint64(0); n < 5000; n += 500 {
		sha, err := rcv.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("Retreived index %d %s\n",
			n, sha.String())
	}
}
<commit_msg>move elkrem tests to chainhash<commit_after>package elkrem

import (
	"testing"

	"github.com/btcsuite/btcd/chaincfg/chainhash"
)

// TestElkremBig tries 10K hashes
func TestElkremBig(t *testing.T) {
	sndr := NewElkremSender(chainhash.DoubleHashH([]byte("elktest")))
	var rcv ElkremReceiver
	//	SenderSerdesTest(t, sndr)
	for n := uint64(0); n < 10000; n++ {
		sha, err := sndr.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		err = rcv.AddNext(sha)
		if err != nil {
			t.Fatal(err)
		}
		if n%1000 == 999 {
			t.Logf("stack with %d received hashes\n", n+1)
			for i, n := range rcv.s {
				t.Logf("Stack element %d: index %d height %d %s\n",
					i, n.i, n.h, n.sha.String())
			}
		}
	}
	//	SenderSerdesTest(t, sndr)
	ReceiverSerdesTest(t, &rcv)
	for n := uint64(0); n < 10000; n += 500 {
		sha, err := rcv.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("Retreived index %d %s\n", n, sha.String())
	}
}

// TestElkremLess tries 10K hashes
func TestElkremLess(t *testing.T) {
	sndr := NewElkremSender(chainhash.DoubleHashH([]byte("elktest2")))
	var rcv ElkremReceiver
	for n := uint64(0); n < 5000; n++ {
		sha, err := sndr.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		err = rcv.AddNext(sha)
		if err != nil {
			t.Fatal(err)
		}
		if n%1000 == 999 {
			t.Logf("stack with %d received hashes\n", n+1)
			for i, n := range rcv.s {
				t.Logf("Stack element %d: index %d height %d %s\n",
					i, n.i, n.h, n.sha.String())
			}
		}
	}
	for n := uint64(0); n < 5000; n += 500 {
		sha, err := rcv.AtIndex(n)
		if err != nil {
			t.Fatal(err)
		}
		t.Logf("Retreived index %d %s\n",
			n, sha.String())
	}
}
<|endoftext|>


<commit_before>package operators

import (
	"math/rand"
	"testing"

	"github.com/aurelien-rainone/evolve/framework"
	"github.com/aurelien-rainone/evolve/number"
	"github.com/stretchr/testify/assert"
)

func TestStringCrossover(t *testing.T) {
	rng := rand.New(rand.NewSource(99))

	crossover, err := NewStringCrossover()
	if assert.NoError(t, err) {
		population := make([]framework.Candidate, 4)
		population[0] = "abcde"
		population[1] = "fghij"
		population[2] = "klmno"
		population[3] = "pqrst"

		for i := 0; i < 20; i++ {
			values := make(map[rune]struct{}, 20) // used as a set of runes
			population = crossover.Apply(population, rng)
			assert.Len(t, population, 4, "Population size changed after cross-over.")
			for _, individual := range population {
				s := individual.(string)
				assert.Lenf(t, s, 5, "Invalid candidate length: %v", len(s))
				for _, value := range s {
					values[value] = struct{}{}
				}
			}
			// All of the individual elements should still be present, just jumbled up
			// between individuals.
			assert.Len(t, values, 20, "Information lost during cross-over.")
		}
	}
}

// The StringCrossover operator is only defined to work on populations
// containing strings of equal lengths. Any attempt to apply the operation to
// populations that contain different length strings should panic. Not panicking
// should be considered a bug since it could lead to hard to trace bugs
// elsewhere.
func TestStringCrossoverWithDifferentLengthParents(t *testing.T) {
	rng := rand.New(rand.NewSource(99))

	crossover, err := NewStringCrossover(
		WithConstantCrossoverPoints(1),
		WithConstantCrossoverProbability(number.ProbabilityOne),
	)
	if assert.NoError(t, err) {
		population := make([]framework.Candidate, 2)
		population[0] = "abcde"
		population[1] = "fghijklm"

		// This should panic since the parents are different lengths.
		// TODO: why panicking and not returning an error?
		assert.Panics(t, func() {
			crossover.Apply(population, rng)
		})
	}
}

func TestStringCrossoverZeroPoints(t *testing.T) {
	rng := rand.New(rand.NewSource(99))

	t.Run("constant_crossover_points_cant_be_zero", func(t *testing.T) {
		// If created with a specified (constant) number of crossover points,
		// this number must be greater than 0 or the operator is a no-op.
		op, err := NewStringCrossover(WithConstantCrossoverPoints(0))
		assert.Error(t, err)
		assert.Nilf(t, op, "want string crossover to be nil if invalid, got %v", op)
	})

	t.Run("zero_crossover_points_is_noop", func(t *testing.T) {
		// If created with a variable number of crossover points,
		// verifies that when this number happens to be 0, the operator is a
		// no-op.
		crossover, err := NewStringCrossover(WithVariableCrossoverPoints(zeroGenerator{}))
		if assert.NoError(t, err) {
			population := []framework.Candidate{"abcde", "fghij"}
			crossed := crossover.Apply([]framework.Candidate{population[0], population[1]}, rng)
			assert.Equal(t, crossed, population)
		}
	})
}

type zeroGenerator struct{}

func (g zeroGenerator) NextValue() int64 {
	return 0
}
<commit_msg>Add test for crossover with zero proability<commit_after>package operators

import (
	"math/rand"
	"testing"

	"github.com/aurelien-rainone/evolve/framework"
	"github.com/aurelien-rainone/evolve/number"
	"github.com/stretchr/testify/assert"
)

func TestStringCrossover(t *testing.T) {
	rng := rand.New(rand.NewSource(99))

	crossover, err := NewStringCrossover()
	if assert.NoError(t, err) {
		population := make([]framework.Candidate, 4)
		population[0] = "abcde"
		population[1] = "fghij"
		population[2] = "klmno"
		population[3] = "pqrst"

		for i := 0; i < 20; i++ {
			values := make(map[rune]struct{}, 20) // used as a set of runes
			population = crossover.Apply(population, rng)
			assert.Len(t, population, 4, "Population size changed after cross-over.")
			for _, individual := range population {
				s := individual.(string)
				assert.Lenf(t, s, 5, "Invalid candidate length: %v", len(s))
				for _, value := range s {
					values[value] = struct{}{}
				}
			}
			// All of the individual elements should still be present, just jumbled up
			// between individuals.
			assert.Len(t, values, 20, "Information lost during cross-over.")
		}
	}
}

// The StringCrossover operator is only defined to work on populations
// containing strings of equal lengths. Any attempt to apply the operation to
// populations that contain different length strings should panic. Not panicking
// should be considered a bug since it could lead to hard to trace bugs
// elsewhere.
func TestStringCrossoverWithDifferentLengthParents(t *testing.T) {
	rng := rand.New(rand.NewSource(99))

	crossover, err := NewStringCrossover(
		WithConstantCrossoverPoints(1),
		WithConstantCrossoverProbability(number.ProbabilityOne),
	)
	if assert.NoError(t, err) {
		population := []framework.Candidate{"abcde", "fghijklm"}

		// This should panic since the parents are different lengths.
		// TODO: why panicking and not returning an error?
		assert.Panics(t, func() {
			crossover.Apply(population, rng)
		})
	}
}

func TestStringCrossoverNoop(t *testing.T) {
	rng := rand.New(rand.NewSource(99))

	t.Run("constant_crossover_points_cant_be_zero", func(t *testing.T) {
		// If created with a specified (constant) number of crossover points,
		// this number must be greater than 0 or the operator is a no-op.
		op, err := NewStringCrossover(WithConstantCrossoverPoints(0))
		assert.Error(t, err)
		assert.Nilf(t, op, "want string crossover to be nil if invalid, got %v", op)
	})

	t.Run("zero_crossover_points_is_noop", func(t *testing.T) {
		// If created with a variable number of crossover points,
		// verifies that when this number happens to be 0, the operator is a
		// no-op.
		crossover, err := NewStringCrossover(WithVariableCrossoverPoints(zeroGenerator{}))
		if assert.NoError(t, err) {
			population := []framework.Candidate{"abcde", "fghij"}
			crossed := crossover.Apply([]framework.Candidate{population[0], population[1]}, rng)
			assert.Equal(t, crossed, population)
		}
	})

	t.Run("zero_crossover_probability_is_noop", func(t *testing.T) {
		// If created wit a variable number of crossover probability,
		// verifies that when this number happens to be 0, the operator is a
		// no-op.
		crossover, err := NewStringCrossover(WithConstantCrossoverProbability(number.ProbabilityZero))
		if assert.NoError(t, err) {
			population := []framework.Candidate{"abcde", "fghij"}
			crossed := crossover.Apply([]framework.Candidate{population[0], population[1]}, rng)
			assert.Equal(t, crossed, population)
		}
	})
}

type zeroGenerator struct{}

func (g zeroGenerator) NextValue() int64 {
	return 0
}
<|endoftext|>


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

package binaryprefix

import (
	"testing"

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

func TestMB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix MB to number of megabytes", t, func() {
		Convey("Given the string 1MB, the number of megabytes is 1", func() {
			str := "1MB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1)
		})
		Convey("Given the string 1mb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1mb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestGB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix GB to number of megabytes", t, func() {
		Convey("Given the string 1GB, the number of megabytes is 1024", func() {
			str := "1GB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1024)
		})
		Convey("Given the string 1gb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1gb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestTB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix TB to number of megabytes", t, func() {
		Convey("Given the string 1TB, the number of megabytes is 1048576", func() {
			str := "1TB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1048576)
		})
		Convey("Given the string 1tb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1tb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestPB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix PB to number of megabytes", t, func() {
		Convey("Given the string 1PB, the number of megabytes is 1073741824", func() {
			str := "1PB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1073741824)
		})
		Convey("Given the string 1pb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1pb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestEB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix EB to number of megabytes", t, func() {
		Convey("Given the string 1EB, the number of megabytes is 1099511627776", func() {
			str := "1EB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1099511627776)
		})
		Convey("Given the string 1eb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1eb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestZB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix ZB to number of megabytes", t, func() {
		Convey("Given the string 1ZB, the number of megabytes is 1125899906842624", func() {
			str := "1ZB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1125899906842624)
		})
		Convey("Given the string 1zb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1zb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestYB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix YB to number of megabytes", t, func() {
		Convey("Given the string 1YB, the number of megabytes is 1152921504606846976", func() {
			str := "1YB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1152921504606846976)
		})
		Convey("Given the string 1yb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1yb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}
<commit_msg>Test invalid format<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package binaryprefix

import (
	"testing"

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

func TestInvalidFormat(t *testing.T) {
	Convey("Scenario: Convert string with an invalid format to number of megabytes", t, func() {
		Convey("Given the string 'foo', fails as it is an invalid format", func() {
			str := "foo"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "strconv.ParseFloat: parsing \"f\": invalid syntax")
		})
	})
}

func TestMB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix MB to number of megabytes", t, func() {
		Convey("Given the string 1MB, the number of megabytes is 1", func() {
			str := "1MB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1)
		})
		Convey("Given the string 1mb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1mb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestGB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix GB to number of megabytes", t, func() {
		Convey("Given the string 1GB, the number of megabytes is 1024", func() {
			str := "1GB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1024)
		})
		Convey("Given the string 1gb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1gb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestTB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix TB to number of megabytes", t, func() {
		Convey("Given the string 1TB, the number of megabytes is 1048576", func() {
			str := "1TB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1048576)
		})
		Convey("Given the string 1tb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1tb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestPB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix PB to number of megabytes", t, func() {
		Convey("Given the string 1PB, the number of megabytes is 1073741824", func() {
			str := "1PB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1073741824)
		})
		Convey("Given the string 1pb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1pb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestEB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix EB to number of megabytes", t, func() {
		Convey("Given the string 1EB, the number of megabytes is 1099511627776", func() {
			str := "1EB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1099511627776)
		})
		Convey("Given the string 1eb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1eb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestZB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix ZB to number of megabytes", t, func() {
		Convey("Given the string 1ZB, the number of megabytes is 1125899906842624", func() {
			str := "1ZB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1125899906842624)
		})
		Convey("Given the string 1zb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1zb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}

func TestYB(t *testing.T) {
	Convey("Scenario: Convert string with the prefix YB to number of megabytes", t, func() {
		Convey("Given the string 1YB, the number of megabytes is 1152921504606846976", func() {
			str := "1YB"
			mb, _ := GetMB(str)
			So(mb, ShouldEqual, 1152921504606846976)
		})
		Convey("Given the string 1yb, the number of megabytes is 0 as lower case denominations are not supported", func() {
			str := "1yb"
			mb, err := GetMB(str)
			So(mb, ShouldEqual, 0)
			So(err.Error(), ShouldEqual, "Unknown Denomination")
		})
	})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"time"
	"errors"
)

import "stathat.com/c/consistent"

// BUFFERSIZE controls the size of the [...]byte array used to read UDP data
// off the wire and into local memory.  Metrics are separated by \n
// characters.  This buffer is passed to a handler to proxy out the metrics
// to the real statsd daemons.
const BUFFERSIZE int = 1 * 1024 * 1024 // 1MiB

// packetLen is the size in bytes of data we stuff into one packet before
// sending it to statsd.  This must be lower than the MTU, IPv4 header size
// and UDP header size.
const packetLen int = 1000

// prefix is the string that will be prefixed onto self generated stats.
// Such as <prefix>.statsProcessed.  Default is "statsrelay"
var prefix string

// udpAddr is a mapping of HOST:PORT:INSTANCE to a UDPAddr object
var udpAddr = make(map[string]*net.UDPAddr)

// hashRing is our consistent hashing ring.
var hashRing *consistent.Consistent

// totalMetrics tracks the totall number of metrics processed
var totalMetrics int = 0

// totalMetricsLock is a mutex gaurding totalMetrics
var totalMetricsLock sync.Mutex

// Time we began
var epochTime int64

// Verbose/Debug output
var verbose bool

// sockBufferMaxSize() returns the maximum size that the UDP receive buffer
// in the kernel can be set to.  In bytes.
func sockBufferMaxSize() int {

	// XXX: This is Linux-only most likely
	data, err := ioutil.ReadFile("/proc/sys/net/core/rmem_max")
	if err != nil {
		log.Panicln(err)
	}

	data = bytes.TrimRight(data, "\n\r")
	i, err := strconv.Atoi(string(data))
	if err != nil {
		log.Printf("Could not parse /proc/sys/net/core/rmem_max\n")
		log.Fatalln(err)
	}

	return i
}

// getMetricName() parses the given []byte metric as a string, extracts
// the metric key name and returns it as a string.
func getMetricName(metric []byte) (string,error) {
	// statsd metrics are of the form:
	//    KEY:VALUE|TYPE|RATE
	length := bytes.IndexByte(metric, byte(':'))
	if (length == -1) {
		return "error", errors.New("Length of -1, must be invalid StatsD data")
	}
	return string(metric[:length]), nil
}

// sendPacket takes a []byte and writes that directly to a UDP socket
// that was assigned for target.
func sendPacket(buff []byte, target string) {
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		log.Panicln(err)
	}

	conn.WriteToUDP(buff, udpAddr[target])
	conn.Close()
}

// buildPacketMap() is a helper function to initiallize a map that represents
// a UDP packet currently being built for each destination we proxy to.  As
// Go forbids taking the address of an object in a map or array so the
// bytes.Buffer object must be stored in the map as a pointer rather than
// a direct object in order to call the pointer methods on it.
func buildPacketMap() map[string]*bytes.Buffer {
	members := hashRing.Members()
	hash := make(map[string]*bytes.Buffer, len(members))

	for _, v := range members {
		hash[v] = new(bytes.Buffer)
	}

	return hash
}

// handleBuff() sorts through a full buffer of metrics and batches metrics
// to remote statsd daemons using a consistent hash.
func handleBuff(buff []byte) {
	packets := buildPacketMap()
	sep := []byte("\n")
	numMetrics := 0
	statsMetric := prefix + ".statsProcessed"

	for offset := 0; offset < len(buff); {
	loop:
		for offset < len(buff) {
			// Find our next value
			switch buff[offset] {
			case '\n':
				offset++
			case '\r':
				offset++
			case 0:
				offset++
			default:
				break loop
			}
		}

		size := bytes.IndexByte(buff[offset:], '\n')

		if size == -1 {
			// last metric in buffer
			size = len(buff) - offset
		}
		if size == 0 {
			// no more metrics
			break
		}

        //Check to ensure we get a metric, and not an invalid Byte sequence
		metric, err := getMetricName(buff[offset : offset+size])

        if (err == nil) {

        	target, err := hashRing.Get(metric)
        	if err != nil {
        		log.Panicln(err)
        	}

			// check built packet size and send if metric doesn't fit
        	if packets[target].Len()+size > packetLen {
        		sendPacket(packets[target].Bytes(), target)
        		packets[target].Reset()
        	}
			// add to packet
        	packets[target].Write(buff[offset : offset+size])
        	packets[target].Write(sep)

        	numMetrics++
	    }

		offset = offset + size + 1
	}

	if numMetrics == 0 {
		// if we haven't handled any metrics, then don't update counters/stats
		// or send packets
		return
	}

	// Update interal counter
	totalMetricsLock.Lock()
	totalMetrics = totalMetrics + numMetrics
	totalMetricsLock.Unlock()

	// Handle reporting our own stats
	stats := fmt.Sprintf("%s:%d|c\n", statsMetric, numMetrics)
	target, err := hashRing.Get(statsMetric)
	if err != nil {
		log.Panicln(err)
	}
	if packets[target].Len()+len(stats) > packetLen {
		sendPacket(packets[target].Bytes(), target)
		packets[target].Reset()
	}
	packets[target].Write([]byte(stats))

	// Empty out any remaining data
	for _, target := range hashRing.Members() {
		if packets[target].Len() > 0 {
			sendPacket(packets[target].Bytes(), target)
		}
	}

	if verbose {
		log.Printf("Procssed %d metrics. Running total: %d. Metrics/sec: %d\n",
			numMetrics, totalMetrics,
			int64(totalMetrics)/(time.Now().Unix()-epochTime))
	}
}

// readUDP() a goroutine that just reads data off of a UDP socket and fills
// buffers.  Once a buffer is full, it passes it to handleBuff().
func readUDP(ip string, port int, c chan []byte) {
	var buff *[BUFFERSIZE]byte
	var offset int
	var timeout bool
	var addr = net.UDPAddr{
		Port: port,
		IP:   net.ParseIP(ip),
	}

	log.Printf("Listening on %s:%d\n", ip, port)
	sock, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Printf("Error opening UDP socket.\n")
		log.Fatalln(err)
	}
	defer sock.Close()

	log.Printf("Setting socket read buffer size to: %d\n", sockBufferMaxSize())
	err = sock.SetReadBuffer(sockBufferMaxSize())
	if err != nil {
		log.Printf("Unable to set read buffer size on socket.  Non-fatal.")
		log.Println(err)
	}
	err = sock.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		log.Printf("Unable to set timeout on socket.\n")
		log.Fatalln(err)
	}

	if verbose {
		log.Printf("Rock and Roll!\n")
	}
	for {
		if buff == nil {
			buff = new([BUFFERSIZE]byte)
			offset = 0
			timeout = false
		}

		i, err := sock.Read(buff[offset:])
		if err == nil {
			buff[offset+i] = '\n'
			offset = offset + i + 1
		} else if err.(net.Error).Timeout() {
			timeout = true
			err = sock.SetDeadline(time.Now().Add(time.Second))
			if err != nil {
				log.Panicln(err)
			}
		} else {
			log.Printf("Read Error: %s\n", err)
			continue
		}

		if offset > BUFFERSIZE-4096 || timeout {
			// Approching make buff size
			// we use a 4KiB margin
			c <- buff[:offset]
			buff = nil
		}
	}
}

// runServer() runs and manages this daemon, deals with OS signals, and handles
// communication channels.
func runServer(host string, port int) {
	var c chan []byte = make(chan []byte, 256)
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	var sig chan os.Signal = make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	// read incming UDP packets
	go readUDP(host, port, c)

	for {
		select {
		case buff := <-c:
			//fmt.Printf("Handling %d length buffer...\n", len(buff))
			go handleBuff(buff)
		case <-sig:
			log.Printf("Signal received.  Shutting down...\n")
			log.Printf("Received %d metrics.\n", totalMetrics)
			return
		}
	}
}

func main() {
	var bindAddress string
	var port int

	flag.IntVar(&port, "port", 9125, "Port to listen on")
	flag.IntVar(&port, "p", 9125, "Port to listen on")

	flag.StringVar(&bindAddress, "bind", "0.0.0.0", "IP Address to listen on")
	flag.StringVar(&bindAddress, "b", "0.0.0.0", "IP Address to listen on")

	flag.StringVar(&prefix, "prefix", "statsrelay", "The prefix to use with self generated stats")

	flag.BoolVar(&verbose, "verbose", false, "Verbose output")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatalf("One or most host specifications are needed to locate statsd daemons.\n")
	}

	hashRing = consistent.New()
	hashRing.NumberOfReplicas = 1

	for _, v := range flag.Args() {
		var addr *net.UDPAddr
		var err error
		host := strings.Split(v, ":")

		switch len(host) {
		case 1:
			log.Printf("Invalid statsd location: %s\n", v)
			log.Fatalf("Must be of the form HOST:PORT or HOST:PORT:INSTANCE\n")
		case 2:
			addr, err = net.ResolveUDPAddr("udp", v)
			if err != nil {
				log.Printf("Error parsing HOST:PORT \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		case 3:
			addr, err = net.ResolveUDPAddr("udp", host[0]+":"+host[1])
			if err != nil {
				log.Printf("Error parsing HOST:PORT:INSTANCE \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		default:
			log.Fatalf("Unrecongnized host specification: %s\n", v)
		}

		if addr != nil {
			udpAddr[v] = addr
			hashRing.Add(v)
		}
	}

	epochTime = time.Now().Unix()
	runServer(bindAddress, port)

	log.Printf("Normal shutdown.\n")
}
<commit_msg>resubmitting with go fmt :)<commit_after>package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"time"
)

import "stathat.com/c/consistent"

// BUFFERSIZE controls the size of the [...]byte array used to read UDP data
// off the wire and into local memory.  Metrics are separated by \n
// characters.  This buffer is passed to a handler to proxy out the metrics
// to the real statsd daemons.
const BUFFERSIZE int = 1 * 1024 * 1024 // 1MiB

// packetLen is the size in bytes of data we stuff into one packet before
// sending it to statsd.  This must be lower than the MTU, IPv4 header size
// and UDP header size.
const packetLen int = 1000

// prefix is the string that will be prefixed onto self generated stats.
// Such as <prefix>.statsProcessed.  Default is "statsrelay"
var prefix string

// udpAddr is a mapping of HOST:PORT:INSTANCE to a UDPAddr object
var udpAddr = make(map[string]*net.UDPAddr)

// hashRing is our consistent hashing ring.
var hashRing *consistent.Consistent

// totalMetrics tracks the totall number of metrics processed
var totalMetrics int = 0

// totalMetricsLock is a mutex gaurding totalMetrics
var totalMetricsLock sync.Mutex

// Time we began
var epochTime int64

// Verbose/Debug output
var verbose bool

// sockBufferMaxSize() returns the maximum size that the UDP receive buffer
// in the kernel can be set to.  In bytes.
func sockBufferMaxSize() int {

	// XXX: This is Linux-only most likely
	data, err := ioutil.ReadFile("/proc/sys/net/core/rmem_max")
	if err != nil {
		log.Panicln(err)
	}

	data = bytes.TrimRight(data, "\n\r")
	i, err := strconv.Atoi(string(data))
	if err != nil {
		log.Printf("Could not parse /proc/sys/net/core/rmem_max\n")
		log.Fatalln(err)
	}

	return i
}

// getMetricName() parses the given []byte metric as a string, extracts
// the metric key name and returns it as a string.
func getMetricName(metric []byte) (string, error) {
	// statsd metrics are of the form:
	//    KEY:VALUE|TYPE|RATE
	length := bytes.IndexByte(metric, byte(':'))
	if length == -1 {
		return "error", errors.New("Length of -1, must be invalid StatsD data")
	}
	return string(metric[:length]), nil
}

// sendPacket takes a []byte and writes that directly to a UDP socket
// that was assigned for target.
func sendPacket(buff []byte, target string) {
	conn, err := net.ListenUDP("udp", nil)
	if err != nil {
		log.Panicln(err)
	}

	conn.WriteToUDP(buff, udpAddr[target])
	conn.Close()
}

// buildPacketMap() is a helper function to initiallize a map that represents
// a UDP packet currently being built for each destination we proxy to.  As
// Go forbids taking the address of an object in a map or array so the
// bytes.Buffer object must be stored in the map as a pointer rather than
// a direct object in order to call the pointer methods on it.
func buildPacketMap() map[string]*bytes.Buffer {
	members := hashRing.Members()
	hash := make(map[string]*bytes.Buffer, len(members))

	for _, v := range members {
		hash[v] = new(bytes.Buffer)
	}

	return hash
}

// handleBuff() sorts through a full buffer of metrics and batches metrics
// to remote statsd daemons using a consistent hash.
func handleBuff(buff []byte) {
	packets := buildPacketMap()
	sep := []byte("\n")
	numMetrics := 0
	statsMetric := prefix + ".statsProcessed"

	for offset := 0; offset < len(buff); {
	loop:
		for offset < len(buff) {
			// Find our next value
			switch buff[offset] {
			case '\n':
				offset++
			case '\r':
				offset++
			case 0:
				offset++
			default:
				break loop
			}
		}

		size := bytes.IndexByte(buff[offset:], '\n')

		if size == -1 {
			// last metric in buffer
			size = len(buff) - offset
		}
		if size == 0 {
			// no more metrics
			break
		}

		//Check to ensure we get a metric, and not an invalid Byte sequence
		metric, err := getMetricName(buff[offset : offset+size])

		if err == nil {

			target, err := hashRing.Get(metric)
			if err != nil {
				log.Panicln(err)
			}

			// check built packet size and send if metric doesn't fit
			if packets[target].Len()+size > packetLen {
				sendPacket(packets[target].Bytes(), target)
				packets[target].Reset()
			}
			// add to packet
			packets[target].Write(buff[offset : offset+size])
			packets[target].Write(sep)

			numMetrics++
		}

		offset = offset + size + 1
	}

	if numMetrics == 0 {
		// if we haven't handled any metrics, then don't update counters/stats
		// or send packets
		return
	}

	// Update interal counter
	totalMetricsLock.Lock()
	totalMetrics = totalMetrics + numMetrics
	totalMetricsLock.Unlock()

	// Handle reporting our own stats
	stats := fmt.Sprintf("%s:%d|c\n", statsMetric, numMetrics)
	target, err := hashRing.Get(statsMetric)
	if err != nil {
		log.Panicln(err)
	}
	if packets[target].Len()+len(stats) > packetLen {
		sendPacket(packets[target].Bytes(), target)
		packets[target].Reset()
	}
	packets[target].Write([]byte(stats))

	// Empty out any remaining data
	for _, target := range hashRing.Members() {
		if packets[target].Len() > 0 {
			sendPacket(packets[target].Bytes(), target)
		}
	}

	if verbose {
		log.Printf("Procssed %d metrics. Running total: %d. Metrics/sec: %d\n",
			numMetrics, totalMetrics,
			int64(totalMetrics)/(time.Now().Unix()-epochTime))
	}
}

// readUDP() a goroutine that just reads data off of a UDP socket and fills
// buffers.  Once a buffer is full, it passes it to handleBuff().
func readUDP(ip string, port int, c chan []byte) {
	var buff *[BUFFERSIZE]byte
	var offset int
	var timeout bool
	var addr = net.UDPAddr{
		Port: port,
		IP:   net.ParseIP(ip),
	}

	log.Printf("Listening on %s:%d\n", ip, port)
	sock, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Printf("Error opening UDP socket.\n")
		log.Fatalln(err)
	}
	defer sock.Close()

	log.Printf("Setting socket read buffer size to: %d\n", sockBufferMaxSize())
	err = sock.SetReadBuffer(sockBufferMaxSize())
	if err != nil {
		log.Printf("Unable to set read buffer size on socket.  Non-fatal.")
		log.Println(err)
	}
	err = sock.SetDeadline(time.Now().Add(time.Second))
	if err != nil {
		log.Printf("Unable to set timeout on socket.\n")
		log.Fatalln(err)
	}

	if verbose {
		log.Printf("Rock and Roll!\n")
	}
	for {
		if buff == nil {
			buff = new([BUFFERSIZE]byte)
			offset = 0
			timeout = false
		}

		i, err := sock.Read(buff[offset:])
		if err == nil {
			buff[offset+i] = '\n'
			offset = offset + i + 1
		} else if err.(net.Error).Timeout() {
			timeout = true
			err = sock.SetDeadline(time.Now().Add(time.Second))
			if err != nil {
				log.Panicln(err)
			}
		} else {
			log.Printf("Read Error: %s\n", err)
			continue
		}

		if offset > BUFFERSIZE-4096 || timeout {
			// Approching make buff size
			// we use a 4KiB margin
			c <- buff[:offset]
			buff = nil
		}
	}
}

// runServer() runs and manages this daemon, deals with OS signals, and handles
// communication channels.
func runServer(host string, port int) {
	var c chan []byte = make(chan []byte, 256)
	// Set up channel on which to send signal notifications.
	// We must use a buffered channel or risk missing the signal
	// if we're not ready to receive when the signal is sent.
	var sig chan os.Signal = make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, os.Kill)

	// read incming UDP packets
	go readUDP(host, port, c)

	for {
		select {
		case buff := <-c:
			//fmt.Printf("Handling %d length buffer...\n", len(buff))
			go handleBuff(buff)
		case <-sig:
			log.Printf("Signal received.  Shutting down...\n")
			log.Printf("Received %d metrics.\n", totalMetrics)
			return
		}
	}
}

func main() {
	var bindAddress string
	var port int

	flag.IntVar(&port, "port", 9125, "Port to listen on")
	flag.IntVar(&port, "p", 9125, "Port to listen on")

	flag.StringVar(&bindAddress, "bind", "0.0.0.0", "IP Address to listen on")
	flag.StringVar(&bindAddress, "b", "0.0.0.0", "IP Address to listen on")

	flag.StringVar(&prefix, "prefix", "statsrelay", "The prefix to use with self generated stats")

	flag.BoolVar(&verbose, "verbose", false, "Verbose output")
	flag.BoolVar(&verbose, "v", false, "Verbose output")

	flag.Parse()

	if len(flag.Args()) == 0 {
		log.Fatalf("One or most host specifications are needed to locate statsd daemons.\n")
	}

	hashRing = consistent.New()
	hashRing.NumberOfReplicas = 1

	for _, v := range flag.Args() {
		var addr *net.UDPAddr
		var err error
		host := strings.Split(v, ":")

		switch len(host) {
		case 1:
			log.Printf("Invalid statsd location: %s\n", v)
			log.Fatalf("Must be of the form HOST:PORT or HOST:PORT:INSTANCE\n")
		case 2:
			addr, err = net.ResolveUDPAddr("udp", v)
			if err != nil {
				log.Printf("Error parsing HOST:PORT \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		case 3:
			addr, err = net.ResolveUDPAddr("udp", host[0]+":"+host[1])
			if err != nil {
				log.Printf("Error parsing HOST:PORT:INSTANCE \"%s\"\n", v)
				log.Fatalf("%s\n", err.Error())
			}
		default:
			log.Fatalf("Unrecongnized host specification: %s\n", v)
		}

		if addr != nil {
			udpAddr[v] = addr
			hashRing.Add(v)
		}
	}

	epochTime = time.Now().Unix()
	runServer(bindAddress, port)

	log.Printf("Normal shutdown.\n")
}
<|endoftext|>


<commit_before>// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package placementhandler

import (
	"net/http"
	"path"
	"time"

	"github.com/gogo/protobuf/jsonpb"
	"go.uber.org/zap"

	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/placementhandler/handleroptions"
	"github.com/m3db/m3/src/query/api/v1/route"
	"github.com/m3db/m3/src/query/generated/proto/admin"
	"github.com/m3db/m3/src/query/util/logging"
	xerrors "github.com/m3db/m3/src/x/errors"
	xhttp "github.com/m3db/m3/src/x/net/http"
)

const (
	// RemoveHTTPMethod is the HTTP method used with this resource.
	RemoveHTTPMethod = http.MethodPost
)

var (
	// M3DBRemoveURL is the url for the placement Remove handler (with the POST method)
	// for the M3DB service.
	M3DBRemoveURL = path.Join(route.Prefix, M3DBServicePlacementPathName)

	// M3AggRemoveURL is the url for the placement Remove handler (with the POST method)
	// for the M3Agg service.
	M3AggRemoveURL = path.Join(route.Prefix, M3AggServicePlacementPathName)

	// M3CoordinatorRemoveURL is the url for the placement Remove handler (with the POST method)
	// for the M3Coordinator service.
	M3CoordinatorRemoveURL = path.Join(route.Prefix, M3CoordinatorServicePlacementPathName)
)

// RemoveHandler is the handler for placement removes.
type RemoveHandler Handler

// NewRemoveHandler returns a new instance of RemoveHandler.
func NewRemoveHandler(opts HandlerOptions) *RemoveHandler {
	return &RemoveHandler{HandlerOptions: opts, nowFn: time.Now}
}

// ServeHTTP serves HTTP requests.
//nolint: dupl
func (h *RemoveHandler) ServeHTTP(
	svc handleroptions.ServiceNameAndDefaults,
	w http.ResponseWriter,
	r *http.Request,
) {
	ctx := r.Context()
	logger := logging.WithContext(ctx, h.instrumentOptions)

	req, rErr := h.parseRequest(r)
	if rErr != nil {
		xhttp.WriteError(w, rErr)
		return
	}

	placement, err := h.Remove(svc, r, req)
	if err != nil {
		logger.Error("unable to Remove placement", zap.Error(err))
		xhttp.WriteError(w, err)
		return
	}

	placementProto, err := placement.Proto()
	if err != nil {
		logger.Error("unable to get placement protobuf", zap.Error(err))
		xhttp.WriteError(w, err)
		return
	}

	resp := &admin.PlacementGetResponse{
		Placement: placementProto,
		Version:   int32(placement.Version()),
	}

	xhttp.WriteProtoMsgJSONResponse(w, resp, logger)
}

func (h *RemoveHandler) parseRequest(r *http.Request) (*admin.PlacementRemoveRequest, error) {
	defer r.Body.Close()

	removeReq := new(admin.PlacementRemoveRequest)
	if err := jsonpb.Unmarshal(r.Body, removeReq); err != nil {
		return nil, xerrors.NewInvalidParamsError(err)
	}

	return removeReq, nil
}

// Remove removes an instance.
func (h *RemoveHandler) Remove(
	svc handleroptions.ServiceNameAndDefaults,
	httpReq *http.Request,
	req *admin.PlacementRemoveRequest,
) (placement.Placement, error) {
	serviceOpts := handleroptions.NewServiceOptions(svc, httpReq.Header,
		h.m3AggServiceOptions)
	var validateFn placement.ValidateFn
	if !req.Force {
		validateFn = validateAllAvailable
	}

	pcfg, err := Handler(*h).PlacementConfigCopy()
	if err != nil {
		return nil, err
	}
	service, _, err := ServiceWithAlgo(
		h.clusterClient,
		serviceOpts,
		pcfg.ApplyOverride(req.OptionOverride),
		h.nowFn(),
		validateFn,
	)
	if err != nil {
		return nil, err
	}
	newPlacement, err := service.RemoveInstances(req.InstanceIds)
	if err != nil {
		return nil, err
	}

	return newPlacement, nil
}
<commit_msg>Use specific path for remove (#3700)<commit_after>// Copyright (c) 2021 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package placementhandler

import (
	"net/http"
	"path"
	"time"

	"github.com/gogo/protobuf/jsonpb"
	"go.uber.org/zap"

	"github.com/m3db/m3/src/cluster/placement"
	"github.com/m3db/m3/src/cluster/placementhandler/handleroptions"
	"github.com/m3db/m3/src/query/api/v1/route"
	"github.com/m3db/m3/src/query/generated/proto/admin"
	"github.com/m3db/m3/src/query/util/logging"
	xerrors "github.com/m3db/m3/src/x/errors"
	xhttp "github.com/m3db/m3/src/x/net/http"
)

const (
	// RemoveHTTPMethod is the HTTP method used with this resource.
	RemoveHTTPMethod = http.MethodPost

	removePathName = "remove"
)

var (
	// M3DBRemoveURL is the url for the placement Remove handler (with the POST method)
	// for the M3DB service.
	M3DBRemoveURL = path.Join(route.Prefix, M3DBServicePlacementPathName, removePathName)

	// M3AggRemoveURL is the url for the placement Remove handler (with the POST method)
	// for the M3Agg service.
	M3AggRemoveURL = path.Join(route.Prefix, M3AggServicePlacementPathName, removePathName)

	// M3CoordinatorRemoveURL is the url for the placement Remove handler (with the POST method)
	// for the M3Coordinator service.
	M3CoordinatorRemoveURL = path.Join(route.Prefix, M3CoordinatorServicePlacementPathName, removePathName)
)

// RemoveHandler is the handler for placement removes.
type RemoveHandler Handler

// NewRemoveHandler returns a new instance of RemoveHandler.
func NewRemoveHandler(opts HandlerOptions) *RemoveHandler {
	return &RemoveHandler{HandlerOptions: opts, nowFn: time.Now}
}

// ServeHTTP serves HTTP requests.
//nolint: dupl
func (h *RemoveHandler) ServeHTTP(
	svc handleroptions.ServiceNameAndDefaults,
	w http.ResponseWriter,
	r *http.Request,
) {
	ctx := r.Context()
	logger := logging.WithContext(ctx, h.instrumentOptions)

	req, rErr := h.parseRequest(r)
	if rErr != nil {
		xhttp.WriteError(w, rErr)
		return
	}

	placement, err := h.Remove(svc, r, req)
	if err != nil {
		logger.Error("unable to Remove placement", zap.Error(err))
		xhttp.WriteError(w, err)
		return
	}

	placementProto, err := placement.Proto()
	if err != nil {
		logger.Error("unable to get placement protobuf", zap.Error(err))
		xhttp.WriteError(w, err)
		return
	}

	resp := &admin.PlacementGetResponse{
		Placement: placementProto,
		Version:   int32(placement.Version()),
	}

	xhttp.WriteProtoMsgJSONResponse(w, resp, logger)
}

func (h *RemoveHandler) parseRequest(r *http.Request) (*admin.PlacementRemoveRequest, error) {
	defer r.Body.Close()

	removeReq := new(admin.PlacementRemoveRequest)
	if err := jsonpb.Unmarshal(r.Body, removeReq); err != nil {
		return nil, xerrors.NewInvalidParamsError(err)
	}

	return removeReq, nil
}

// Remove removes an instance.
func (h *RemoveHandler) Remove(
	svc handleroptions.ServiceNameAndDefaults,
	httpReq *http.Request,
	req *admin.PlacementRemoveRequest,
) (placement.Placement, error) {
	serviceOpts := handleroptions.NewServiceOptions(svc, httpReq.Header,
		h.m3AggServiceOptions)
	var validateFn placement.ValidateFn
	if !req.Force {
		validateFn = validateAllAvailable
	}

	pcfg, err := Handler(*h).PlacementConfigCopy()
	if err != nil {
		return nil, err
	}
	service, _, err := ServiceWithAlgo(
		h.clusterClient,
		serviceOpts,
		pcfg.ApplyOverride(req.OptionOverride),
		h.nowFn(),
		validateFn,
	)
	if err != nil {
		return nil, err
	}
	newPlacement, err := service.RemoveInstances(req.InstanceIds)
	if err != nil {
		return nil, err
	}

	return newPlacement, nil
}
<|endoftext|>


<commit_before>package providers

import (
	"context"
	"time"

	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
)

const (
	cloudInit = `#cloud-config

  coreos:
    units:
      - name: nimona.service
        command: start
        enable: true
        content: |
          [Unit]
          Description=nimona
          After=docker.service
          Requires=docker.service
          After=docker.redis.service

          [Service]
          TimeoutStartSec=0
          Restart=always
          ExecStartPre=-/usr/bin/docker stop nimona
          ExecStartPre=-/usr/bin/docker rm nimona
          ExecStartPre=/usr/bin/docker pull nimona/nimona:latest
          ExecStart=/usr/bin/docker run --name nimona --rm -p 21013:21013 nimona/nimona daemon --port=21013 --api-port=8080

          [Install]
          WantedBy=multi-user.target`
)

// DigitalOceanProvider prodides a DO operations
type DigitalOceanProvider struct {
	client *godo.Client
}

type tokenSource struct {
	AccessToken string
}

func (t *tokenSource) Token() (*oauth2.Token, error) {
	token := &oauth2.Token{
		AccessToken: t.AccessToken,
	}
	return token, nil
}

// NewDigitalocean creates a new DigitalOcean Provider
func NewDigitalocean(token string) (Provider, error) {
	if token == "" {
		return nil, ErrNoToken
	}

	tokenSource := &tokenSource{
		AccessToken: token,
	}

	oauthClient := oauth2.NewClient(context.Background(), tokenSource)
	client := godo.NewClient(oauthClient)

	return &DigitalOceanProvider{
		client: client,
	}, nil
}

// NewInstance creates a new DO Droplet
func (dp *DigitalOceanProvider) NewInstance(name, sshFingerprint,
	size, region string) (string, error) {
	if size == "" {
		size = "s-1vcpu-1gb"
	}

	if region == "" {
		region = "lon1"
	}

	ctx := context.Background()
	createRequest := &godo.DropletCreateRequest{
		Name:   name,
		Region: region,
		Size:   size,
		Image: godo.DropletCreateImage{
			Slug: "coreos-stable",
		},
		SSHKeys: []godo.DropletCreateSSHKey{godo.DropletCreateSSHKey{
			Fingerprint: sshFingerprint,
		}},
		UserData: cloudInit,
	}

	// Create server
	drop, _, err := dp.client.Droplets.Create(
		ctx, createRequest)
	if err != nil {
		return "", err
	}

	// Wait for the API to return an IP
	for {
		d, _, err := dp.client.Droplets.Get(ctx, drop.ID)
		if err != nil {
			return "", err
		}

		ip, err := d.PublicIPv4()
		if err != nil {
			return "", err
		}
		if ip != "" {
			return ip, nil
		}

		time.Sleep(2 * time.Second)
	}

}
<commit_msg>Add counter for newInstance for loop<commit_after>package providers

import (
	"context"
	"time"

	"github.com/digitalocean/godo"
	"golang.org/x/oauth2"
)

const (
	cloudInit = `#cloud-config

  coreos:
    units:
      - name: nimona.service
        command: start
        enable: true
        content: |
          [Unit]
          Description=nimona
          After=docker.service
          Requires=docker.service
          After=docker.redis.service

          [Service]
          TimeoutStartSec=0
          Restart=always
          ExecStartPre=-/usr/bin/docker stop nimona
          ExecStartPre=-/usr/bin/docker rm nimona
          ExecStartPre=/usr/bin/docker pull nimona/nimona:latest
          ExecStart=/usr/bin/docker run --name nimona --rm -p 21013:21013 nimona/nimona daemon --port=21013 --api-port=8080

          [Install]
          WantedBy=multi-user.target`
)

// DigitalOceanProvider prodides a DO operations
type DigitalOceanProvider struct {
	client *godo.Client
}

type tokenSource struct {
	AccessToken string
}

func (t *tokenSource) Token() (*oauth2.Token, error) {
	token := &oauth2.Token{
		AccessToken: t.AccessToken,
	}
	return token, nil
}

// NewDigitalocean creates a new DigitalOcean Provider
func NewDigitalocean(token string) (Provider, error) {
	if token == "" {
		return nil, ErrNoToken
	}

	tokenSource := &tokenSource{
		AccessToken: token,
	}

	oauthClient := oauth2.NewClient(context.Background(), tokenSource)
	client := godo.NewClient(oauthClient)

	return &DigitalOceanProvider{
		client: client,
	}, nil
}

// NewInstance creates a new DO Droplet
func (dp *DigitalOceanProvider) NewInstance(name, sshFingerprint,
	size, region string) (string, error) {
	if size == "" {
		size = "s-1vcpu-1gb"
	}

	if region == "" {
		region = "lon1"
	}

	ctx := context.Background()
	createRequest := &godo.DropletCreateRequest{
		Name:   name,
		Region: region,
		Size:   size,
		Image: godo.DropletCreateImage{
			Slug: "coreos-stable",
		},
		SSHKeys: []godo.DropletCreateSSHKey{godo.DropletCreateSSHKey{
			Fingerprint: sshFingerprint,
		}},
		UserData: cloudInit,
	}

	// Create server
	drop, _, err := dp.client.Droplets.Create(
		ctx, createRequest)
	if err != nil {
		return "", err
	}

	wn := 0

	// Wait for the API to return an IP
	for {
		d, _, err := dp.client.Droplets.Get(ctx, drop.ID)
		if err != nil {
			return "", err
		}

		ip, err := d.PublicIPv4()
		if err != nil {
			return "", err
		}
		if ip != "" {
			return ip, nil
		}

		if wn == 60 {
			break
		}

		wn++
		time.Sleep(2 * time.Second)

	}

}
<|endoftext|>


<commit_before>package eval

import (
	"errors"
	"strconv"

	"github.com/elves/elvish/daemon/api"
)

var ErrDaemonOffline = errors.New("daemon is offline")

func makeDaemonNamespace(daemon *api.Client) Namespace {
	// Obtain process ID
	daemonPid := func() Value {
		pid, err := daemon.Pid()
		maybeThrow(err)
		return String(strconv.Itoa(pid))
	}

	return Namespace{
		"pid":  MakeRoVariableFromCallback(daemonPid),
		"sock": NewRoVariable(String(daemon.SockPath())),

		FnPrefix + "spawn": NewRoVariable(&BuiltinFn{"daemon:spawn", daemonSpawn}),
	}
}

func daemonSpawn(ec *EvalCtx, args []Value, opts map[string]Value) {
	TakeNoArg(args)
	TakeNoOpt(opts)
	ec.ToSpawn.Spawn()
}
<commit_msg>eval: daemon:spawn now throws any spawning error.<commit_after>package eval

import (
	"errors"
	"strconv"

	"github.com/elves/elvish/daemon/api"
)

var ErrDaemonOffline = errors.New("daemon is offline")

func makeDaemonNamespace(daemon *api.Client) Namespace {
	// Obtain process ID
	daemonPid := func() Value {
		pid, err := daemon.Pid()
		maybeThrow(err)
		return String(strconv.Itoa(pid))
	}

	return Namespace{
		"pid":  MakeRoVariableFromCallback(daemonPid),
		"sock": NewRoVariable(String(daemon.SockPath())),

		FnPrefix + "spawn": NewRoVariable(&BuiltinFn{"daemon:spawn", daemonSpawn}),
	}
}

func daemonSpawn(ec *EvalCtx, args []Value, opts map[string]Value) {
	TakeNoArg(args)
	TakeNoOpt(opts)
	maybeThrow(ec.ToSpawn.Spawn())
}
<|endoftext|>


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

package rpc

import (
	"fmt"
	"math"
	"sort"
	"sync"
	"time"

	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/stop"
)

// How often the cluster offset is measured.
var monitorInterval = defaultHeartbeatInterval * 10

// RemoteClockMonitor keeps track of the most recent measurements of remote
// offsets from this node to connected nodes.
type RemoteClockMonitor struct {
	offsets map[string]RemoteOffset // Maps remote string addr to offset.
	lClock  *hlc.Clock              // The server clock.
	mu      sync.Mutex
	// Wall time in nanoseconds when we last monitored cluster offset.
	lastMonitoredAt int64
}

type clusterOffsetInterval struct {
	lowerbound, upperbound int64
}

func (i clusterOffsetInterval) String() string {
	return fmt.Sprintf("{%s, %s}", time.Duration(i.lowerbound), time.Duration(i.upperbound))
}

// majorityIntervalNotFoundError indicates that we could not find a majority
// overlap in our estimate of remote clocks.
type majorityIntervalNotFoundError struct {
	endpoints endpointList
}

func (m *majorityIntervalNotFoundError) Error() string {
	return fmt.Sprintf("unable to determine the true cluster time from remote clock endpoints %v", m.endpoints)
}

// endpoint represents an endpoint in the interval estimation of a single
// remote clock. It could be either the lowpoint or the highpoint of the
// interval.
//
// For example, if the remote clock offset bounds are [-5, 10], then it
// will be converted into two endpoints:
// endpoint{offset: -5, endType: -1}
// endpoint{offset: 10, endType: +1}
type endpoint struct {
	offset  int64 // The boundary offset represented by this endpoint.
	endType int   // -1 if lowpoint, +1 if highpoint.
}

// endpointList is a slice of endpoints, sorted by endpoint offset.
type endpointList []endpoint

// Implementation of sort.Interface.
func (l endpointList) Len() int {
	return len(l)
}
func (l endpointList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}
func (l endpointList) Less(i, j int) bool {
	if l[i].offset == l[j].offset {
		return l[i].endType < l[j].endType
	}
	return l[i].offset < l[j].offset
}

// newRemoteClockMonitor returns a monitor with the given server clock.
func newRemoteClockMonitor(clock *hlc.Clock) *RemoteClockMonitor {
	return &RemoteClockMonitor{
		offsets: map[string]RemoteOffset{},
		lClock:  clock,
	}
}

// UpdateOffset is a thread-safe way to update the remote clock measurements.
//
// It only updates the offset for addr if one the following three cases holds:
// 1. There is no prior offset for that address.
// 2. The old offset for addr was measured before r.lastMonitoredAt. We never
// use values during monitoring that are older than r.lastMonitoredAt.
// 3. The new offset's error is smaller than the old offset's error.
//
// The third case allows the monitor to use the most precise clock reading of
// the remote addr during the next findOffsetInterval() invocation. We may
// measure the remote clock several times before we next calculate the cluster
// offset. When we do the measurement, we want to use the reading with the
// smallest error. Because monitorInterval > heartbeatInterval, this gives us
// several chances to accurately read the remote clock. Note that we don't want
// monitorInterval to be too large, else we might end up relying on old
// information.
func (r *RemoteClockMonitor) UpdateOffset(addr string, offset RemoteOffset) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if oldOffset, ok := r.offsets[addr]; !ok {
		r.offsets[addr] = offset
	} else if oldOffset.MeasuredAt < r.lastMonitoredAt {
		// No matter what offset is, we weren't going to use oldOffset again,
		// because it was measured before the last cluster offset calculation.
		r.offsets[addr] = offset
	} else if offset.Uncertainty < oldOffset.Uncertainty {
		r.offsets[addr] = offset
	}

	if log.V(2) {
		log.Infof("update offset: %s %v", addr, r.offsets[addr])
	}
}

// MonitorRemoteOffsets periodically checks that the offset of this server's
// clock from the true cluster time is within MaxOffset. If the offset exceeds
// MaxOffset, then this method will trigger a fatal error, causing the node to
// suicide.
func (r *RemoteClockMonitor) MonitorRemoteOffsets(stopper *stop.Stopper) {
	if log.V(1) {
		log.Infof("monitoring cluster offset")
	}
	var monitorTimer util.Timer
	defer monitorTimer.Stop()
	for {
		monitorTimer.Reset(monitorInterval)
		select {
		case <-stopper.ShouldStop():
			return
		case <-monitorTimer.C:
			monitorTimer.Read = true
			offsetInterval, err := r.findOffsetInterval()
			// By the contract of the hlc, if the value is 0, then safety checking
			// of the max offset is disabled. However we may still want to
			// propagate the information to a status node.
			// TODO(embark): once there is a framework for collecting timeseries
			// data about the db, propagate the offset status to that.
			if r.lClock.MaxOffset() != 0 {
				if err != nil {
					log.Fatalf("clock offset from the cluster time "+
						"for remote clocks %v could not be determined: %s",
						r.offsets, err)
				}

				if !isHealthyOffsetInterval(offsetInterval, r.lClock.MaxOffset()) {
					log.Fatalf("clock offset from the cluster time "+
						"for remote clocks: %v is in interval: %s, which "+
						"indicates that the true offset is greater than %s",
						r.offsets, offsetInterval, r.lClock.MaxOffset())
				}
				if log.V(1) {
					log.Infof("healthy cluster offset: %s", offsetInterval)
				}
			}
			r.mu.Lock()
			r.lastMonitoredAt = r.lClock.PhysicalNow()
			r.mu.Unlock()
		}
	}
}

// isHealthyOffsetInterval returns true if the clusterOffsetInterval indicates
// that the node's offset is within maxOffset, else false. For example, if the
// offset interval is [-20, -11] and the maxOffset is 10 nanoseconds, then the
// clock offset must be too great, because no point in the interval is within
// the maxOffset.
func isHealthyOffsetInterval(i clusterOffsetInterval, maxOffset time.Duration) bool {
	return i.lowerbound <= maxOffset.Nanoseconds() && i.upperbound >= -maxOffset.Nanoseconds()
}

// The routine that measures this node's probable offset from the rest of the
// cluster. This offset is measured as a clusterOffsetInterval. For example,
// the output might be [-5, 10], which would indicate that this node's offset
// is likely between -5 and 10 nanoseconds from the average clock of the
// cluster.
//
// The intersection algorithm used here is documented at:
// http://infolab.stanford.edu/pub/cstr/reports/csl/tr/83/247/CSL-TR-83-247.pdf,
// commonly known as Marzullo's algorithm. If a remote clock is correct, its
// offset interval should encompass this clock's offset from the cluster time
// (see buildEndpointList()). If the majority of remote clock are correct, then
// their intervals should overlap over some region, which should include the
// true offset from the cluster time. This algorithm returns this region.
//
// If an interval cannot be found, an error is returned, indicating that
// a majority of remote node offset intervals do not overlap the cluster time.
func (r *RemoteClockMonitor) findOffsetInterval() (clusterOffsetInterval, error) {
	endpoints := r.buildEndpointList()
	sort.Sort(endpoints)
	numClocks := len(endpoints) / 2
	if log.V(1) {
		log.Infof("finding offset interval for monitorInterval: %s, numOffsets %d",
			monitorInterval, numClocks)
	}
	if numClocks == 0 {
		return clusterOffsetInterval{
			lowerbound: 0,
			upperbound: 0,
		}, nil
	}

	best := 0
	count := 0
	var lowerbound int64
	var upperbound int64

	// Find the interval which the most offset intervals overlap.
	for i, endpoint := range endpoints {
		count -= endpoint.endType
		if count > best {
			best = count
			lowerbound = endpoint.offset
			// Note the endType of the last endpoint is +1, so count < best.
			// Thus this code will never run when i = len(endpoint)-1.
			upperbound = endpoints[i+1].offset
		}
	}

	// Indicates that fewer than a majority of connected remote clocks seem to
	// encompass the central offset from the cluster, an error condition.
	if best <= numClocks/2 {
		return clusterOffsetInterval{
				lowerbound: math.MaxInt64,
				upperbound: math.MaxInt64,
			}, &majorityIntervalNotFoundError{
				endpoints: endpoints,
			}
	}

	// A majority of offset intervals overlap at this interval, which should
	// contain the true cluster offset.
	return clusterOffsetInterval{
		lowerbound: lowerbound,
		upperbound: upperbound,
	}, nil
}

// buildEndpointList() takes all the RemoteOffsets that are in the monitor, and
// turns these offsets into intervals which should encompass this node's true
// offset from the cluster time. It returns a list including the two endpoints
// of each interval.
//
// As a side effect, any RemoteOffsets that haven't been
// updated since the last monitoring are removed. (Side effects are nasty, but
// prevent us from running through the list an extra time under a lock).
//
// A RemoteOffset r is represented by this interval:
// [r.Offset - r.Uncertainty - MaxOffset, r.Offset + r.Uncertainty + MaxOffset],
// where MaxOffset is the furthest a node's clock can deviate from the cluster
// time. While the offset between this node and the remote time is actually
// within [r.Offset - r.Uncertainty, r.Offset + r.Uncertainty], we also must expand the
// interval by MaxOffset. This accounts for the fact that the remote clock is at
// most MaxOffset distance from the cluster time. Thus the expanded interval
// ought to contain this node's offset from the true cluster time, not just the
// offset from the remote clock's time.
func (r *RemoteClockMonitor) buildEndpointList() endpointList {
	r.mu.Lock()
	defer r.mu.Unlock()

	endpoints := make(endpointList, 0, len(r.offsets)*2)
	for addr, o := range r.offsets {
		// Remove anything that hasn't been updated since the last time offest
		// was measured. This indicates that we no longer have a connection to
		// that addr.
		if o.MeasuredAt < r.lastMonitoredAt {
			delete(r.offsets, addr)
			continue
		}

		lowpoint := endpoint{
			offset:  o.Offset - o.Uncertainty - r.lClock.MaxOffset().Nanoseconds(),
			endType: -1,
		}
		highpoint := endpoint{
			offset:  o.Offset + o.Uncertainty + r.lClock.MaxOffset().Nanoseconds(),
			endType: +1,
		}
		endpoints = append(endpoints, lowpoint, highpoint)
	}
	return endpoints
}
<commit_msg>rpc: print a proper interval<commit_after>// Copyright 2014 The Cockroach Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied. See the License for the specific language governing
// permissions and limitations under the License.
//
// Author: Kathy Spradlin (kathyspradlin@gmail.com)

package rpc

import (
	"fmt"
	"math"
	"sort"
	"sync"
	"time"

	"github.com/cockroachdb/cockroach/util"
	"github.com/cockroachdb/cockroach/util/hlc"
	"github.com/cockroachdb/cockroach/util/log"
	"github.com/cockroachdb/cockroach/util/stop"
)

// How often the cluster offset is measured.
var monitorInterval = defaultHeartbeatInterval * 10

// RemoteClockMonitor keeps track of the most recent measurements of remote
// offsets from this node to connected nodes.
type RemoteClockMonitor struct {
	offsets map[string]RemoteOffset // Maps remote string addr to offset.
	lClock  *hlc.Clock              // The server clock.
	mu      sync.Mutex
	// Wall time in nanoseconds when we last monitored cluster offset.
	lastMonitoredAt int64
}

type clusterOffsetInterval struct {
	lowerbound, upperbound int64
}

func (i clusterOffsetInterval) String() string {
	return fmt.Sprintf("[%s, %s]", time.Duration(i.lowerbound), time.Duration(i.upperbound))
}

// majorityIntervalNotFoundError indicates that we could not find a majority
// overlap in our estimate of remote clocks.
type majorityIntervalNotFoundError struct {
	endpoints endpointList
}

func (m *majorityIntervalNotFoundError) Error() string {
	return fmt.Sprintf("unable to determine the true cluster time from remote clock endpoints %v", m.endpoints)
}

// endpoint represents an endpoint in the interval estimation of a single
// remote clock. It could be either the lowpoint or the highpoint of the
// interval.
//
// For example, if the remote clock offset bounds are [-5, 10], then it
// will be converted into two endpoints:
// endpoint{offset: -5, endType: -1}
// endpoint{offset: 10, endType: +1}
type endpoint struct {
	offset  int64 // The boundary offset represented by this endpoint.
	endType int   // -1 if lowpoint, +1 if highpoint.
}

// endpointList is a slice of endpoints, sorted by endpoint offset.
type endpointList []endpoint

// Implementation of sort.Interface.
func (l endpointList) Len() int {
	return len(l)
}
func (l endpointList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}
func (l endpointList) Less(i, j int) bool {
	if l[i].offset == l[j].offset {
		return l[i].endType < l[j].endType
	}
	return l[i].offset < l[j].offset
}

// newRemoteClockMonitor returns a monitor with the given server clock.
func newRemoteClockMonitor(clock *hlc.Clock) *RemoteClockMonitor {
	return &RemoteClockMonitor{
		offsets: map[string]RemoteOffset{},
		lClock:  clock,
	}
}

// UpdateOffset is a thread-safe way to update the remote clock measurements.
//
// It only updates the offset for addr if one the following three cases holds:
// 1. There is no prior offset for that address.
// 2. The old offset for addr was measured before r.lastMonitoredAt. We never
// use values during monitoring that are older than r.lastMonitoredAt.
// 3. The new offset's error is smaller than the old offset's error.
//
// The third case allows the monitor to use the most precise clock reading of
// the remote addr during the next findOffsetInterval() invocation. We may
// measure the remote clock several times before we next calculate the cluster
// offset. When we do the measurement, we want to use the reading with the
// smallest error. Because monitorInterval > heartbeatInterval, this gives us
// several chances to accurately read the remote clock. Note that we don't want
// monitorInterval to be too large, else we might end up relying on old
// information.
func (r *RemoteClockMonitor) UpdateOffset(addr string, offset RemoteOffset) {
	r.mu.Lock()
	defer r.mu.Unlock()

	if oldOffset, ok := r.offsets[addr]; !ok {
		r.offsets[addr] = offset
	} else if oldOffset.MeasuredAt < r.lastMonitoredAt {
		// No matter what offset is, we weren't going to use oldOffset again,
		// because it was measured before the last cluster offset calculation.
		r.offsets[addr] = offset
	} else if offset.Uncertainty < oldOffset.Uncertainty {
		r.offsets[addr] = offset
	}

	if log.V(2) {
		log.Infof("update offset: %s %v", addr, r.offsets[addr])
	}
}

// MonitorRemoteOffsets periodically checks that the offset of this server's
// clock from the true cluster time is within MaxOffset. If the offset exceeds
// MaxOffset, then this method will trigger a fatal error, causing the node to
// suicide.
func (r *RemoteClockMonitor) MonitorRemoteOffsets(stopper *stop.Stopper) {
	if log.V(1) {
		log.Infof("monitoring cluster offset")
	}
	var monitorTimer util.Timer
	defer monitorTimer.Stop()
	for {
		monitorTimer.Reset(monitorInterval)
		select {
		case <-stopper.ShouldStop():
			return
		case <-monitorTimer.C:
			monitorTimer.Read = true
			offsetInterval, err := r.findOffsetInterval()
			// By the contract of the hlc, if the value is 0, then safety checking
			// of the max offset is disabled. However we may still want to
			// propagate the information to a status node.
			// TODO(embark): once there is a framework for collecting timeseries
			// data about the db, propagate the offset status to that.
			if r.lClock.MaxOffset() != 0 {
				if err != nil {
					log.Fatalf("clock offset from the cluster time "+
						"for remote clocks %v could not be determined: %s",
						r.offsets, err)
				}

				if !isHealthyOffsetInterval(offsetInterval, r.lClock.MaxOffset()) {
					log.Fatalf("clock offset from the cluster time "+
						"for remote clocks: %v is in interval: %s, which "+
						"indicates that the true offset is greater than %s",
						r.offsets, offsetInterval, r.lClock.MaxOffset())
				}
				if log.V(1) {
					log.Infof("healthy cluster offset: %s", offsetInterval)
				}
			}
			r.mu.Lock()
			r.lastMonitoredAt = r.lClock.PhysicalNow()
			r.mu.Unlock()
		}
	}
}

// isHealthyOffsetInterval returns true if the clusterOffsetInterval indicates
// that the node's offset is within maxOffset, else false. For example, if the
// offset interval is [-20, -11] and the maxOffset is 10 nanoseconds, then the
// clock offset must be too great, because no point in the interval is within
// the maxOffset.
func isHealthyOffsetInterval(i clusterOffsetInterval, maxOffset time.Duration) bool {
	return i.lowerbound <= maxOffset.Nanoseconds() && i.upperbound >= -maxOffset.Nanoseconds()
}

// The routine that measures this node's probable offset from the rest of the
// cluster. This offset is measured as a clusterOffsetInterval. For example,
// the output might be [-5, 10], which would indicate that this node's offset
// is likely between -5 and 10 nanoseconds from the average clock of the
// cluster.
//
// The intersection algorithm used here is documented at:
// http://infolab.stanford.edu/pub/cstr/reports/csl/tr/83/247/CSL-TR-83-247.pdf,
// commonly known as Marzullo's algorithm. If a remote clock is correct, its
// offset interval should encompass this clock's offset from the cluster time
// (see buildEndpointList()). If the majority of remote clock are correct, then
// their intervals should overlap over some region, which should include the
// true offset from the cluster time. This algorithm returns this region.
//
// If an interval cannot be found, an error is returned, indicating that
// a majority of remote node offset intervals do not overlap the cluster time.
func (r *RemoteClockMonitor) findOffsetInterval() (clusterOffsetInterval, error) {
	endpoints := r.buildEndpointList()
	sort.Sort(endpoints)
	numClocks := len(endpoints) / 2
	if log.V(1) {
		log.Infof("finding offset interval for monitorInterval: %s, numOffsets %d",
			monitorInterval, numClocks)
	}
	if numClocks == 0 {
		return clusterOffsetInterval{
			lowerbound: 0,
			upperbound: 0,
		}, nil
	}

	best := 0
	count := 0
	var lowerbound int64
	var upperbound int64

	// Find the interval which the most offset intervals overlap.
	for i, endpoint := range endpoints {
		count -= endpoint.endType
		if count > best {
			best = count
			lowerbound = endpoint.offset
			// Note the endType of the last endpoint is +1, so count < best.
			// Thus this code will never run when i = len(endpoint)-1.
			upperbound = endpoints[i+1].offset
		}
	}

	// Indicates that fewer than a majority of connected remote clocks seem to
	// encompass the central offset from the cluster, an error condition.
	if best <= numClocks/2 {
		return clusterOffsetInterval{
				lowerbound: math.MaxInt64,
				upperbound: math.MaxInt64,
			}, &majorityIntervalNotFoundError{
				endpoints: endpoints,
			}
	}

	// A majority of offset intervals overlap at this interval, which should
	// contain the true cluster offset.
	return clusterOffsetInterval{
		lowerbound: lowerbound,
		upperbound: upperbound,
	}, nil
}

// buildEndpointList() takes all the RemoteOffsets that are in the monitor, and
// turns these offsets into intervals which should encompass this node's true
// offset from the cluster time. It returns a list including the two endpoints
// of each interval.
//
// As a side effect, any RemoteOffsets that haven't been
// updated since the last monitoring are removed. (Side effects are nasty, but
// prevent us from running through the list an extra time under a lock).
//
// A RemoteOffset r is represented by this interval:
// [r.Offset - r.Uncertainty - MaxOffset, r.Offset + r.Uncertainty + MaxOffset],
// where MaxOffset is the furthest a node's clock can deviate from the cluster
// time. While the offset between this node and the remote time is actually
// within [r.Offset - r.Uncertainty, r.Offset + r.Uncertainty], we also must expand the
// interval by MaxOffset. This accounts for the fact that the remote clock is at
// most MaxOffset distance from the cluster time. Thus the expanded interval
// ought to contain this node's offset from the true cluster time, not just the
// offset from the remote clock's time.
func (r *RemoteClockMonitor) buildEndpointList() endpointList {
	r.mu.Lock()
	defer r.mu.Unlock()

	endpoints := make(endpointList, 0, len(r.offsets)*2)
	for addr, o := range r.offsets {
		// Remove anything that hasn't been updated since the last time offest
		// was measured. This indicates that we no longer have a connection to
		// that addr.
		if o.MeasuredAt < r.lastMonitoredAt {
			delete(r.offsets, addr)
			continue
		}

		lowpoint := endpoint{
			offset:  o.Offset - o.Uncertainty - r.lClock.MaxOffset().Nanoseconds(),
			endType: -1,
		}
		highpoint := endpoint{
			offset:  o.Offset + o.Uncertainty + r.lClock.MaxOffset().Nanoseconds(),
			endType: +1,
		}
		endpoints = append(endpoints, lowpoint, highpoint)
	}
	return endpoints
}
<|endoftext|>


<commit_before>// Package dgmux provides a simple Discord message route multiplexer that
// parses messages and then executes a matching registered handler, if found.
// dgMux can be used with both Disgord and the DiscordGo library.
package dgmux

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/bwmarrin/discordgo"
)

// Route holds information about a specific message route handler
type Route struct {
	Pattern     string      // match pattern that should trigger this route handler
	Description string      // short description of this route
	Help        string      // detailed help string for this route
	Run         HandlerFunc // route handler function to call
}

// Context holds a bit of extra data we pass along to route handlers
// This way processing some of this only needs to happen once.
type Context struct {
	Fields          []string
	Content         string
	GuildID         string
	IsDirected      bool
	IsPrivate       bool
	HasPrefix       bool
	HasMention      bool
	HasMentionFirst bool
}

// HandlerFunc is the function signature required for a message route handler.
type HandlerFunc func(*discordgo.Session, *discordgo.Message, *Context)

// Mux is the main struct for all dgMux methods.
type Mux struct {
	Routes  []*Route
	Default *Route
	Prefix  string
}

// New returns a new Discord message route mux
func New() *Mux {
	m := &Mux{}
	m.Prefix = "-dg "
	return &Mux{}
}

// Route allows you to register a route
func (m *Mux) Route(pattern, desc string, cb HandlerFunc) (*Route, error) {

	r := Route{}
	r.Pattern = pattern
	r.Description = desc
	r.Run = cb
	m.Routes = append(m.Routes, &r)

	return &r, nil
}

// FuzzyMatch attepts to find the best route match for a givin message.
func (m *Mux) FuzzyMatch(msg string) (*Route, []string) {

	// Tokenize the msg string into a slice of words
	fields := strings.Fields(msg)

	// no point to continue if there's no fields
	if len(fields) == 0 {
		return nil, nil
	}

	// Search though the command list for a match
	var r *Route
	var rank int

	var fk int
	for fk, fv := range fields {

		for _, rv := range m.Routes {

			// If we find an exact match, return that immediately.
			if rv.Pattern == fv {
				return rv, fields[fk:]
			}

			// Some "Fuzzy" searching...
			if strings.HasPrefix(rv.Pattern, fv) {
				if len(fv) > rank {
					r = rv
					rank = len(fv)
				}
			}
		}
	}
	return r, fields[fk:]
}

// OnMessageCreate is a DiscordGo Event Handler function.  This must be
// registered using the DiscordGo.Session.AddHandler function.  This function
// will receive all Discord messages and parse them for matches to registered
// routes.
func (m *Mux) OnMessageCreate(ds *discordgo.Session, mc *discordgo.MessageCreate) {

	var err error

	// Ignore all messages created by the Bot account itself
	if mc.Author.ID == ds.State.User.ID {
		return
	}

	// Fetch the channel for this Message
	var c *discordgo.Channel
	c, err = ds.State.Channel(mc.ChannelID)
	if err != nil {
		// Try fetching via REST API
		c, err = ds.Channel(mc.ChannelID)
		if err != nil {
			log.Printf("unable to fetch Channel for Message")
			return
		}
		// Attempt to add this channel into our State
		err = ds.State.ChannelAdd(c)
		if err != nil {
			log.Printf("error updating State with Channel")
		}
	}

	ctx := &Context{
		Content:   strings.TrimSpace(mc.Content),
		GuildID:   c.GuildID,
		IsPrivate: c.IsPrivate,
	}

	// Detect Private Message
	if c.IsPrivate {
		ctx.IsDirected = true
	}

	// Detect @name or @nick mentions
	if !ctx.IsDirected {

		// Detect if Bot was @mentioned
		for _, v := range mc.Mentions {

			if v.ID == ds.State.User.ID {

				ctx.IsDirected, ctx.HasMention = true, true

				reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", ds.State.User.ID))

				// Was the @mention the first part of the string?
				if reg.FindStringIndex(ctx.Content)[0] == 0 {
					ctx.HasMentionFirst = true
				}

				// strip bot mention tags from content string
				ctx.Content = reg.ReplaceAllString(ctx.Content, "")

				break
			}
		}
	}

	// Detect prefix mention
	if !ctx.IsDirected {

		// TODO : Must be changed to support a per-guild user defined prefix
		if strings.HasPrefix(ctx.Content, m.Prefix) {
			ctx.IsDirected, ctx.HasPrefix, ctx.HasMentionFirst = true, true, true
			ctx.Content = strings.TrimPrefix(ctx.Content, m.Prefix)
		}
	}

	// For now, if we're not specifically mentioned we do nothing.
	// later I might add an option for global non-mentioned command words
	if !ctx.IsDirected {
		return
	}

	// Try to find the "best match" command out of the message.
	r, fl := m.FuzzyMatch(ctx.Content)
	if r != nil {
		ctx.Fields = fl
		r.Run(ds, mc.Message, ctx)
		return
	}

	// If no command match was found, call the default.
	// Ignore if only @mentioned in the middle of a message
	if m.Default != nil && (ctx.HasMentionFirst) {
		// TODO: This could use a ratelimit
		// or should the ratelimit be inside the cmd handler?..
		// In the case of "talking" to another bot, this can create an endless
		// loop.  Probably most common in private messages.
		m.Default.Run(ds, mc.Message, ctx)
	}

}
<commit_msg>return the right Mux and prevent no Prefix spam<commit_after>// Package dgmux provides a simple Discord message route multiplexer that
// parses messages and then executes a matching registered handler, if found.
// dgMux can be used with both Disgord and the DiscordGo library.
package dgmux

import (
	"fmt"
	"log"
	"regexp"
	"strings"

	"github.com/bwmarrin/discordgo"
)

// Route holds information about a specific message route handler
type Route struct {
	Pattern     string      // match pattern that should trigger this route handler
	Description string      // short description of this route
	Help        string      // detailed help string for this route
	Run         HandlerFunc // route handler function to call
}

// Context holds a bit of extra data we pass along to route handlers
// This way processing some of this only needs to happen once.
type Context struct {
	Fields          []string
	Content         string
	GuildID         string
	IsDirected      bool
	IsPrivate       bool
	HasPrefix       bool
	HasMention      bool
	HasMentionFirst bool
}

// HandlerFunc is the function signature required for a message route handler.
type HandlerFunc func(*discordgo.Session, *discordgo.Message, *Context)

// Mux is the main struct for all dgMux methods.
type Mux struct {
	Routes  []*Route
	Default *Route
	Prefix  string
}

// New returns a new Discord message route mux
func New() *Mux {
	m := &Mux{}
	m.Prefix = "-dg "
	return m
}

// Route allows you to register a route
func (m *Mux) Route(pattern, desc string, cb HandlerFunc) (*Route, error) {

	r := Route{}
	r.Pattern = pattern
	r.Description = desc
	r.Run = cb
	m.Routes = append(m.Routes, &r)

	return &r, nil
}

// FuzzyMatch attepts to find the best route match for a givin message.
func (m *Mux) FuzzyMatch(msg string) (*Route, []string) {

	// Tokenize the msg string into a slice of words
	fields := strings.Fields(msg)

	// no point to continue if there's no fields
	if len(fields) == 0 {
		return nil, nil
	}

	// Search though the command list for a match
	var r *Route
	var rank int

	var fk int
	for fk, fv := range fields {

		for _, rv := range m.Routes {

			// If we find an exact match, return that immediately.
			if rv.Pattern == fv {
				return rv, fields[fk:]
			}

			// Some "Fuzzy" searching...
			if strings.HasPrefix(rv.Pattern, fv) {
				if len(fv) > rank {
					r = rv
					rank = len(fv)
				}
			}
		}
	}
	return r, fields[fk:]
}

// OnMessageCreate is a DiscordGo Event Handler function.  This must be
// registered using the DiscordGo.Session.AddHandler function.  This function
// will receive all Discord messages and parse them for matches to registered
// routes.
func (m *Mux) OnMessageCreate(ds *discordgo.Session, mc *discordgo.MessageCreate) {

	var err error

	// Ignore all messages created by the Bot account itself
	if mc.Author.ID == ds.State.User.ID {
		return
	}

	// Fetch the channel for this Message
	var c *discordgo.Channel
	c, err = ds.State.Channel(mc.ChannelID)
	if err != nil {
		// Try fetching via REST API
		c, err = ds.Channel(mc.ChannelID)
		if err != nil {
			log.Printf("unable to fetch Channel for Message")
			return
		}
		// Attempt to add this channel into our State
		err = ds.State.ChannelAdd(c)
		if err != nil {
			log.Printf("error updating State with Channel")
		}
	}

	ctx := &Context{
		Content:   strings.TrimSpace(mc.Content),
		GuildID:   c.GuildID,
		IsPrivate: c.IsPrivate,
	}

	// Detect Private Message
	if c.IsPrivate {
		ctx.IsDirected = true
	}

	// Detect @name or @nick mentions
	if !ctx.IsDirected {

		// Detect if Bot was @mentioned
		for _, v := range mc.Mentions {

			if v.ID == ds.State.User.ID {

				ctx.IsDirected, ctx.HasMention = true, true

				reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", ds.State.User.ID))

				// Was the @mention the first part of the string?
				if reg.FindStringIndex(ctx.Content)[0] == 0 {
					ctx.HasMentionFirst = true
				}

				// strip bot mention tags from content string
				ctx.Content = reg.ReplaceAllString(ctx.Content, "")

				break
			}
		}
	}

	// Detect prefix mention
	if !ctx.IsDirected && len(m.Prefix) > 0 {

		// TODO : Must be changed to support a per-guild user defined prefix
		if strings.HasPrefix(ctx.Content, m.Prefix) {
			ctx.IsDirected, ctx.HasPrefix, ctx.HasMentionFirst = true, true, true
			ctx.Content = strings.TrimPrefix(ctx.Content, m.Prefix)
		}
	}

	// For now, if we're not specifically mentioned we do nothing.
	// later I might add an option for global non-mentioned command words
	if !ctx.IsDirected {
		return
	}

	// Try to find the "best match" command out of the message.
	r, fl := m.FuzzyMatch(ctx.Content)
	if r != nil {
		ctx.Fields = fl
		r.Run(ds, mc.Message, ctx)
		return
	}

	// If no command match was found, call the default.
	// Ignore if only @mentioned in the middle of a message
	if m.Default != nil && (ctx.HasMentionFirst) {
		// TODO: This could use a ratelimit
		// or should the ratelimit be inside the cmd handler?..
		// In the case of "talking" to another bot, this can create an endless
		// loop.  Probably most common in private messages.
		m.Default.Run(ds, mc.Message, ctx)
	}

}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"

	"github.com/dmnlk/gomadare"
	"github.com/dmnlk/stringUtils"
	"github.com/rem7/goprowl"
)

var (
	CONSUMER_KEY        string
	CONSUMER_KEY_SECRET string
	ACCESS_TOKEN        string
	ACCESS_TOKEN_SECRET string
	PROWL_API_KEY       string
)

func main() {
	error := configureToken()
	if error != nil {
		fmt.Println(error)
		return
	}

	client := gomadare.NewClient(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET)

	client.GetUserStream(nil, func(s gomadare.Status, e gomadare.Event) {
		if &s != nil {

		}
		if &e != nil {
			sendEventToProwl(e)
		}
	})
}

func configureToken() error {
	CONSUMER_KEY = os.Getenv("CONSUMER_KEY")
	CONSUMER_KEY_SECRET = os.Getenv("CONSUMER_KEY_SECRET")
	ACCESS_TOKEN = os.Getenv("ACCESS_TOKEN")
	ACCESS_TOKEN_SECRET = os.Getenv("ACCESS_TOKEN_SECRET")
	PROWL_API_KEY = os.Getenv("PROWL_API_KEY")
	if ng := stringUtils.IsAnyEmpty(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET); ng {
		return fmt.Errorf("some key invalid")
	}

	return nil
}

func sendEventToProwl(e gomadare.Event) {

	var p goprowl.Goprowl
	err := p.RegisterKey(PROWL_API_KEY)
	if err != nil {
		fmt.Println(err)
		return
	}

	n := &goprowl.Notification{
		Application: "Twitter",
		Description:  e.Event + e.TargetObject.Text,
		Event:       e.Event + e.Source.Name,
		Priority:    "1",
		Providerkey: "",
		Url:         "www.foobar.com",
	}

	p.Push(n)

}
<commit_msg>check api key<commit_after>package main

import (
	"fmt"
	"os"

	"github.com/dmnlk/gomadare"
	"github.com/dmnlk/stringUtils"
	"github.com/rem7/goprowl"
)

var (
	CONSUMER_KEY        string
	CONSUMER_KEY_SECRET string
	ACCESS_TOKEN        string
	ACCESS_TOKEN_SECRET string
	PROWL_API_KEY       string
)

func main() {
	error := configureToken()
	if error != nil {
		fmt.Println(error)
		return
	}

	client := gomadare.NewClient(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET)

	client.GetUserStream(nil, func(s gomadare.Status, e gomadare.Event) {
		if &s != nil {

		}
		if &e != nil {
			sendEventToProwl(e)
		}
	})
}

func configureToken() error {
	CONSUMER_KEY = os.Getenv("CONSUMER_KEY")
	CONSUMER_KEY_SECRET = os.Getenv("CONSUMER_KEY_SECRET")
	ACCESS_TOKEN = os.Getenv("ACCESS_TOKEN")
	ACCESS_TOKEN_SECRET = os.Getenv("ACCESS_TOKEN_SECRET")
	PROWL_API_KEY = os.Getenv("PROWL_API_KEY")
	if ng := stringUtils.IsAnyEmpty(CONSUMER_KEY, CONSUMER_KEY_SECRET, ACCESS_TOKEN, ACCESS_TOKEN_SECRET, PROWL_API_KEY); ng {
		return fmt.Errorf("some key invalid")
	}

	return nil
}

func sendEventToProwl(e gomadare.Event) {

	var p goprowl.Goprowl
	err := p.RegisterKey(PROWL_API_KEY)
	if err != nil {
		fmt.Println(err)
		return
	}

	n := &goprowl.Notification{
		Application: "Twitter",
		Description:  e.Event + e.TargetObject.Text,
		Event:       e.Event + e.Source.Name,
		Priority:    "1",
		Providerkey: "",
		Url:         "www.foobar.com",
	}

	p.Push(n)

}
<|endoftext|>