<commit_before>package gockl

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
)

type DocumentInfo struct {
	Data         string
	ElementNames []string
}

var documents map[string]DocumentInfo = map[string]DocumentInfo{
	// taken from https://github.com/golang/go/issues/10158
	"doctype subset": {
		Data: `<?xml version="1.0" encoding="UTF-8" ?>
<!DOCTYPE doc [
    <!ELEMENT doc ANY>
]>
<doc>
</doc>`,
		ElementNames: []string{"doc", "doc"},
	},
	"empty start": {
		Data:         `<<a`,
		ElementNames: []string{},
	},
	"empty end": {
		Data:         `</<a`,
		ElementNames: []string{""},
	},
	"empty doctype": {
		Data:         `<!DOCTYPE[`,
		ElementNames: []string{},
	},
	"simple-svg": {
		Data: `<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" version="1.1" width="100%" height="100%" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 1920 1080">
  <style>
/* This is a comment. */
.test {
	fill: 'black';
}
  </style>
  <rect width="1920" height="1080" class="test" fill="red"></rect>
  <defs>
    <linearGradient id="grad">
      <stop stop-color="white" offset="0"></stop>
      <stop stop-opacity="0" stop-color="white" offset="1"></stop>
    </linearGradient>
  </defs>
</svg>`,
		ElementNames: []string{
			"svg", "style", "style", "rect", "rect", "defs", "linearGradient", "stop", "stop", "stop", "stop", "linearGradient", "defs", "svg",
		},
	},
}

func passthrough(data string) string {
	buf := bytes.Buffer{}
	z := New(data)

	for {
		t, err := z.Next()
		if err != nil {
			break
		}

		buf.WriteString(t.Raw())
	}

	return buf.String()
}

func elements(data string) []string {
	r := []string{}
	z := New(data)

	for {
		t, err := z.Next()
		if err != nil {
			break
		}

		if el, ok := t.(ElementToken); ok {
			r = append(r, el.Name())
		}
	}

	return r
}

func Test_NoChange(t *testing.T) {
	for name, info := range documents {
		if info.Data != passthrough(info.Data) {
			t.Errorf("Error processing document '%s'", name)
		}
	}
}

func Test_ElementNames(t *testing.T) {
	for name, info := range documents {
		elements := elements(info.Data)
		for pos, expected := range info.ElementNames {
			if pos >= len(elements) {
				t.Errorf("Element pos %d not existing for document %s", pos, name)
			} else if actual := elements[pos]; actual != expected {
				t.Errorf("Element name not matching at pos %d for document %s: %s (actual) != %s (expected)", pos, name, actual, expected)
			}
		}
	}
}

func Test_BrokenStartElement(t *testing.T) {
	input := "<elem"
	decoder := New(input)
	tok, err := decoder.Next()
	if err != nil {
		t.Error("Error while getting token.")
	}
	if _, ok := tok.(StartElementToken); !ok {
		t.Errorf("Not a start element token: %s, %s", input, reflect.TypeOf(tok))
	}
	if tok.Raw() != input {
		t.Errorf("Token text not matching: %s (expected) != %s (actual)", input, tok.Raw())
	}
	if next, err := decoder.Next(); next != nil || err != io.EOF {
		t.Errorf("Wanted EOF, got: '%s'/%s", next, err)
	}
}

func Test_BrokenTextElement(t *testing.T) {
	input := "/asdkjlh"
	decoder := New(input)
	tok, err := decoder.Next()
	if err != nil {
		t.Error("Error while getting token.")
	}
	if _, ok := tok.(TextToken); !ok {
		t.Errorf("Not a text token: %s, %s", input, reflect.TypeOf(tok))
	}
	if tok.Raw() != input {
		t.Errorf("Token text not matching: %s (expected) != %s (actual)", input, tok.Raw())
	}
	if next, err := decoder.Next(); next != nil || err != io.EOF {
		t.Errorf("Wanted EOF, got: '%s'/%s", next, err)
	}
}

// I have a directory of ~15.000 XML files created using: find ~ -name '*.xml' -exec cp {} ./XML \;
func Test_RealLifeFiles(t *testing.T) {
	files, err := ioutil.ReadDir(filepath.Join("testdata", "xml"))
	if err != nil {
		return
	}

	for _, fi := range files {
		raw, err := ioutil.ReadFile(filepath.Join("testdata", "xml", fi.Name()))
		if err != nil {
			t.Logf("Error reading %s: %s", fi.Name(), err)
			continue
		}

		input := string(raw)
		if output := passthrough(input); output != input {
			outfile := fmt.Sprintf("%s.tmp", filepath.Join("testdata", "xml", fi.Name()))
			t.Errorf("Error processing document '%s', writing output to %s for comparing to ", fi.Name(), outfile)
			ioutil.WriteFile(outfile, []byte(output), 0644)
			return
		}

		//t.Logf("%d/%d files checked.", i+1, len(files))
	}
}

func Test_Attributes(t *testing.T) {
	svg := StartElementToken(`<svg xmlns="http://www.w3.org/2000/svg" version=1.1 width='100%' height='a + b' xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 1920 1080" bla=blub bla>`)
	attrs := []Attribute{
		Attribute{"xmlns", "http://www.w3.org/2000/svg"},
		Attribute{"version", "1.1"},
		Attribute{"width", "100%"},
		Attribute{"height", "a + b"},
		Attribute{"xmlns:xlink", "http://www.w3.org/1999/xlink"},
		Attribute{"viewBox", "0 0 1920 1080"},
		Attribute{"bla", "blub"},
		Attribute{"bla", ""},
	}

	result := svg.Attributes()
	if len(result) != len(attrs) {
		t.Error("Attribute count not matching")
		t.Errorf("Expected: %v", attrs)
		t.Errorf("Got: %v", result)
	}

	for i, a := range result {
		if i >= len(attrs) {
			t.Errorf("Unexpected attribute: %s", a)
		} else if attrs[i] != a {
			t.Errorf("Attributes not matching. Expected %s, got %s", attrs[i], a)
		}
	}
}

func Test_AttributesInEmptyElements(t *testing.T) {
	svg := EmptyElementToken(`<circle cx="50" cy="25" r="20" fill="yellow" />`)
	attrs := []Attribute{
		Attribute{"cx", "50"},
		Attribute{"cy", "25"},
		Attribute{"r", "20"},
		Attribute{"fill", "yellow"},
	}

	if "circle" != svg.Name() {
		t.Error("No a circle")
	}

	result := svg.Attributes()
	if len(result) != len(attrs) {
		t.Error("Attribute count not matching")
		t.Errorf("Expected: %v", attrs)
		t.Errorf("Got: %v", result)
	}

	for i, a := range result {
		if i >= len(attrs) {
			t.Errorf("Unexpected attribute: %s", a)
		} else if attrs[i] != a {
			t.Errorf("Attributes not matching. Expected %s, got %s", attrs[i], a)
		}
	}
}

func TestGettingAttributesByName(t *testing.T) {
	type AttribTest struct {
		Token         Token
		Attributes    []Attribute
		NonAttributes []string
	}
	testdata := []AttribTest{
		AttribTest{
			EmptyElementToken(`<circle cx="50" cy="25" r="20" fill="yellow" />`),
			[]Attribute{
				Attribute{Name: "cx", Content: "50"},
				Attribute{Name: "R", Content: "20"},
			},
			[]string{"bla"},
		},
		AttribTest{
			StartElementToken(`<group style="fill: none;" style="nope">`),
			[]Attribute{
				Attribute{Name: "style", Content: "fill: none;"},
			},
			[]string{"r"},
		},
	}

	for _, i := range testdata {
		tok, ok := i.Token.(StartOrEmptyElementToken)
		if !ok {
			t.FailNow()
		}

		for _, a := range i.Attributes {
			for _, fn := range []func(string) string{
				strings.ToLower,
				strings.ToUpper,
				func(x string) string { return x },
			} {
				name := fn(a.Name)
				res, ok := tok.Attribute(name)
				if !ok {
					t.Errorf("Missing attribute %s in token %s", name, i.Token)
					continue
				}
				if res != a.Content {
					t.Errorf("Wrong attribute %s in token %s: Expected %s, got %s", name, i.Token, a.Content, res)
				}
			}
		}
	}
}

func Test_CDATA(t *testing.T) {
	doc := New(`<p><![CDATA[</p>]]><!-- </p> --></p>`)

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if start, ok := tok.(StartElementToken); !ok {
		t.Errorf("Expected start element token, got: %s", start)
	}

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if cdata, ok := tok.(CDATAToken); !ok {
		t.Errorf("Expected CDATA token, got: %s", cdata)
	} else if raw := cdata.Raw(); raw != `<![CDATA[</p>]]>` {
		t.Errorf("Wrong content for CDATA token: %s", raw)
	}

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if comment, ok := tok.(CommentToken); !ok {
		t.Errorf("Expected end element token, got: %s", comment)
	} else if raw := comment.Raw(); raw != `<!-- </p> -->` {
		t.Errorf("Wrong content for end token: %s", raw)
	}

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if end, ok := tok.(EndElementToken); !ok {
		t.Errorf("Expected end element token, got: %s", end)
	} else if raw := end.Raw(); raw != `</p>` {
		t.Errorf("Wrong content for end token: %s", raw)
	}
}
<commit_msg>Simplify tests.<commit_after>package gockl

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"path/filepath"
	"reflect"
	"strings"
	"testing"
)

type DocumentInfo struct {
	Data         string
	ElementNames []string
}

var documents map[string]DocumentInfo = map[string]DocumentInfo{
	// taken from https://github.com/golang/go/issues/10158
	"doctype subset": {
		Data: `<?xml version="1.0" encoding="UTF-8" ?>
<!DOCTYPE doc [
    <!ELEMENT doc ANY>
]>
<doc>
</doc>`,
		ElementNames: []string{"doc", "doc"},
	},
	"empty start": {
		Data:         `<<a`,
		ElementNames: []string{},
	},
	"empty end": {
		Data:         `</<a`,
		ElementNames: []string{""},
	},
	"empty doctype": {
		Data:         `<!DOCTYPE[`,
		ElementNames: []string{},
	},
	"simple-svg": {
		Data: `<?xml version="1.0" encoding="UTF-8"?>
<svg xmlns="http://www.w3.org/2000/svg" version="1.1" width="100%" height="100%" xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 1920 1080">
  <style>
/* This is a comment. */
.test {
	fill: 'black';
}
  </style>
  <rect width="1920" height="1080" class="test" fill="red"></rect>
  <defs>
    <linearGradient id="grad">
      <stop stop-color="white" offset="0"></stop>
      <stop stop-opacity="0" stop-color="white" offset="1"></stop>
    </linearGradient>
  </defs>
</svg>`,
		ElementNames: []string{
			"svg", "style", "style", "rect", "rect", "defs", "linearGradient", "stop", "stop", "stop", "stop", "linearGradient", "defs", "svg",
		},
	},
}

func passthrough(data string) string {
	buf := bytes.Buffer{}
	z := New(data)

	for {
		t, err := z.Next()
		if err != nil {
			break
		}

		buf.WriteString(t.Raw())
	}

	return buf.String()
}

func elements(data string) []string {
	r := []string{}
	z := New(data)

	for {
		t, err := z.Next()
		if err != nil {
			break
		}

		if el, ok := t.(ElementToken); ok {
			r = append(r, el.Name())
		}
	}

	return r
}

func Test_NoChange(t *testing.T) {
	for name, info := range documents {
		if info.Data != passthrough(info.Data) {
			t.Errorf("Error processing document '%s'", name)
		}
	}
}

func Test_ElementNames(t *testing.T) {
	for name, info := range documents {
		elements := elements(info.Data)
		for pos, expected := range info.ElementNames {
			if pos >= len(elements) {
				t.Errorf("Element pos %d not existing for document %s", pos, name)
			} else if actual := elements[pos]; actual != expected {
				t.Errorf("Element name not matching at pos %d for document %s: %s (actual) != %s (expected)", pos, name, actual, expected)
			}
		}
	}
}

func Test_BrokenStartElement(t *testing.T) {
	input := "<elem"
	decoder := New(input)
	tok, err := decoder.Next()
	if err != nil {
		t.Error("Error while getting token.")
	}
	if _, ok := tok.(StartElementToken); !ok {
		t.Errorf("Not a start element token: %s, %s", input, reflect.TypeOf(tok))
	}
	if tok.Raw() != input {
		t.Errorf("Token text not matching: %s (expected) != %s (actual)", input, tok.Raw())
	}
	if next, err := decoder.Next(); next != nil || err != io.EOF {
		t.Errorf("Wanted EOF, got: '%s'/%s", next, err)
	}
}

func Test_BrokenTextElement(t *testing.T) {
	input := "/asdkjlh"
	decoder := New(input)
	tok, err := decoder.Next()
	if err != nil {
		t.Error("Error while getting token.")
	}
	if _, ok := tok.(TextToken); !ok {
		t.Errorf("Not a text token: %s, %s", input, reflect.TypeOf(tok))
	}
	if tok.Raw() != input {
		t.Errorf("Token text not matching: %s (expected) != %s (actual)", input, tok.Raw())
	}
	if next, err := decoder.Next(); next != nil || err != io.EOF {
		t.Errorf("Wanted EOF, got: '%s'/%s", next, err)
	}
}

// I have a directory of ~15.000 XML files created using: find ~ -name '*.xml' -exec cp {} ./XML \;
func Test_RealLifeFiles(t *testing.T) {
	files, err := ioutil.ReadDir(filepath.Join("testdata", "xml"))
	if err != nil {
		return
	}

	for _, fi := range files {
		raw, err := ioutil.ReadFile(filepath.Join("testdata", "xml", fi.Name()))
		if err != nil {
			t.Logf("Error reading %s: %s", fi.Name(), err)
			continue
		}

		input := string(raw)
		if output := passthrough(input); output != input {
			outfile := fmt.Sprintf("%s.tmp", filepath.Join("testdata", "xml", fi.Name()))
			t.Errorf("Error processing document '%s', writing output to %s for comparing to ", fi.Name(), outfile)
			ioutil.WriteFile(outfile, []byte(output), 0644)
			return
		}

		//t.Logf("%d/%d files checked.", i+1, len(files))
	}
}

func Test_Attributes(t *testing.T) {
	svg := StartElementToken(`<svg xmlns="http://www.w3.org/2000/svg" version=1.1 width='100%' height='a + b' xmlns:xlink="http://www.w3.org/1999/xlink" viewBox="0 0 1920 1080" bla=blub bla>`)
	attrs := []Attribute{
		{"xmlns", "http://www.w3.org/2000/svg"},
		{"version", "1.1"},
		{"width", "100%"},
		{"height", "a + b"},
		{"xmlns:xlink", "http://www.w3.org/1999/xlink"},
		{"viewBox", "0 0 1920 1080"},
		{"bla", "blub"},
		{"bla", ""},
	}

	result := svg.Attributes()
	if len(result) != len(attrs) {
		t.Error("Attribute count not matching")
		t.Errorf("Expected: %v", attrs)
		t.Errorf("Got: %v", result)
	}

	for i, a := range result {
		if i >= len(attrs) {
			t.Errorf("Unexpected attribute: %s", a)
		} else if attrs[i] != a {
			t.Errorf("Attributes not matching. Expected %s, got %s", attrs[i], a)
		}
	}
}

func Test_AttributesInEmptyElements(t *testing.T) {
	svg := EmptyElementToken(`<circle cx="50" cy="25" r="20" fill="yellow" />`)
	attrs := []Attribute{
		{"cx", "50"},
		{"cy", "25"},
		{"r", "20"},
		{"fill", "yellow"},
	}

	if "circle" != svg.Name() {
		t.Error("No a circle")
	}

	result := svg.Attributes()
	if len(result) != len(attrs) {
		t.Error("Attribute count not matching")
		t.Errorf("Expected: %v", attrs)
		t.Errorf("Got: %v", result)
	}

	for i, a := range result {
		if i >= len(attrs) {
			t.Errorf("Unexpected attribute: %s", a)
		} else if attrs[i] != a {
			t.Errorf("Attributes not matching. Expected %s, got %s", attrs[i], a)
		}
	}
}

func TestGettingAttributesByName(t *testing.T) {
	type AttribTest struct {
		Token         Token
		Attributes    []Attribute
		NonAttributes []string
	}
	testdata := []AttribTest{
		{
			EmptyElementToken(`<circle cx="50" cy="25" r="20" fill="yellow" />`),
			[]Attribute{
				{Name: "cx", Content: "50"},
				{Name: "R", Content: "20"},
			},
			[]string{"bla"},
		},
		{
			StartElementToken(`<group style="fill: none;" style="nope">`),
			[]Attribute{
				{Name: "style", Content: "fill: none;"},
			},
			[]string{"r"},
		},
	}

	for _, i := range testdata {
		tok, ok := i.Token.(StartOrEmptyElementToken)
		if !ok {
			t.FailNow()
		}

		for _, a := range i.Attributes {
			for _, fn := range []func(string) string{
				strings.ToLower,
				strings.ToUpper,
				func(x string) string { return x },
			} {
				name := fn(a.Name)
				res, ok := tok.Attribute(name)
				if !ok {
					t.Errorf("Missing attribute %s in token %s", name, i.Token)
					continue
				}
				if res != a.Content {
					t.Errorf("Wrong attribute %s in token %s: Expected %s, got %s", name, i.Token, a.Content, res)
				}
			}
		}
	}
}

func Test_CDATA(t *testing.T) {
	doc := New(`<p><![CDATA[</p>]]><!-- </p> --></p>`)

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if start, ok := tok.(StartElementToken); !ok {
		t.Errorf("Expected start element token, got: %s", start)
	}

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if cdata, ok := tok.(CDATAToken); !ok {
		t.Errorf("Expected CDATA token, got: %s", cdata)
	} else if raw := cdata.Raw(); raw != `<![CDATA[</p>]]>` {
		t.Errorf("Wrong content for CDATA token: %s", raw)
	}

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if comment, ok := tok.(CommentToken); !ok {
		t.Errorf("Expected end element token, got: %s", comment)
	} else if raw := comment.Raw(); raw != `<!-- </p> -->` {
		t.Errorf("Wrong content for end token: %s", raw)
	}

	if tok, err := doc.Next(); err != nil {
		t.Error(err)
	} else if end, ok := tok.(EndElementToken); !ok {
		t.Errorf("Expected end element token, got: %s", end)
	} else if raw := end.Raw(); raw != `</p>` {
		t.Errorf("Wrong content for end token: %s", raw)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"testing"
)

type docCheck struct {
	ct []*CheckedLexeme
}

func (dc *docCheck) get(t *testing.T, path string) *docCheck {
	if len(dc.ct) == 0 {
		sp, err := NewSpellcheck("")
		if err != nil {
			t.Fatal(err)
		}
		defer sp.Close()
		ct, err := sp.Check([]string{path})
		if err != nil {
			t.Fatal(err)
		}
		dc.ct = ct
	}
	return dc
}

// this will trigger a warning that TestGoDocNoFuncNameReject isn't the function name
func TestGoDocNoFuncNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocNoFuncNameReject")
}

// TestGoDocNoFuncNameReject has the wrong name; want TestGoDocWrongFuncNameReject.
func TestGoDocWrongFuncNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocWrongFuncNameReject")
}

// this will trigger TestGoDocNoTypeNameReject
type godocStructReject struct {
	// this will trigger TestGoDocNoFieldNameReject
	oopsie int
	// okthough will pass TestGoDocFieldNamePass
	okthough int
}

// TestGoDocNoTypeNameReject rejects a type missing name in comments
func TestGoDocNoTypeNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocNoFuncNameReject")
}

// TestGoDocNoFieldNameReject rejects a field missing name in comments
func TestGoDocNoFieldNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocNoFieldNameReject")
}

// TestGoDocFieldNamePass accepts a godoc of a field
func TestGoDocFieldNamePass(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocFieldNamePass")
}

// TestGoDocFuncPass will not trigger a warning
func TestGoDocFuncPass(t *testing.T) {
	accept(t, "godoc_test.go", "TestGoDocFuncPass")
}

// TestGoDocMultiLinePass should pass
// even though this is a multiple line comment
func TestGoDocMultiLinePass(t *testing.T) {
	accept(t, "godoc_test.go", "TestGoDocMultiLinePass")
}

// this comment not part of the function documentation

// TestGoDocCommentBreak will not trigger a warning
func TestGoDocCommentBreak(t *testing.T) {
}
<commit_msg>godoc: fill out empty test case<commit_after>package main

import (
	"testing"
)

type docCheck struct {
	ct []*CheckedLexeme
}

func (dc *docCheck) get(t *testing.T, path string) *docCheck {
	if len(dc.ct) == 0 {
		sp, err := NewSpellcheck("")
		if err != nil {
			t.Fatal(err)
		}
		defer sp.Close()
		ct, err := sp.Check([]string{path})
		if err != nil {
			t.Fatal(err)
		}
		dc.ct = ct
	}
	return dc
}

// this will trigger a warning that TestGoDocNoFuncNameReject isn't the function name
func TestGoDocNoFuncNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocNoFuncNameReject")
}

// TestGoDocNoFuncNameReject has the wrong name; want TestGoDocWrongFuncNameReject.
func TestGoDocWrongFuncNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocWrongFuncNameReject")
}

// this will trigger TestGoDocNoTypeNameReject
type godocStructReject struct {
	// this will trigger TestGoDocNoFieldNameReject
	oopsie int
	// okthough will pass TestGoDocFieldNamePass
	okthough int
}

// TestGoDocNoTypeNameReject rejects a type missing name in comments
func TestGoDocNoTypeNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocNoFuncNameReject")
}

// TestGoDocNoFieldNameReject rejects a field missing name in comments
func TestGoDocNoFieldNameReject(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocNoFieldNameReject")
}

// TestGoDocFieldNamePass accepts a godoc of a field
func TestGoDocFieldNamePass(t *testing.T) {
	reject(t, "godoc_test.go", "TestGoDocFieldNamePass")
}

// TestGoDocFuncPass will not trigger a warning
func TestGoDocFuncPass(t *testing.T) {
	accept(t, "godoc_test.go", "TestGoDocFuncPass")
}

// TestGoDocMultiLinePass should pass
// even though this is a multiple line comment
func TestGoDocMultiLinePass(t *testing.T) {
	accept(t, "godoc_test.go", "TestGoDocMultiLinePass")
}

// this comment not part of the function documentation

// TestGoDocCommentBreak will not trigger a warning
func TestGoDocCommentBreak(t *testing.T) {
	accept(t, "godoc_test.go", "TestGoDocCommentBreak")
}
<|endoftext|>


<commit_before>package gologspace

import(
	"math"
)

func LogAddAll(logX float64, logYs ...float64) float64 {
	result := logX
	for _, logY := range logYs {
		result = LogAdd(result, logY)
	}
	return result
}

func LogAdd(logX, logY float64) float64 {
	if logY > logX {
		temp := logX
		logX = logY
		logY = temp
	}
	if logX == math.Inf(-1) {
		return logX
	}
	negDiff := logY - logX
	if negDiff < 20.0 {
		return logX
	}
	return logX + math.Log(1.0 + math.Exp(negDiff))
}

func LogDivAll(logX float64, logYs ...float64) float64 {
	result := logX
	for _, logY := range logYs {
		result = LogDiv(result, logY)
	}
	return result
}

func LogDiv(logX, logY float64) float64 {
	return logX - logY
}

func LogMultAll(logX float64, logYs ...float64) float64 {
	result := logX
	for _, logY := range logYs {
		result = LogMult(result, logY)
	}
	return result
}

func LogMult(logX, logY float64) float64 {
	return logX + logY
}

func LogProb(p float64) float64 {
	return math.Log(p)
}
<commit_msg>nasty bug<commit_after>package gologspace

import(
	"math"
)

func LogAddAll(logX float64, logYs ...float64) float64 {
	result := logX
	for _, logY := range logYs {
		result = LogAdd(result, logY)
	}
	return result
}

func LogAdd(logX, logY float64) float64 {
	if logY > logX {
		temp := logX
		logX = logY
		logY = temp
	}
	if logX == math.Inf(-1) {
		return logX
	}
	negDiff := logY - logX
	if negDiff < -20.0 {
		return logX
	}
	return logX + math.Log(1.0 + math.Exp(negDiff))
}

func LogDivAll(logX float64, logYs ...float64) float64 {
	result := logX
	for _, logY := range logYs {
		result = LogDiv(result, logY)
	}
	return result
}

func LogDiv(logX, logY float64) float64 {
	return logX - logY
}

func LogMultAll(logX float64, logYs ...float64) float64 {
	result := logX
	for _, logY := range logYs {
		result = LogMult(result, logY)
	}
	return result
}

func LogMult(logX, logY float64) float64 {
	return logX + logY
}

func LogProb(p float64) float64 {
	return math.Log(p)
}
<|endoftext|>


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

package main

/*
*  File: gonew_main.go
*  Author: Bryan Matsuo [bmatsuo@soe.ucsc.edu]
*  Created: Sat Jul  2 19:17:53 PDT 2011
*  Usage: gonew [options]
 */
import (
	"github.com/bmatsuo/gonew/config"
	"github.com/bmatsuo/gonew/project"
	"github.com/bmatsuo/gonew/templates"

	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"runtime"
	"strings"
	"syscall"
	"text/template"
	"time"
	"unicode"
)

var (
	GonewRoot    string // The Gonew source directory
	TemplateRoot string // The "templates" subdirectory of GonewRoot.
)

// Returns a path to the directory containing Gonew's source (and the templates/ directory).
// This function first searches locates the 'gonew' executable in PATH. Then locates the directory
// it was built from (either GOROOT or a subdirectory of GOPATH).
func FindGonew() (dir string, err error) {
	var bin string
	if bin, err = exec.LookPath("gonew"); err != nil {
		return
	}
	var bindir string
	if bindir, err = filepath.Abs(filepath.Dir(bin)); err != nil {
		return
	}
	var gobin string
	var gopath string
	for _, env := range os.Environ() {
		if strings.HasPrefix(env, "GOBIN=") {
			if gobin, err = filepath.Abs(env[6:]); err != nil {
				return
			}
		}
		if strings.HasPrefix(env, "GOPATH=") {
			if gopath, err = filepath.Abs(env[7:]); err != nil {
				return
			}
		}
		if len(gobin) > 0 && len(gopath) > 0 {
			break
		}
	}
	if bindir == gobin {
		rootdir := filepath.Join(runtime.GOROOT(), "github.com", "bmtasuo", "gonew")
		var stat os.FileInfo
		switch stat, err = os.Stat(rootdir); {
		case err != nil:
			if err != syscall.ENOENT {
				return
			}
		case stat != nil:
			dir = rootdir
			return
		default:
			panic("unreachable")
		}
	}
	var stat os.FileInfo
	if filepath.Base(bindir) != "bin" {
		err = fmt.Errorf("%s not under a GOPATH", bindir)
		return
	}
	gosrc := filepath.Join(filepath.Dir(bindir), "src")
	stat, err = os.Stat(gosrc)
	if err != nil || !stat.IsDir() {
		err = fmt.Errorf("%s does not exist", gosrc)
		return
	}

	dirCandidates := [...]string{
		filepath.Join(gosrc, "pkg", "github.com", "bmatsuo", "gonew"),
		filepath.Join(gosrc, "github.com", "bmatsuo", "gonew"),
		filepath.Join(gosrc, "gonew"),
	}
	for _, dir = range dirCandidates {
		stat, err = os.Stat(dir)
		if err == nil && stat.IsDir() {
			// found
			return
		}
	}
	err = fmt.Errorf("Couldn't find gonew source directory")
	return
}

func init() {
	root, err := FindGonew()
	GonewRoot = root
	if err != nil {
		fmt.Fprintf(os.Stderr, "error finding gonew source directory: %v\n", err)
		os.Exit(1)
	}
}

func check(err error) error {
	if err != nil {
		fmt.Println(err)
	}
	return err
}

func checkFatal(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func logJson(v ...interface{}) {
	w := make([]interface{}, 0, len(v))
	w = append(w, v[:len(v)-1]...)
	p, _ := json.MarshalIndent(v[len(v)-1], " ", "\t")
	w = append(w, string(p))
	fmt.Println(w...)
}

func executeHooks(ts templates.Interface, tenv templates.Environment, hooks ...*config.HookConfig) {
	for _, hook := range hooks {
		cwd, err := tenv.RenderTextAsString(ts, "cwd_", hook.Cwd)
		checkFatal(err)
		// fmt.Println("cd", cwd)
		for _, _cmd := range hook.Commands {
			cmd, err := tenv.RenderTextAsString(ts, "cmd_", _cmd)
			checkFatal(err)
			// fmt.Println("bash", "-c", cmd)
			shell := exec.Command("bash", "-c", cmd)
			shell.Dir = cwd
			shell.Stdin = os.Stdin
			shell.Stdout = os.Stdout
			shell.Stderr = os.Stderr
			checkFatal(shell.Run()) // TODO clean exit
		}
	}
}

type File2 struct {
	path    string
	content []byte
}

func funcsV2(env *config.EnvironmentConfig) template.FuncMap {
	return template.FuncMap{
		"name":  func() string { return env.User.Name },
		"email": func() string { return env.User.Email },

		"year": func() string { return time.Now().Format("2006") },
		"time": func(format ...string) string {
			if len(format) == 0 {
				format = append(format, time.RFC1123)
			}
			return time.Now().Format(format[0])
		},
		"date": func(format ...string) string {
			if len(format) == 0 {
				format = append(format, "Jan 02, 2006")
			}
			return time.Now().Format(format[0])
		},

		"import": func(pkgs ...string) string {
			if len(pkgs) == 0 {
				return `import ()`
			}
			if len(pkgs) == 1 {
				return `import "` + pkgs[0] + `"`
			}
			s := "import (\n"
			for _, pkg := range pkgs {
				s += "\t" + pkg + "\n"
			}
			s += ")"
			return s
		},
		"equal": func(v1, v2 interface{}) bool {
			return reflect.DeepEqual(reflect.ValueOf(v1), reflect.ValueOf(v2))
		},
	}
}

type options struct {
	env     string
	project string
	target  string
	pkg     string
	config  string
}

func parseOptionsV2() *options {
	opts := new(options)
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	fs.StringVar(&opts.env, "env", "", "specify a user environment")
	fs.StringVar(&opts.pkg, "pkg", "", "specify a package name")
	fs.StringVar(&opts.config, "config", "", "specify config path")
	fs.Parse(os.Args[1:])

	args := fs.Args()
	if len(args) < 2 {
		fmt.Fprintln(os.Stderr, "usage:", os.Args[0], "[options] project target")
		os.Exit(1)
	}
	opts.project, opts.target = args[0], args[1]
	if opts.pkg == "" {
		opts.pkg = opts.target
		if strings.HasPrefix(opts.pkg, "go-") {
			opts.pkg = opts.pkg[3:]
		}
		if strings.HasSuffix(opts.pkg, ".go") {
			opts.pkg = opts.pkg[:len(opts.pkg)-3]
		}
	}

	return opts
}

func initConfigV2(path string) (*config.GonewConfig2, error) {
	if path == "" {
		home := os.Getenv("HOME")
		path = filepath.Join(home, ".config", "gonew.json")
	}
	conf := new(config.GonewConfig2)
	err := conf.UnmarshalFileJSON(path)
	if err != nil {
		if err, ok := err.(*os.PathError); ok {
			if err.Err == syscall.ENOENT || err.Err == os.ErrNotExist {
				fmt.Fprintln(os.Stderr, "couldn't find config")

				bufr := bufio.NewReader(os.Stdin)
				fmt.Print("Enter your name: ")
				p, _, err := bufr.ReadLine()
				checkFatal(err)
				name := strings.TrimFunc(string(p), unicode.IsSpace)
				fmt.Print("Enter your email: ")
				p, _, err = bufr.ReadLine()
				checkFatal(err)
				email := strings.TrimFunc(string(p), unicode.IsSpace)
				fmt.Print("Base import path (e.g. github.com/bmatsuo): ")
				p, _, err = bufr.ReadLine()
				checkFatal(err)
				baseImportPath := strings.TrimFunc(string(p), unicode.IsSpace)
				// TODO locate gonew.json.example
				examplePath := "gonew.json.example"
				checkFatal(conf.UnmarshalFileJSON(examplePath))
				conf.Environments = config.EnvironmentsConfig{
					"norm": &config.EnvironmentConfig{
						BaseImportPath: baseImportPath,
						User: &config.EnvironmentUserConfig{
							Name:  name,
							Email: email,
						},
					},
				}

				err = conf.MarshalFileJSON(path)
				if err != nil {
					return nil, err
				}
			} else {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	return conf, nil
}

func mainv2() {

	// parse command line options/args
	opts := parseOptionsV2()
	projectName := opts.target
	packageName := opts.pkg
	envName := opts.env
	projType := opts.project

	// read the config file
	conf, err := initConfigV2(opts.config)
	checkFatal(err)

	// initialize project
	env, err := conf.Environment(envName)
	checkFatal(err)
	project.BaseImportPath = env.BaseImportPath
	proj := project.New(projectName, packageName, env)
	projContext := project.Context("", "", proj)
	projTemplEnv := templates.Env(projContext)
	projConfig, err := conf.Project(projType)
	checkFatal(err)

	// initialize template environment
	ts := templates.New(".t2")
	checkFatal(ts.Funcs(funcsV2(env)))

	// read templates
	src := templates.SourceDirectory(filepath.Join(GonewRoot,"templates"))
	checkFatal(ts.Source(src))
	for i := len(conf.ExternalTemplates) - 1; i >= 0; i-- {
		src := templates.SourceDirectory(conf.ExternalTemplates[i])
		checkFatal(ts.Source(src))
	}

	if projConfig.Hooks != nil {
		// fmt.Println("PRE")
		executeHooks(ts, projTemplEnv, projConfig.Hooks.Pre...)
	}

	// generate files. buffer all output then write.
	files := make([]*File2, 0, len(projConfig.Files))
	for _, file := range projConfig.Files {
		_relpath, err := projTemplEnv.RenderTextAsString(ts, "pre_", file.Path)
		checkFatal(err)
		relpath := string(_relpath)
		filename := filepath.Base(relpath)
		filetype := file.Type

		fileContext := project.Context(filename, filetype, proj)
		fileTemplEnv := templates.Env(fileContext)
		fileBuf := new(bytes.Buffer)
		for _, t := range file.Templates {
			if nil != check(fileTemplEnv.Render(fileBuf, ts, t)) {
				fileBuf = nil
				break
			}
		}

		if fileBuf != nil {
			f := &File2{relpath, fileBuf.Bytes()}
			files = append(files, f)
		} else {
			// TODO clean exit
		}
	}
	for _, file := range files {
		dir := filepath.Dir(file.path)

		// fmt.Println("mkdir", "-p", dir)
		err := os.MkdirAll(dir, 0755|os.ModeDir)
		checkFatal(err) // TODO clean exit

		// fmt.Println("cat", ">", file.path)
		// fmt.Println(string(file.content))
		handle, err := os.OpenFile(file.path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
		checkFatal(err) // TODO clean exit
		_, err = handle.Write(file.content)
		checkFatal(err) // TODO clean exit
		err = handle.Close()
	}

	if projConfig.Hooks != nil {
		// fmt.Println("POST")
		executeHooks(ts, projTemplEnv, projConfig.Hooks.Post...)
	}
}

func main() {
	mainv2()
}
<commit_msg>streamline FindGonew for GOPATH enforcement in go 1.1<commit_after>// Copyright 2011, Bryan Matsuo. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

/*
*  File: gonew_main.go
*  Author: Bryan Matsuo [bmatsuo@soe.ucsc.edu]
*  Created: Sat Jul  2 19:17:53 PDT 2011
*  Usage: gonew [options]
 */
import (
	"github.com/bmatsuo/gonew/config"
	"github.com/bmatsuo/gonew/project"
	"github.com/bmatsuo/gonew/templates"

	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"reflect"
	"strings"
	"syscall"
	"text/template"
	"time"
	"unicode"
)

// The directory containing Gonew's source code.
var GonewRoot string // The Gonew source directory
func FindGonew() error {
	gopath := os.Getenv("GOPATH")
	if gopath == "" {
		return fmt.Errorf("GOPATH is not set")
	}
	gopath = strings.SplitN(gopath, ":", 2)[0]
	GonewRoot = filepath.Join(gopath, "src", "github.com", "bmatsuo", "gonew")
	stat, err := os.Stat(GonewRoot)
	if err == nil && !stat.IsDir() {
		err = fmt.Errorf("file is not a directory: %s", GonewRoot)
	}
	return err
}

func check(err error) error {
	if err != nil {
		fmt.Println(err)
	}
	return err
}

func checkFatal(err error) {
	if err != nil {
		fmt.Println(err)
		os.Exit(1)
	}
}

func logJson(v ...interface{}) {
	w := make([]interface{}, 0, len(v))
	w = append(w, v[:len(v)-1]...)
	p, _ := json.MarshalIndent(v[len(v)-1], " ", "\t")
	w = append(w, string(p))
	fmt.Println(w...)
}

func executeHooks(ts templates.Interface, tenv templates.Environment, hooks ...*config.HookConfig) {
	for _, hook := range hooks {
		cwd, err := tenv.RenderTextAsString(ts, "cwd_", hook.Cwd)
		checkFatal(err)
		// fmt.Println("cd", cwd)
		for _, _cmd := range hook.Commands {
			cmd, err := tenv.RenderTextAsString(ts, "cmd_", _cmd)
			checkFatal(err)
			// fmt.Println("bash", "-c", cmd)
			shell := exec.Command("bash", "-c", cmd)
			shell.Dir = cwd
			shell.Stdin = os.Stdin
			shell.Stdout = os.Stdout
			shell.Stderr = os.Stderr
			checkFatal(shell.Run()) // TODO clean exit
		}
	}
}

type File2 struct {
	path    string
	content []byte
}

func funcsV2(env *config.EnvironmentConfig) template.FuncMap {
	return template.FuncMap{
		"name":  func() string { return env.User.Name },
		"email": func() string { return env.User.Email },

		"year": func() string { return time.Now().Format("2006") },
		"time": func(format ...string) string {
			if len(format) == 0 {
				format = append(format, time.RFC1123)
			}
			return time.Now().Format(format[0])
		},
		"date": func(format ...string) string {
			if len(format) == 0 {
				format = append(format, "Jan 02, 2006")
			}
			return time.Now().Format(format[0])
		},

		"import": func(pkgs ...string) string {
			if len(pkgs) == 0 {
				return `import ()`
			}
			if len(pkgs) == 1 {
				return `import "` + pkgs[0] + `"`
			}
			s := "import (\n"
			for _, pkg := range pkgs {
				s += "\t" + pkg + "\n"
			}
			s += ")"
			return s
		},
		"equal": func(v1, v2 interface{}) bool {
			return reflect.DeepEqual(reflect.ValueOf(v1), reflect.ValueOf(v2))
		},
	}
}

type options struct {
	env     string
	project string
	target  string
	pkg     string
	config  string
}

func parseOptionsV2() *options {
	opts := new(options)
	fs := flag.NewFlagSet(os.Args[0], flag.ExitOnError)
	fs.StringVar(&opts.env, "env", "", "specify a user environment")
	fs.StringVar(&opts.pkg, "pkg", "", "specify a package name")
	fs.StringVar(&opts.config, "config", "", "specify config path")
	fs.Parse(os.Args[1:])

	args := fs.Args()
	if len(args) < 2 {
		fmt.Fprintln(os.Stderr, "usage:", os.Args[0], "[options] project target")
		os.Exit(1)
	}
	opts.project, opts.target = args[0], args[1]
	if opts.pkg == "" {
		opts.pkg = opts.target
		if strings.HasPrefix(opts.pkg, "go-") {
			opts.pkg = opts.pkg[3:]
		}
		if strings.HasSuffix(opts.pkg, ".go") {
			opts.pkg = opts.pkg[:len(opts.pkg)-3]
		}
	}

	return opts
}

func initConfigV2(path string) (*config.GonewConfig2, error) {
	if path == "" {
		home := os.Getenv("HOME")
		path = filepath.Join(home, ".config", "gonew.json")
	}
	conf := new(config.GonewConfig2)
	err := conf.UnmarshalFileJSON(path)
	if err != nil {
		if err, ok := err.(*os.PathError); ok {
			if err.Err == syscall.ENOENT || err.Err == os.ErrNotExist {
				fmt.Fprintln(os.Stderr, "couldn't find config")

				bufr := bufio.NewReader(os.Stdin)
				fmt.Print("Enter your name: ")
				p, _, err := bufr.ReadLine()
				checkFatal(err)
				name := strings.TrimFunc(string(p), unicode.IsSpace)
				fmt.Print("Enter your email: ")
				p, _, err = bufr.ReadLine()
				checkFatal(err)
				email := strings.TrimFunc(string(p), unicode.IsSpace)
				fmt.Print("Base import path (e.g. github.com/bmatsuo): ")
				p, _, err = bufr.ReadLine()
				checkFatal(err)
				baseImportPath := strings.TrimFunc(string(p), unicode.IsSpace)
				// TODO locate gonew.json.example
				examplePath := "gonew.json.example"
				checkFatal(conf.UnmarshalFileJSON(examplePath))
				conf.Environments = config.EnvironmentsConfig{
					"norm": &config.EnvironmentConfig{
						BaseImportPath: baseImportPath,
						User: &config.EnvironmentUserConfig{
							Name:  name,
							Email: email,
						},
					},
				}

				err = conf.MarshalFileJSON(path)
				if err != nil {
					return nil, err
				}
			} else {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	return conf, nil
}

func main() {
	checkFatal(FindGonew())

	// parse command line options/args
	opts := parseOptionsV2()
	projectName := opts.target
	packageName := opts.pkg
	envName := opts.env
	projType := opts.project

	// read the config file
	conf, err := initConfigV2(opts.config)
	checkFatal(err)

	// initialize project
	env, err := conf.Environment(envName)
	checkFatal(err)
	project.BaseImportPath = env.BaseImportPath
	proj := project.New(projectName, packageName, env)
	projContext := project.Context("", "", proj)
	projTemplEnv := templates.Env(projContext)
	projConfig, err := conf.Project(projType)
	checkFatal(err)

	// initialize template environment
	ts := templates.New(".t2")
	checkFatal(ts.Funcs(funcsV2(env)))

	// read templates
	src := templates.SourceDirectory(filepath.Join(GonewRoot, "templates"))
	checkFatal(ts.Source(src))
	for i := len(conf.ExternalTemplates) - 1; i >= 0; i-- {
		src := templates.SourceDirectory(conf.ExternalTemplates[i])
		checkFatal(ts.Source(src))
	}

	if projConfig.Hooks != nil {
		// fmt.Println("PRE")
		executeHooks(ts, projTemplEnv, projConfig.Hooks.Pre...)
	}

	// generate files. buffer all output then write.
	files := make([]*File2, 0, len(projConfig.Files))
	for _, file := range projConfig.Files {
		_relpath, err := projTemplEnv.RenderTextAsString(ts, "pre_", file.Path)
		checkFatal(err)
		relpath := string(_relpath)
		filename := filepath.Base(relpath)
		filetype := file.Type

		fileContext := project.Context(filename, filetype, proj)
		fileTemplEnv := templates.Env(fileContext)
		fileBuf := new(bytes.Buffer)
		for _, t := range file.Templates {
			if nil != check(fileTemplEnv.Render(fileBuf, ts, t)) {
				fileBuf = nil
				break
			}
		}

		if fileBuf != nil {
			f := &File2{relpath, fileBuf.Bytes()}
			files = append(files, f)
		} else {
			// TODO clean exit
		}
	}
	for _, file := range files {
		dir := filepath.Dir(file.path)

		// fmt.Println("mkdir", "-p", dir)
		err := os.MkdirAll(dir, 0755|os.ModeDir)
		checkFatal(err) // TODO clean exit

		// fmt.Println("cat", ">", file.path)
		// fmt.Println(string(file.content))
		handle, err := os.OpenFile(file.path, os.O_WRONLY|os.O_CREATE|os.O_EXCL, 0644)
		checkFatal(err) // TODO clean exit
		_, err = handle.Write(file.content)
		checkFatal(err) // TODO clean exit
		err = handle.Close()
	}

	if projConfig.Hooks != nil {
		// fmt.Println("POST")
		executeHooks(ts, projTemplEnv, projConfig.Hooks.Post...)
	}
}
<|endoftext|>


<commit_before>package raftgorums

import (
	"fmt"
	"log"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/grpc"

	"github.com/relab/raft"
	"github.com/relab/raft/commonpb"
	gorums "github.com/relab/raft/raftgorums/gorumspb"
	pb "github.com/relab/raft/raftgorums/raftpb"
)

// Node ties an instance of Raft to the network.
type Node struct {
	id uint64

	Raft    *Raft
	storage Storage

	lookup map[uint64]int
	peers  []string

	catchingUp map[uint32]chan uint64
	catchUp    chan *catchUpRequest

	mgr  *gorums.Manager
	conf *gorums.Configuration
}

// NewNode returns a Node with an instance of Raft given the configuration.
func NewNode(server *grpc.Server, sm raft.StateMachine, cfg *Config) *Node {
	peers := make([]string, len(cfg.Nodes))
	// We don't want to mutate cfg.Nodes.
	copy(peers, cfg.Nodes)

	id := cfg.ID
	// Exclude self.
	peers = append(peers[:id-1], peers[id:]...)

	var pos int
	lookup := make(map[uint64]int)

	for i := 0; i < len(cfg.Nodes); i++ {
		if uint64(i) == id {
			continue
		}

		lookup[uint64(i)] = pos
		pos++
	}

	n := &Node{
		id:         id,
		Raft:       NewRaft(sm, cfg),
		storage:    cfg.Storage,
		lookup:     lookup,
		peers:      peers,
		catchingUp: make(map[uint32]chan uint64),
		catchUp:    make(chan *catchUpRequest),
	}

	gorums.RegisterRaftServer(server, n)

	return n
}

// Run start listening for incoming messages, and delivers outgoing messages.
func (n *Node) Run() error {
	opts := []gorums.ManagerOption{
		gorums.WithGrpcDialOptions(
			grpc.WithBlock(),
			grpc.WithInsecure(),
			grpc.WithTimeout(TCPConnect*time.Millisecond)),
	}

	mgr, err := gorums.NewManager(n.peers, opts...)

	if err != nil {
		return err
	}

	n.mgr = mgr
	n.conf, err = mgr.NewConfiguration(mgr.NodeIDs(), NewQuorumSpec(len(n.peers)+1))

	if err != nil {
		return err
	}

	go n.Raft.Run()

	for {
		rvreqout := n.Raft.RequestVoteRequestChan()
		aereqout := n.Raft.AppendEntriesRequestChan()
		sreqout := n.Raft.SnapshotRequestChan()

		select {
		case req := <-sreqout:
			ctx, cancel := context.WithTimeout(context.Background(), TCPConnect*time.Millisecond)
			node, _ := n.mgr.Node(n.getNodeID(req.FollowerID))
			snapshot, err := node.RaftClient.GetState(ctx, req)
			cancel()

			if err != nil {
				// TODO Better error message.
				log.Println(fmt.Sprintf("Snapshot request failed = %v", err))

			}

			n.Raft.restoreCh <- snapshot

		case req := <-rvreqout:
			ctx, cancel := context.WithTimeout(context.Background(), TCPHeartbeat*time.Millisecond)
			res, err := n.conf.RequestVote(ctx, req)
			cancel()

			if err != nil {
				// TODO Better error message.
				log.Println(fmt.Sprintf("RequestVote failed = %v", err))

			}

			if res.RequestVoteResponse == nil {
				continue
			}

			n.Raft.HandleRequestVoteResponse(res.RequestVoteResponse)

		case req := <-aereqout:
			ctx, cancel := context.WithTimeout(context.Background(), TCPHeartbeat*time.Millisecond)
			res, err := n.conf.AppendEntries(ctx, req)

			if err != nil {
				// TODO Better error message.
				log.Println(fmt.Sprintf("AppendEntries failed = %v", err))

				if res.AppendEntriesResponse == nil {
					continue
				}
			}

			// Cancel on abort.
			if !res.AppendEntriesResponse.Success {
				cancel()
			}

			n.Raft.HandleAppendEntriesResponse(res.AppendEntriesResponse, len(res.NodeIDs))

			for nodeID, matchIndex := range n.catchingUp {
				select {
				case index, ok := <-matchIndex:
					if !ok {
						n.addBackNode(nodeID)
						continue
					}

					matchIndex <- res.MatchIndex

					if index == res.MatchIndex {
						n.addBackNode(nodeID)
					}
				default:
				}
			}

		case creq := <-n.catchUp:
			oldSet := n.conf.NodeIDs()

			// Don't remove servers when we would've gone below the
			// quorum size. The quorum function handles recovery
			// when a majority fails.
			if len(oldSet)-1 < len(n.peers)/2 {
				continue
			}

			node := n.getNodeID(creq.followerID)
			// We use 2 peers as we need to count the leader.
			single, err := n.mgr.NewConfiguration([]uint32{node}, NewQuorumSpec(2))

			if err != nil {
				panic(fmt.Sprintf("tried to catch up node %d->%d: %v", creq.followerID, node, err))
			}

			tmpSet := make([]uint32, len(oldSet)-1)

			var i int
			for _, id := range oldSet {
				if id == node {
					continue
				}

				tmpSet[i] = id
				i++
			}

			// It's important not to change the quorum size when
			// removing the server. We reduce N by one so we don't
			// wait on the recovering server.
			n.conf, err = mgr.NewConfiguration(tmpSet, &QuorumSpec{
				N: len(tmpSet),
				Q: (len(n.peers) + 1) / 2,
			})

			if err != nil {
				panic(fmt.Sprintf("tried to create new configuration %v: %v", tmpSet, err))
			}

			matchIndex := make(chan uint64)
			n.catchingUp[node] = matchIndex
			go n.doCatchUp(single, creq.nextIndex, matchIndex)
		}
	}
}

func (n *Node) addBackNode(nodeID uint32) {
	n.Raft.Lock()
	n.Raft.catchingUp = false
	n.Raft.Unlock()
	delete(n.catchingUp, nodeID)

	newSet := append(n.conf.NodeIDs(), nodeID)
	var err error
	n.conf, err = n.mgr.NewConfiguration(newSet, &QuorumSpec{
		N: len(newSet),
		Q: (len(n.peers) + 1) / 2,
	})

	if err != nil {
		panic(fmt.Sprintf("tried to create new configuration %v: %v", newSet, err))
	}
}

// RequestVote implements gorums.RaftServer.
func (n *Node) RequestVote(ctx context.Context, req *pb.RequestVoteRequest) (*pb.RequestVoteResponse, error) {
	return n.Raft.HandleRequestVoteRequest(req), nil
}

// AppendEntries implements gorums.RaftServer.
func (n *Node) AppendEntries(ctx context.Context, req *pb.AppendEntriesRequest) (*pb.AppendEntriesResponse, error) {
	return n.Raft.HandleAppendEntriesRequest(req), nil
}

func (n *Node) GetState(ctx context.Context, req *pb.SnapshotRequest) (*commonpb.Snapshot, error) {
	future := make(chan *commonpb.Snapshot)
	n.Raft.snapCh <- future

	select {
	case snapshot := <-future:
		n.catchUp <- &catchUpRequest{
			nextIndex:  snapshot.Index,
			followerID: req.FollowerID,
		}
		return snapshot, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

type catchUpRequest struct {
	nextIndex  uint64
	followerID uint64
}

func (n *Node) doCatchUp(conf *gorums.Configuration, nextIndex uint64, matchIndex chan uint64) {
	for {
		state := n.Raft.State()

		if state != Leader {
			close(matchIndex)
			return
		}

		n.Raft.Lock()
		entries := n.Raft.getNextEntries(nextIndex)
		request := n.Raft.getAppendEntriesRequest(nextIndex, entries)
		n.Raft.Unlock()

		ctx, cancel := context.WithTimeout(context.Background(), TCPHeartbeat*time.Millisecond)
		res, err := conf.AppendEntries(ctx, request)

		if err != nil {
			// TODO Better error message.
			log.Println(fmt.Sprintf("Catch-up AppendEntries failed = %v", err))

			if res.AppendEntriesResponse == nil {
				continue
			}
		}

		cancel()

		response := res.AppendEntriesResponse

		if response.Success {
			matchIndex <- response.MatchIndex
			index := <-matchIndex

			if response.MatchIndex == index {
				close(matchIndex)
				return
			}

			nextIndex = response.MatchIndex + 1

			continue
		}

		// If AppendEntries was unsuccessful a new catch-up process will
		// start.
		close(matchIndex)
		return
	}
}

func (n *Node) getNodeID(raftID uint64) uint32 {
	return n.mgr.NodeIDs()[n.lookup[raftID]]
}
<commit_msg>raftgorums/node.go: Fix snapshot not being sent to leader<commit_after>package raftgorums

import (
	"fmt"
	"log"
	"time"

	"golang.org/x/net/context"
	"google.golang.org/grpc"

	"github.com/relab/raft"
	"github.com/relab/raft/commonpb"
	gorums "github.com/relab/raft/raftgorums/gorumspb"
	pb "github.com/relab/raft/raftgorums/raftpb"
)

// Node ties an instance of Raft to the network.
type Node struct {
	id uint64

	Raft    *Raft
	storage Storage

	lookup map[uint64]int
	peers  []string

	catchingUp map[uint32]chan uint64
	catchUp    chan *catchUpRequest

	mgr  *gorums.Manager
	conf *gorums.Configuration
}

// NewNode returns a Node with an instance of Raft given the configuration.
func NewNode(server *grpc.Server, sm raft.StateMachine, cfg *Config) *Node {
	peers := make([]string, len(cfg.Nodes))
	// We don't want to mutate cfg.Nodes.
	copy(peers, cfg.Nodes)

	id := cfg.ID
	// Exclude self.
	peers = append(peers[:id-1], peers[id:]...)

	var pos int
	lookup := make(map[uint64]int)

	for i := 0; i < len(cfg.Nodes); i++ {
		if uint64(i) == id {
			continue
		}

		lookup[uint64(i)] = pos
		pos++
	}

	n := &Node{
		id:         id,
		Raft:       NewRaft(sm, cfg),
		storage:    cfg.Storage,
		lookup:     lookup,
		peers:      peers,
		catchingUp: make(map[uint32]chan uint64),
		catchUp:    make(chan *catchUpRequest),
	}

	gorums.RegisterRaftServer(server, n)

	return n
}

// Run start listening for incoming messages, and delivers outgoing messages.
func (n *Node) Run() error {
	opts := []gorums.ManagerOption{
		gorums.WithGrpcDialOptions(
			grpc.WithBlock(),
			grpc.WithInsecure(),
			grpc.WithTimeout(TCPConnect*time.Millisecond)),
	}

	mgr, err := gorums.NewManager(n.peers, opts...)

	if err != nil {
		return err
	}

	n.mgr = mgr
	n.conf, err = mgr.NewConfiguration(mgr.NodeIDs(), NewQuorumSpec(len(n.peers)+1))

	if err != nil {
		return err
	}

	go n.Raft.Run()

	for {
		rvreqout := n.Raft.RequestVoteRequestChan()
		aereqout := n.Raft.AppendEntriesRequestChan()
		sreqout := n.Raft.SnapshotRequestChan()

		select {
		case req := <-sreqout:
			n.Raft.Lock()
			leader := n.Raft.leader
			n.Raft.Unlock()
			ctx, cancel := context.WithTimeout(context.Background(), TCPConnect*time.Millisecond)
			node, _ := n.mgr.Node(n.getNodeID(leader))
			snapshot, err := node.RaftClient.GetState(ctx, req)
			cancel()

			if err != nil {
				// TODO Better error message.
				log.Println(fmt.Sprintf("Snapshot request failed = %v", err))

			}

			n.Raft.restoreCh <- snapshot

		case req := <-rvreqout:
			ctx, cancel := context.WithTimeout(context.Background(), TCPHeartbeat*time.Millisecond)
			res, err := n.conf.RequestVote(ctx, req)
			cancel()

			if err != nil {
				// TODO Better error message.
				log.Println(fmt.Sprintf("RequestVote failed = %v", err))

			}

			if res.RequestVoteResponse == nil {
				continue
			}

			n.Raft.HandleRequestVoteResponse(res.RequestVoteResponse)

		case req := <-aereqout:
			ctx, cancel := context.WithTimeout(context.Background(), TCPHeartbeat*time.Millisecond)
			res, err := n.conf.AppendEntries(ctx, req)

			if err != nil {
				// TODO Better error message.
				log.Println(fmt.Sprintf("AppendEntries failed = %v", err))

				if res.AppendEntriesResponse == nil {
					continue
				}
			}

			// Cancel on abort.
			if !res.AppendEntriesResponse.Success {
				cancel()
			}

			n.Raft.HandleAppendEntriesResponse(res.AppendEntriesResponse, len(res.NodeIDs))

			for nodeID, matchIndex := range n.catchingUp {
				select {
				case index, ok := <-matchIndex:
					if !ok {
						n.addBackNode(nodeID)
						continue
					}

					matchIndex <- res.MatchIndex

					if index == res.MatchIndex {
						n.addBackNode(nodeID)
					}
				default:
				}
			}

		case creq := <-n.catchUp:
			oldSet := n.conf.NodeIDs()

			// Don't remove servers when we would've gone below the
			// quorum size. The quorum function handles recovery
			// when a majority fails.
			if len(oldSet)-1 < len(n.peers)/2 {
				continue
			}

			node := n.getNodeID(creq.followerID)
			// We use 2 peers as we need to count the leader.
			single, err := n.mgr.NewConfiguration([]uint32{node}, NewQuorumSpec(2))

			if err != nil {
				panic(fmt.Sprintf("tried to catch up node %d->%d: %v", creq.followerID, node, err))
			}

			tmpSet := make([]uint32, len(oldSet)-1)

			var i int
			for _, id := range oldSet {
				if id == node {
					continue
				}

				tmpSet[i] = id
				i++
			}

			// It's important not to change the quorum size when
			// removing the server. We reduce N by one so we don't
			// wait on the recovering server.
			n.conf, err = mgr.NewConfiguration(tmpSet, &QuorumSpec{
				N: len(tmpSet),
				Q: (len(n.peers) + 1) / 2,
			})

			if err != nil {
				panic(fmt.Sprintf("tried to create new configuration %v: %v", tmpSet, err))
			}

			matchIndex := make(chan uint64)
			n.catchingUp[node] = matchIndex
			go n.doCatchUp(single, creq.nextIndex, matchIndex)
		}
	}
}

func (n *Node) addBackNode(nodeID uint32) {
	n.Raft.Lock()
	n.Raft.catchingUp = false
	n.Raft.Unlock()
	delete(n.catchingUp, nodeID)

	newSet := append(n.conf.NodeIDs(), nodeID)
	var err error
	n.conf, err = n.mgr.NewConfiguration(newSet, &QuorumSpec{
		N: len(newSet),
		Q: (len(n.peers) + 1) / 2,
	})

	if err != nil {
		panic(fmt.Sprintf("tried to create new configuration %v: %v", newSet, err))
	}
}

// RequestVote implements gorums.RaftServer.
func (n *Node) RequestVote(ctx context.Context, req *pb.RequestVoteRequest) (*pb.RequestVoteResponse, error) {
	return n.Raft.HandleRequestVoteRequest(req), nil
}

// AppendEntries implements gorums.RaftServer.
func (n *Node) AppendEntries(ctx context.Context, req *pb.AppendEntriesRequest) (*pb.AppendEntriesResponse, error) {
	return n.Raft.HandleAppendEntriesRequest(req), nil
}

func (n *Node) GetState(ctx context.Context, req *pb.SnapshotRequest) (*commonpb.Snapshot, error) {
	future := make(chan *commonpb.Snapshot)
	n.Raft.snapCh <- future

	select {
	case snapshot := <-future:
		n.catchUp <- &catchUpRequest{
			nextIndex:  snapshot.Index,
			followerID: req.FollowerID,
		}
		return snapshot, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

type catchUpRequest struct {
	nextIndex  uint64
	followerID uint64
}

func (n *Node) doCatchUp(conf *gorums.Configuration, nextIndex uint64, matchIndex chan uint64) {
	for {
		state := n.Raft.State()

		if state != Leader {
			close(matchIndex)
			return
		}

		n.Raft.Lock()
		entries := n.Raft.getNextEntries(nextIndex)
		request := n.Raft.getAppendEntriesRequest(nextIndex, entries)
		n.Raft.Unlock()

		ctx, cancel := context.WithTimeout(context.Background(), TCPHeartbeat*time.Millisecond)
		res, err := conf.AppendEntries(ctx, request)

		if err != nil {
			// TODO Better error message.
			log.Println(fmt.Sprintf("Catch-up AppendEntries failed = %v", err))

			if res.AppendEntriesResponse == nil {
				continue
			}
		}

		cancel()

		response := res.AppendEntriesResponse

		if response.Success {
			matchIndex <- response.MatchIndex
			index := <-matchIndex

			if response.MatchIndex == index {
				close(matchIndex)
				return
			}

			nextIndex = response.MatchIndex + 1

			continue
		}

		// If AppendEntries was unsuccessful a new catch-up process will
		// start.
		close(matchIndex)
		return
	}
}

func (n *Node) getNodeID(raftID uint64) uint32 {
	return n.mgr.NodeIDs()[n.lookup[raftID-1]]
}
<|endoftext|>


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

// Package google provides support for making OAuth2 authorized and authenticated
// HTTP requests to Google APIs. It supports the Web server flow, client-side
// credentials, service accounts, Google Compute Engine service accounts,
// Google App Engine service accounts and workload identity federation
// from non-Google cloud platforms.
//
// A brief overview of the package follows. For more information, please read
// https://developers.google.com/accounts/docs/OAuth2
// and
// https://developers.google.com/accounts/docs/application-default-credentials.
// For more information on using workload identity federation, refer to
// https://cloud.google.com/iam/docs/how-to#using-workload-identity-federation.
//
// # OAuth2 Configs
//
// Two functions in this package return golang.org/x/oauth2.Config values from Google credential
// data. Google supports two JSON formats for OAuth2 credentials: one is handled by ConfigFromJSON,
// the other by JWTConfigFromJSON. The returned Config can be used to obtain a TokenSource or
// create an http.Client.
//
// # Workload Identity Federation
//
// Using workload identity federation, your application can access Google Cloud
// resources from Amazon Web Services (AWS), Microsoft Azure or any identity
// provider that supports OpenID Connect (OIDC).
// Traditionally, applications running outside Google Cloud have used service
// account keys to access Google Cloud resources. Using identity federation,
// you can allow your workload to impersonate a service account.
// This lets you access Google Cloud resources directly, eliminating the
// maintenance and security burden associated with service account keys.
//
// Follow the detailed instructions on how to configure Workload Identity Federation
// in various platforms:
//
//	Amazon Web Services (AWS): https://cloud.google.com/iam/docs/access-resources-aws
//	Microsoft Azure: https://cloud.google.com/iam/docs/access-resources-azure
//	OIDC identity provider: https://cloud.google.com/iam/docs/access-resources-oidc
//
// For OIDC providers, the library can retrieve OIDC tokens either from a
// local file location (file-sourced credentials) or from a local server
// (URL-sourced credentials).
// For file-sourced credentials, a background process needs to be continuously
// refreshing the file location with a new OIDC token prior to expiration.
// For tokens with one hour lifetimes, the token needs to be updated in the file
// every hour. The token can be stored directly as plain text or in JSON format.
// For URL-sourced credentials, a local server needs to host a GET endpoint to
// return the OIDC token. The response can be in plain text or JSON.
// Additional required request headers can also be specified.
//
// # Credentials
//
// The Credentials type represents Google credentials, including Application Default
// Credentials.
//
// Use FindDefaultCredentials to obtain Application Default Credentials.
// FindDefaultCredentials looks in some well-known places for a credentials file, and
// will call AppEngineTokenSource or ComputeTokenSource as needed.
//
// Application Default Credentials also support workload identity federation to
// access Google Cloud resources from non-Google Cloud platforms including Amazon
// Web Services (AWS), Microsoft Azure or any identity provider that supports
// OpenID Connect (OIDC). Workload identity federation is recommended for
// non-Google Cloud environments as it avoids the need to download, manage and
// store service account private keys locally.
//
// DefaultClient and DefaultTokenSource are convenience methods. They first call FindDefaultCredentials,
// then use the credentials to construct an http.Client or an oauth2.TokenSource.
//
// Use CredentialsFromJSON to obtain credentials from either of the two JSON formats
// described in OAuth2 Configs, above. The TokenSource in the returned value is the
// same as the one obtained from the oauth2.Config returned from ConfigFromJSON or
// JWTConfigFromJSON, but the Credentials may contain additional information
// that is useful is some circumstances.
package google // import "golang.org/x/oauth2/google"
<commit_msg>google/google: update documentation for workload identity federation<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package google provides support for making OAuth2 authorized and authenticated
// HTTP requests to Google APIs. It supports the Web server flow, client-side
// credentials, service accounts, Google Compute Engine service accounts,
// Google App Engine service accounts and workload identity federation
// from non-Google cloud platforms.
//
// A brief overview of the package follows. For more information, please read
// https://developers.google.com/accounts/docs/OAuth2
// and
// https://developers.google.com/accounts/docs/application-default-credentials.
// For more information on using workload identity federation, refer to
// https://cloud.google.com/iam/docs/how-to#using-workload-identity-federation.
//
// # OAuth2 Configs
//
// Two functions in this package return golang.org/x/oauth2.Config values from Google credential
// data. Google supports two JSON formats for OAuth2 credentials: one is handled by ConfigFromJSON,
// the other by JWTConfigFromJSON. The returned Config can be used to obtain a TokenSource or
// create an http.Client.
//
// # Workload Identity Federation
//
// Using workload identity federation, your application can access Google Cloud
// resources from Amazon Web Services (AWS), Microsoft Azure or any identity
// provider that supports OpenID Connect (OIDC).
// Traditionally, applications running outside Google Cloud have used service
// account keys to access Google Cloud resources. Using identity federation,
// you can allow your workload to impersonate a service account.
// This lets you access Google Cloud resources directly, eliminating the
// maintenance and security burden associated with service account keys.
//
// Follow the detailed instructions on how to configure Workload Identity Federation
// in various platforms:
//
//	Amazon Web Services (AWS): https://cloud.google.com/iam/docs/access-resources-aws
//	Microsoft Azure: https://cloud.google.com/iam/docs/access-resources-azure
//	OIDC identity provider: https://cloud.google.com/iam/docs/access-resources-oidc
//
// For OIDC and SAML providers, the library can retrieve tokens in three ways:
// from a local file location (file-sourced credentials), from a server
// (URL-sourced credentials), or from a local executable (executable-sourced
// credentials).
// For file-sourced credentials, a background process needs to be continuously
// refreshing the file location with a new OIDC token prior to expiration.
// For tokens with one hour lifetimes, the token needs to be updated in the file
// every hour. The token can be stored directly as plain text or in JSON format.
// For URL-sourced credentials, a local server needs to host a GET endpoint to
// return the OIDC token. The response can be in plain text or JSON.
// Additional required request headers can also be specified.
// For executable-sourced credentials, an application needs to be available to
// output the OIDC token and other information in a JSON format.
// For more information on how these work (and how to implement
// executable-sourced credentials), please check out:
// https://cloud.google.com/iam/docs/using-workload-identity-federation#oidc
//
// # Credentials
//
// The Credentials type represents Google credentials, including Application Default
// Credentials.
//
// Use FindDefaultCredentials to obtain Application Default Credentials.
// FindDefaultCredentials looks in some well-known places for a credentials file, and
// will call AppEngineTokenSource or ComputeTokenSource as needed.
//
// Application Default Credentials also support workload identity federation to
// access Google Cloud resources from non-Google Cloud platforms including Amazon
// Web Services (AWS), Microsoft Azure or any identity provider that supports
// OpenID Connect (OIDC). Workload identity federation is recommended for
// non-Google Cloud environments as it avoids the need to download, manage and
// store service account private keys locally.
//
// DefaultClient and DefaultTokenSource are convenience methods. They first call FindDefaultCredentials,
// then use the credentials to construct an http.Client or an oauth2.TokenSource.
//
// Use CredentialsFromJSON to obtain credentials from either of the two JSON formats
// described in OAuth2 Configs, above. The TokenSource in the returned value is the
// same as the one obtained from the oauth2.Config returned from ConfigFromJSON or
// JWTConfigFromJSON, but the Credentials may contain additional information
// that is useful is some circumstances.
package google // import "golang.org/x/oauth2/google"
<|endoftext|>


<commit_before>// package grain provides helpers for acting as a sandstorm grain.
package grain // import "zenhack.net/go/sandstorm/grain"

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

import (
	"golang.org/x/net/context"
	"net"
	"os"
	capnp "zenhack.net/go/sandstorm/capnp/grain"
	"zombiezen.com/go/capnproto2/rpc"
)

// Connect to the sandstorm API, give it a capability to uiview and return the api object.
//
// This assumes that file descriptor #3 is a socket with the API on the other end. This
// will be the case if the program was launched directly as a sandstorm app.
//
// TODO: explain error cases.
func ConnectAPI(ctx context.Context, uiview capnp.UiView_Server) (capnp.SandstormApi, error) {
	file := os.NewFile(3, "<sandstorm-api-socket>")
	conn, err := net.FileConn(file)
	if err != nil {
		return capnp.SandstormApi{}, err
	}
	uiviewClient := capnp.UiView_ServerToClient(uiview)
	transport := rpc.StreamTransport(conn)
	client := rpc.NewConn(transport, rpc.MainInterface(uiviewClient.Client)).Bootstrap(ctx)
	return capnp.SandstormApi{Client: client}, nil
}

// A HasSessionContext has a SessionContext associated with it, which may be retrieved by
// calling GetSessionContext().
type HasSessionContext interface {
	GetSessionContext() capnp.SessionContext
}
<commit_msg>grain: replace HasSessionContext interface<commit_after>// package grain provides helpers for acting as a sandstorm grain.
package grain // import "zenhack.net/go/sandstorm/grain"

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

import (
	"golang.org/x/net/context"
	"net"
	"os"
	capnp "zenhack.net/go/sandstorm/capnp/grain"
	"zombiezen.com/go/capnproto2/rpc"
)

type ctxKey int

const (
	sessionCtxKey ctxKey = iota
)

// Connect to the sandstorm API, give it a capability to uiview and return the api object.
//
// This assumes that file descriptor #3 is a socket with the API on the other end. This
// will be the case if the program was launched directly as a sandstorm app.
//
// TODO: explain error cases.
func ConnectAPI(ctx context.Context, uiview capnp.UiView_Server) (capnp.SandstormApi, error) {
	file := os.NewFile(3, "<sandstorm-api-socket>")
	conn, err := net.FileConn(file)
	if err != nil {
		return capnp.SandstormApi{}, err
	}
	uiviewClient := capnp.UiView_ServerToClient(uiview)
	transport := rpc.StreamTransport(conn)
	client := rpc.NewConn(transport, rpc.MainInterface(uiviewClient.Client)).Bootstrap(ctx)
	return capnp.SandstormApi{Client: client}, nil
}

// Add `sessionCtx` to the values of `parent`, returning the new context.
//
// `sessionCtx` can be retrieved using `GetSessionContext`, below.
func WithSessionContext(parent context.Context, sessionCtx capnp.SessionContext) context.Context {
	return context.WithValue(parent, sessionCtxKey, sessionCtx)
}

// Retrieve the SessionContext from `ctx`'s values, which must have been previously set via
// `WithSessionContext`.
//
// panics if there is no SessionContext.
func GetSessionContext(ctx context.Context) capnp.SessionContext {
	val := ctx.Value(sessionCtxKey)
	if val == nil {
		panic("No Session Context")
	}
	return val.(capnp.SessionContext)
}
<|endoftext|>


<commit_before>// TODO: better name? network topology?
package graph

import (
	"container/ring"
	"crypto/md5"
	"encoding/hex"
	"io"
	"sync"

	"github.com/Sirupsen/logrus"
)

type NetworkGraph struct {
	// nodeName -> Node
	NodesMap  map[string]*NetworkNode `json:"nodes"`
	NodesLock *sync.RWMutex           `json:"-"`

	// nodeName,nodeName -> NetworkLink
	LinksMap  map[string]*NetworkLink `json:"edges"`
	LinksLock *sync.RWMutex           `json:"-"`

	RoutesMap  map[string]*NetworkRoute `json:"routes"`
	RoutesLock *sync.RWMutex            `json:"-"`

	// event stuff
	eventChannels     map[chan *Event]bool
	eventRegistration chan chan *Event
	internalEvents    chan *Event
}

func Create() *NetworkGraph {
	g := &NetworkGraph{
		NodesMap:   make(map[string]*NetworkNode),
		NodesLock:  &sync.RWMutex{},
		LinksMap:   make(map[string]*NetworkLink),
		LinksLock:  &sync.RWMutex{},
		RoutesMap:  make(map[string]*NetworkRoute),
		RoutesLock: &sync.RWMutex{},

		eventChannels:     make(map[chan *Event]bool),
		eventRegistration: make(chan chan *Event),
		internalEvents:    make(chan *Event),
	}

	go g.publisher()

	return g
}

// TODO: buffer messages?
// goroutine target to do all the publishing of events
func (g *NetworkGraph) publisher() {
	for {
		select {
		case newChannel := <-g.eventRegistration:
			g.eventChannels[newChannel] = true
		case newEvent := <-g.internalEvents:
			//logrus.Infof("got a new event! %v", newEvent)
			for subscriberChannel := range g.eventChannels {
				select {
				case subscriberChannel <- newEvent:
					//logrus.Infof("sent that event to a channel")
				default:
					//logrus.Infof("Unable to send event to that subscriber, killing")
					delete(g.eventChannels, subscriberChannel)
					close(subscriberChannel)
				}
			}
		}
	}
}

// TODO: locking here
// Dump everything in the NetworkGraph into a channel
func (g *NetworkGraph) EventDumpChannel() chan *Event {
	// TODO: buffer?
	c := make(chan *Event)
	go func(c chan *Event) {
		// dump nodes
		for _, n := range g.NodesMap {
			c <- &Event{
				E:    addEvent,
				Item: n,
			}
		}
		// dump links
		for _, l := range g.LinksMap {
			c <- &Event{
				E:    addEvent,
				Item: l,
			}
		}
		// dump routes
		for _, route := range g.RoutesMap {
			c <- &Event{
				E:    addEvent,
				Item: route,
			}
		}

		close(c)
	}(c)
	return c
}

// add subscriber to our events
func (g *NetworkGraph) Subscribe(c chan *Event) {
	g.eventRegistration <- c
}

func (g *NetworkGraph) IncrNode(name string, newNode *NetworkNode) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]
	// if this one doesn't exist, lets add it
	if !ok {
		if newNode == nil {
			n = NewNetworkNode(name, g.internalEvents)
		} else {
			n = newNode
		}
		g.NodesMap[name] = n

		// TODO: there is an obvious race here-- the NewNetworkNode() call spawns
		// a gouroutine in the background-- we need to either hold the event (ordering problems)
		// or send an update event after the fact
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: n,
		}
	}
	n.refCount++
	return n, !ok
}

func (g *NetworkGraph) GetNode(name string) *NetworkNode {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	n, _ := g.NodesMap[name]
	return n
}

func (g *NetworkGraph) GetNodeCount() int {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	return len(g.NodesMap)
}

func (g *NetworkGraph) DecrNode(name string) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]

	if !ok {
		logrus.Warningf("Attempted to remove node with ip %v which wasn't in the graph", name)
		return nil, false
	}

	n.refCount--
	if n.refCount == 0 {
		delete(g.NodesMap, name)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: n,
		}
		return n, true
	}
	return nil, false
}

func (g *NetworkGraph) IncrLink(src, dst string, newLink *NetworkLink) (*NetworkLink, bool) {
	key := src + "," + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		if newLink == nil {
			srcNode, _ := g.IncrNode(src, nil)
			dstNode, _ := g.IncrNode(dst, nil)
			l = &NetworkLink{
				Src: srcNode,
				Dst: dstNode,
			}
		} else {
			srcNode, _ := g.IncrNode(src, newLink.Src)
			dstNode, _ := g.IncrNode(dst, newLink.Dst)
			// update child pointers
			newLink.Src = srcNode
			newLink.Dst = dstNode
			l = newLink
		}
		g.LinksMap[key] = l
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: l,
		}
	}
	l.refCount++
	return l, !ok
}

func (g *NetworkGraph) GetLink(src, dst string) *NetworkLink {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	key := src + "," + dst
	l, _ := g.LinksMap[key]
	return l
}

func (g *NetworkGraph) GetLinkCount() int {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	return len(g.LinksMap)
}

func (g *NetworkGraph) DecrLink(src, dst string) (*NetworkLink, bool) {
	key := src + "," + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove link %v which wasn't in the graph", key)
		return nil, false
	}
	// decrement ourselves
	l.refCount--
	if l.refCount == 0 {
		// Decrement our children
		g.DecrNode(src)
		g.DecrNode(dst)
		delete(g.LinksMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: l,
		}
		return l, true
	}
	return nil, false
}

func (g *NetworkGraph) pathKey(hops []string) string {
	h := md5.New()
	for _, hop := range hops {
		io.WriteString(h, hop)
	}
	return hex.EncodeToString(h.Sum(nil))
}

func (g *NetworkGraph) IncrRoute(hops []string, newRoute *NetworkRoute) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()

	// check if we have a route for this already
	route, ok := g.RoutesMap[key]
	// if we don't have it, lets make it
	if !ok {
		logrus.Debugf("New Route: key=%s %v", key, hops)
		if newRoute == nil {
			path := make([]*NetworkNode, 0, len(hops))
			for i, hop := range hops {
				hopNode, _ := g.IncrNode(hop, nil)
				path = append(path, hopNode)
				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(hops[i-1], hop, nil)
				}
			}
			route = &NetworkRoute{
				Path:       path,
				State:      Up,
				metricRing: ring.New(100), // TODO: config
				mLock:      &sync.RWMutex{},
				updateChan: g.internalEvents,
			}
		} else {
			for i, node := range newRoute.Path {
				node, _ := g.IncrNode(node.Name, node)
				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(newRoute.Path[i-1].Name, node.Name, nil)
				}
				// re-set the node in the path, in case its not the same pointer
				newRoute.Path[i] = node
			}
			route = newRoute
		}
		g.RoutesMap[key] = route

		g.internalEvents <- &Event{
			E:    addEvent,
			Item: route,
		}
	}

	// increment route's refcount
	route.refCount++

	return route, !ok
}

func (g *NetworkGraph) GetRoute(hops []string) *NetworkRoute {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	r, _ := g.RoutesMap[g.pathKey(hops)]
	return r
}

func (g *NetworkGraph) GetRouteCount() int {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	return len(g.RoutesMap)
}

func (g *NetworkGraph) DecrRoute(hops []string) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()
	r, ok := g.RoutesMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove route %v which wasn't in the graph", key)
		return nil, false
	}

	r.refCount--
	if r.refCount == 0 {
		// decrement all the links/nodes as well
		for i, node := range r.Path {
			g.DecrNode(node.Name)
			if i-1 >= 0 {
				g.DecrLink(r.Path[i-1].Name, node.Name)
			}
		}

		delete(g.RoutesMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: r,
		}
		return r, true
	}
	return r, false
}
<commit_msg>Properly set the update channels-- this will be important as states change and updates are needed<commit_after>// TODO: better name? network topology?
package graph

import (
	"container/ring"
	"crypto/md5"
	"encoding/hex"
	"io"
	"sync"

	"github.com/Sirupsen/logrus"
)

type NetworkGraph struct {
	// nodeName -> Node
	NodesMap  map[string]*NetworkNode `json:"nodes"`
	NodesLock *sync.RWMutex           `json:"-"`

	// nodeName,nodeName -> NetworkLink
	LinksMap  map[string]*NetworkLink `json:"edges"`
	LinksLock *sync.RWMutex           `json:"-"`

	RoutesMap  map[string]*NetworkRoute `json:"routes"`
	RoutesLock *sync.RWMutex            `json:"-"`

	// event stuff
	eventChannels     map[chan *Event]bool
	eventRegistration chan chan *Event
	internalEvents    chan *Event
}

func Create() *NetworkGraph {
	g := &NetworkGraph{
		NodesMap:   make(map[string]*NetworkNode),
		NodesLock:  &sync.RWMutex{},
		LinksMap:   make(map[string]*NetworkLink),
		LinksLock:  &sync.RWMutex{},
		RoutesMap:  make(map[string]*NetworkRoute),
		RoutesLock: &sync.RWMutex{},

		eventChannels:     make(map[chan *Event]bool),
		eventRegistration: make(chan chan *Event),
		internalEvents:    make(chan *Event),
	}

	go g.publisher()

	return g
}

// TODO: buffer messages?
// goroutine target to do all the publishing of events
func (g *NetworkGraph) publisher() {
	for {
		select {
		case newChannel := <-g.eventRegistration:
			g.eventChannels[newChannel] = true
		case newEvent := <-g.internalEvents:
			//logrus.Infof("got a new event! %v", newEvent)
			for subscriberChannel := range g.eventChannels {
				select {
				case subscriberChannel <- newEvent:
					//logrus.Infof("sent that event to a channel")
				default:
					//logrus.Infof("Unable to send event to that subscriber, killing")
					delete(g.eventChannels, subscriberChannel)
					close(subscriberChannel)
				}
			}
		}
	}
}

// TODO: locking here
// Dump everything in the NetworkGraph into a channel
func (g *NetworkGraph) EventDumpChannel() chan *Event {
	// TODO: buffer?
	c := make(chan *Event)
	go func(c chan *Event) {
		// dump nodes
		for _, n := range g.NodesMap {
			c <- &Event{
				E:    addEvent,
				Item: n,
			}
		}
		// dump links
		for _, l := range g.LinksMap {
			c <- &Event{
				E:    addEvent,
				Item: l,
			}
		}
		// dump routes
		for _, route := range g.RoutesMap {
			c <- &Event{
				E:    addEvent,
				Item: route,
			}
		}

		close(c)
	}(c)
	return c
}

// add subscriber to our events
func (g *NetworkGraph) Subscribe(c chan *Event) {
	g.eventRegistration <- c
}

func (g *NetworkGraph) IncrNode(name string, newNode *NetworkNode) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]
	// if this one doesn't exist, lets add it
	if !ok {
		if newNode == nil {
			n = NewNetworkNode(name, g.internalEvents)
		} else {
			n = newNode
			n.updateChan = g.internalEvents
		}
		g.NodesMap[name] = n

		// Now that there is a new thing we fire an addEvent.
		// Note: if the background DNS lookup was initiated an updateEvent
		// will fire as soon as the lookup completes
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: n,
		}
	}
	n.refCount++
	return n, !ok
}

func (g *NetworkGraph) GetNode(name string) *NetworkNode {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	n, _ := g.NodesMap[name]
	return n
}

func (g *NetworkGraph) GetNodeCount() int {
	g.NodesLock.RLock()
	defer g.NodesLock.RUnlock()
	return len(g.NodesMap)
}

func (g *NetworkGraph) DecrNode(name string) (*NetworkNode, bool) {
	g.NodesLock.Lock()
	defer g.NodesLock.Unlock()
	n, ok := g.NodesMap[name]

	if !ok {
		logrus.Warningf("Attempted to remove node with ip %v which wasn't in the graph", name)
		return nil, false
	}

	n.refCount--
	if n.refCount == 0 {
		delete(g.NodesMap, name)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: n,
		}
		return n, true
	}
	return nil, false
}

func (g *NetworkGraph) IncrLink(src, dst string, newLink *NetworkLink) (*NetworkLink, bool) {
	key := src + "," + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		if newLink == nil {
			srcNode, _ := g.IncrNode(src, nil)
			dstNode, _ := g.IncrNode(dst, nil)
			l = &NetworkLink{
				Src: srcNode,
				Dst: dstNode,
			}
		} else {
			srcNode, _ := g.IncrNode(src, newLink.Src)
			dstNode, _ := g.IncrNode(dst, newLink.Dst)
			// update child pointers
			newLink.Src = srcNode
			newLink.Dst = dstNode
			l = newLink
		}
		g.LinksMap[key] = l
		g.internalEvents <- &Event{
			E:    addEvent,
			Item: l,
		}
	}
	l.refCount++
	return l, !ok
}

func (g *NetworkGraph) GetLink(src, dst string) *NetworkLink {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	key := src + "," + dst
	l, _ := g.LinksMap[key]
	return l
}

func (g *NetworkGraph) GetLinkCount() int {
	g.LinksLock.RLock()
	defer g.LinksLock.RUnlock()
	return len(g.LinksMap)
}

func (g *NetworkGraph) DecrLink(src, dst string) (*NetworkLink, bool) {
	key := src + "," + dst
	g.LinksLock.Lock()
	defer g.LinksLock.Unlock()
	l, ok := g.LinksMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove link %v which wasn't in the graph", key)
		return nil, false
	}
	// decrement ourselves
	l.refCount--
	if l.refCount == 0 {
		// Decrement our children
		g.DecrNode(src)
		g.DecrNode(dst)
		delete(g.LinksMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: l,
		}
		return l, true
	}
	return nil, false
}

func (g *NetworkGraph) pathKey(hops []string) string {
	h := md5.New()
	for _, hop := range hops {
		io.WriteString(h, hop)
	}
	return hex.EncodeToString(h.Sum(nil))
}

func (g *NetworkGraph) IncrRoute(hops []string, newRoute *NetworkRoute) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()

	// check if we have a route for this already
	route, ok := g.RoutesMap[key]
	// if we don't have it, lets make it
	if !ok {
		logrus.Debugf("New Route: key=%s %v", key, hops)
		if newRoute == nil {
			path := make([]*NetworkNode, 0, len(hops))
			for i, hop := range hops {
				hopNode, _ := g.IncrNode(hop, nil)
				path = append(path, hopNode)
				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(hops[i-1], hop, nil)
				}
			}
			route = &NetworkRoute{
				Path:       path,
				State:      Up,
				metricRing: ring.New(100), // TODO: config
				mLock:      &sync.RWMutex{},
				updateChan: g.internalEvents,
			}
		} else {
			for i, node := range newRoute.Path {
				node, _ := g.IncrNode(node.Name, node)
				// If there was something prior-- lets add the link as well
				if i-1 >= 0 {
					g.IncrLink(newRoute.Path[i-1].Name, node.Name, nil)
				}
				// re-set the node in the path, in case its not the same pointer
				newRoute.Path[i] = node
			}
			route = newRoute
			route.updateChan = g.internalEvents
		}
		g.RoutesMap[key] = route

		g.internalEvents <- &Event{
			E:    addEvent,
			Item: route,
		}
	}

	// increment route's refcount
	route.refCount++

	return route, !ok
}

func (g *NetworkGraph) GetRoute(hops []string) *NetworkRoute {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	r, _ := g.RoutesMap[g.pathKey(hops)]
	return r
}

func (g *NetworkGraph) GetRouteCount() int {
	g.RoutesLock.RLock()
	defer g.RoutesLock.RUnlock()
	return len(g.RoutesMap)
}

func (g *NetworkGraph) DecrRoute(hops []string) (*NetworkRoute, bool) {
	key := g.pathKey(hops)

	g.RoutesLock.Lock()
	defer g.RoutesLock.Unlock()
	r, ok := g.RoutesMap[key]
	if !ok {
		logrus.Warningf("Attempted to remove route %v which wasn't in the graph", key)
		return nil, false
	}

	r.refCount--
	if r.refCount == 0 {
		// decrement all the links/nodes as well
		for i, node := range r.Path {
			g.DecrNode(node.Name)
			if i-1 >= 0 {
				g.DecrLink(r.Path[i-1].Name, node.Name)
			}
		}

		delete(g.RoutesMap, key)
		g.internalEvents <- &Event{
			E:    removeEvent,
			Item: r,
		}
		return r, true
	}
	return r, false
}
<|endoftext|>


<commit_before>package gotcp

import (
	"bytes"
	"encoding/binary"
	"net"
	"reflect"
	"time"
)

type HandleMessageFunc func(task *Task, data []byte)

type HanldeMessageMap [255][256]HandleMessageFunc

type ReadData struct {
	Length int
	Data   []byte
}
type Task struct {
	Entry
	Conn            *net.TCPConn
	in              chan []ReadData
	StopChan        chan bool
	handleReadFun   func(task *Task) ([]ReadData, error)
	handleWriteFun  func(task *Task, data []byte) error
	handleParse     func(task *Task, data []byte) bool
	handleHeartBeat func(task *Task)
	heartBeatTime   time.Duration
	HeartBeatReturn bool
	handleMessage   *HanldeMessageMap
}

func NewTask(c *net.TCPConn, name string) *Task {
	task := &Task{
		Conn:            c,
		in:              make(chan []ReadData),
		StopChan:        make(chan bool),
		handleReadFun:   handleReadFunDefault,
		handleWriteFun:  handleWriteFunDefault,
		handleParse:     handleParseDefault,
		handleHeartBeat: handleHeartBeatRequestDefault,
		heartBeatTime:   time.Second * 10,
		HeartBeatReturn: true,
	}
	task.GetEntryName = func() string { return name }
	return task
}
func (self *Task) SetHandleMessage(handle *HanldeMessageMap) {
	self.handleMessage = handle
}
func (self *Task) ParseMessage(data []byte) {
	if self.handleMessage == nil {
		self.Error("no handleMessage:%d,%d", data[0], data[1])
		return
	}
	if len(data) < 2 {
		self.Error("message can not less then 2 bytes")
		return
	}
	switch fun := self.handleMessage[data[0]][data[1]]; fun {
	case nil:
		self.Error("no parse func for message2:%d,%d", data[0], data[1])
		return
	default:
		fun(self, data)
	}
}
func (self *Task) SetHandleReadFun(fun func(task *Task) ([]ReadData, error)) {
	self.handleReadFun = fun
}
func (self *Task) SetHandleWriteFun(fun func(task *Task, data []byte) error) {
	self.handleWriteFun = fun
}
func (self *Task) SetHandleParseFun(fun func(task *Task, data []byte) bool) {
	self.handleParse = fun
}
func (self *Task) SetHandleHeartBteaFun(fun func(task *Task), dur time.Duration) {
	self.handleHeartBeat = fun
	self.heartBeatTime = dur
}

/*
func (self *Task) SendCmd(ptr unsafe.Pointer, length int) {
	data := make([]byte, 0, 0)
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	header.Data = (*reflect.SliceHeader)(ptr).Data
	header.Len = length
	header.Cap = length
	self.handleWriteFun(self, data)

}
//*/
func (self *Task) SendCmd(v interface{}) {
	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.LittleEndian, v); err != nil {
		self.Error("SendCmd err:%s", err.Error())
		self.Conn.Close()
		return
	}
	self.handleWriteFun(self, buf.Bytes())

}

func (self *Task) GetCmd(data []byte, v interface{}) error {
	cmdsize := reflect.TypeOf(v).Size()
	if len(data) < cmdsize {
		self.Debug("err size:%s,%d,%d", reflect.TypeOf(v).Name(), len(data), cmdsize)
		return nil
	}
	buf := bytes.NewBuffer(data[:cmdsize])
	self.Debug("msg size:%s,%d,%d", reflect.TypeOf(v).Name(), len(data), cmdsize)
	if err := binary.Read(buf, binary.BigEndian, v); err != nil {
		self.Error("GetCmd err:%s", err.Error())
		return err
	}
	return nil
}
func (self *Task) Start() {
	go self.startRead()
	go self.startWrite()
}
func (self *Task) Stop() {
	self.Conn.Close()
	close(self.StopChan)
}
func (self *Task) startRead() {
	defer close(self.in)
	defer self.Debug("startRead close")
	for {
		data, err := self.handleReadFun(self)
		if err != nil {
			self.Error("read error:%s", err.Error())
			self.Stop()
			break
		}
		self.in <- data
	}
}
func (self *Task) startWrite() {
	heartBeat := time.Tick(self.heartBeatTime)
LOOP:
	for {
		select {
		case <-self.StopChan:
			self.Debug("startWrite close")
			return
		case <-heartBeat:
			self.handleHeartBeat(self)
			if !self.HeartBeatReturn {
				self.Debug("timeout error")
				break LOOP
			}
			self.HeartBeatReturn = false
		case datas, ok := <-self.in:
			if !ok {
				self.Debug("self.in chan err")
				break LOOP
			}
			for _, data := range datas {
				if ok := self.handleParse(self, data.Data); !ok {
					self.ParseMessage(data.Data)
				}
			}
		}
	}
}
<commit_msg>test code<commit_after>package gotcp

import (
	"bytes"
	"encoding/binary"
	"net"
	"reflect"
	"time"
)

type HandleMessageFunc func(task *Task, data []byte)

type HanldeMessageMap [255][256]HandleMessageFunc

type ReadData struct {
	Length int
	Data   []byte
}
type Task struct {
	Entry
	Conn            *net.TCPConn
	in              chan []ReadData
	StopChan        chan bool
	handleReadFun   func(task *Task) ([]ReadData, error)
	handleWriteFun  func(task *Task, data []byte) error
	handleParse     func(task *Task, data []byte) bool
	handleHeartBeat func(task *Task)
	heartBeatTime   time.Duration
	HeartBeatReturn bool
	handleMessage   *HanldeMessageMap
}

func NewTask(c *net.TCPConn, name string) *Task {
	task := &Task{
		Conn:            c,
		in:              make(chan []ReadData),
		StopChan:        make(chan bool),
		handleReadFun:   handleReadFunDefault,
		handleWriteFun:  handleWriteFunDefault,
		handleParse:     handleParseDefault,
		handleHeartBeat: handleHeartBeatRequestDefault,
		heartBeatTime:   time.Second * 10,
		HeartBeatReturn: true,
	}
	task.GetEntryName = func() string { return name }
	return task
}
func (self *Task) SetHandleMessage(handle *HanldeMessageMap) {
	self.handleMessage = handle
}
func (self *Task) ParseMessage(data []byte) {
	if self.handleMessage == nil {
		self.Error("no handleMessage:%d,%d", data[0], data[1])
		return
	}
	if len(data) < 2 {
		self.Error("message can not less then 2 bytes")
		return
	}
	switch fun := self.handleMessage[data[0]][data[1]]; fun {
	case nil:
		self.Error("no parse func for message2:%d,%d", data[0], data[1])
		return
	default:
		fun(self, data)
	}
}
func (self *Task) SetHandleReadFun(fun func(task *Task) ([]ReadData, error)) {
	self.handleReadFun = fun
}
func (self *Task) SetHandleWriteFun(fun func(task *Task, data []byte) error) {
	self.handleWriteFun = fun
}
func (self *Task) SetHandleParseFun(fun func(task *Task, data []byte) bool) {
	self.handleParse = fun
}
func (self *Task) SetHandleHeartBteaFun(fun func(task *Task), dur time.Duration) {
	self.handleHeartBeat = fun
	self.heartBeatTime = dur
}

/*
func (self *Task) SendCmd(ptr unsafe.Pointer, length int) {
	data := make([]byte, 0, 0)
	header := (*reflect.SliceHeader)(unsafe.Pointer(&data))
	header.Data = (*reflect.SliceHeader)(ptr).Data
	header.Len = length
	header.Cap = length
	self.handleWriteFun(self, data)

}
//*/
func (self *Task) SendCmd(v interface{}) {
	buf := new(bytes.Buffer)
	if err := binary.Write(buf, binary.LittleEndian, v); err != nil {
		self.Error("SendCmd err:%s", err.Error())
		self.Conn.Close()
		return
	}
	self.handleWriteFun(self, buf.Bytes())

}

func (self *Task) GetCmd(data []byte, v interface{}) error {
	cmdsize := int(reflect.TypeOf(v).Size())
	if len(data) < cmdsize {
		self.Debug("err size:%s,%d,%d", reflect.TypeOf(v).Name(), len(data), cmdsize)
		return nil
	}
	buf := bytes.NewBuffer(data[:cmdsize])
	self.Debug("msg size:%s,%d,%d", reflect.TypeOf(v).Name(), len(data), cmdsize)
	if err := binary.Read(buf, binary.BigEndian, v); err != nil {
		self.Error("GetCmd err:%s", err.Error())
		return err
	}
	return nil
}
func (self *Task) Start() {
	go self.startRead()
	go self.startWrite()
}
func (self *Task) Stop() {
	self.Conn.Close()
	close(self.StopChan)
}
func (self *Task) startRead() {
	defer close(self.in)
	defer self.Debug("startRead close")
	for {
		data, err := self.handleReadFun(self)
		if err != nil {
			self.Error("read error:%s", err.Error())
			self.Stop()
			break
		}
		self.in <- data
	}
}
func (self *Task) startWrite() {
	heartBeat := time.Tick(self.heartBeatTime)
LOOP:
	for {
		select {
		case <-self.StopChan:
			self.Debug("startWrite close")
			return
		case <-heartBeat:
			self.handleHeartBeat(self)
			if !self.HeartBeatReturn {
				self.Debug("timeout error")
				break LOOP
			}
			self.HeartBeatReturn = false
		case datas, ok := <-self.in:
			if !ok {
				self.Debug("self.in chan err")
				break LOOP
			}
			for _, data := range datas {
				if ok := self.handleParse(self, data.Data); !ok {
					self.ParseMessage(data.Data)
				}
			}
		}
	}
}
<|endoftext|>