<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 objectpath_test

import (
	"bytes"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"testing"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/gcexportdata"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types/objectpath"
)

func TestPaths(t *testing.T) {
	pkgs := map[string]map[string]string{
		"b": {"b.go": `
package b

import "a"

const C = a.Int(0)

func F(a, b, c int, d a.T)

type T struct{ A int; b int; a.T }

func (T) M() *interface{ f() }

type U T

type A = struct{ x int }

var V []*a.T

type M map[struct{x int}]struct{y int}

func unexportedFunc()
type unexportedType struct{}
`},
		"a": {"a.go": `
package a

type Int int

type T struct{x, y int}

`},
	}
	conf := loader.Config{Build: buildutil.FakeContext(pkgs)}
	conf.Import("a")
	conf.Import("b")
	prog, err := conf.Load()
	if err != nil {
		t.Fatal(err)
	}
	a := prog.Imported["a"].Pkg
	b := prog.Imported["b"].Pkg

	// We test objectpath by enumerating a set of paths
	// and ensuring that Path(pkg, Object(pkg, path)) == path.
	//
	// It might seem more natural to invert the test:
	// identify a set of objects and for each one,
	// ensure that Object(pkg, Path(pkg, obj)) == obj.
	// However, for most interesting test cases there is no
	// easy way to identify the object short of applying
	// a series of destructuring operations to pkg---which
	// is essentially what objectpath.Object does.
	// (We do a little of that when testing bad paths, below.)
	//
	// The downside is that the test depends on the path encoding.
	// The upside is that the test exercises the encoding.

	// good paths
	for _, test := range []struct {
		pkg     *types.Package
		path    objectpath.Path
		wantobj string
	}{
		{b, "C", "const b.C a.Int"},
		{b, "F", "func b.F(a int, b int, c int, d a.T)"},
		{b, "F.PA0", "var a int"},
		{b, "F.PA1", "var b int"},
		{b, "F.PA2", "var c int"},
		{b, "F.PA3", "var d a.T"},
		{b, "T", "type b.T struct{A int; b int; a.T}"},
		{b, "T.O", "type b.T struct{A int; b int; a.T}"},
		{b, "T.UF0", "field A int"},
		{b, "T.UF1", "field b int"},
		{b, "T.UF2", "field T a.T"},
		{b, "U.UF2", "field T a.T"}, // U.U... are aliases for T.U...
		{b, "A", "type b.A = struct{x int}"},
		{b, "A.F0", "field x int"},
		{b, "V", "var b.V []*a.T"},
		{b, "M", "type b.M map[struct{x int}]struct{y int}"},
		{b, "M.UKF0", "field x int"},
		{b, "M.UEF0", "field y int"},
		{b, "T.M0", "func (b.T).M() *interface{f()}"}, // concrete method
		{b, "T.M0.RA0", "var  *interface{f()}"},       // parameter
		{b, "T.M0.RA0.EM0", "func (interface).f()"},   // interface method
		{b, "unexportedType", "type b.unexportedType struct{}"},
		{a, "T", "type a.T struct{x int; y int}"},
		{a, "T.UF0", "field x int"},
	} {
		// check path -> object
		obj, err := objectpath.Object(test.pkg, test.path)
		if err != nil {
			t.Errorf("Object(%s, %q) failed: %v",
				test.pkg.Path(), test.path, err)
			continue
		}
		if obj.String() != test.wantobj {
			t.Errorf("Object(%s, %q) = %v, want %s",
				test.pkg.Path(), test.path, obj, test.wantobj)
			continue
		}
		if obj.Pkg() != test.pkg {
			t.Errorf("Object(%s, %q) = %v, which belongs to package %s",
				test.pkg.Path(), test.path, obj, obj.Pkg().Path())
			continue
		}

		// check object -> path
		path2, err := objectpath.For(obj)
		if err != nil {
			t.Errorf("For(%v) failed: %v, want %q", obj, err, test.path)
			continue
		}
		// We do not require that test.path == path2. Aliases are legal.
		// But we do require that Object(path2) finds the same object.
		obj2, err := objectpath.Object(test.pkg, path2)
		if err != nil {
			t.Errorf("Object(%s, %q) failed: %v (roundtrip from %q)",
				test.pkg.Path(), path2, err, test.path)
			continue
		}
		if obj2 != obj {
			t.Errorf("Object(%s, For(obj)) != obj: got %s, obj is %s (path1=%q, path2=%q)",
				test.pkg.Path(), obj2, obj, test.path, path2)
			continue
		}
	}

	// bad paths (all relative to package b)
	for _, test := range []struct {
		pkg     *types.Package
		path    objectpath.Path
		wantErr string
	}{
		{b, "", "empty path"},
		{b, "missing", `package b does not contain "missing"`},
		{b, "F.U", "invalid path: ends with 'U', want [AFMO]"},
		{b, "F.PA3.O", "path denotes type a.T struct{x int; y int}, which belongs to a different package"},
		{b, "F.PA!", `invalid path: bad numeric operand "" for code 'A'`},
		{b, "F.PA3.UF0", "path denotes field x int, which belongs to a different package"},
		{b, "F.PA3.UF5", "field index 5 out of range [0-2)"},
		{b, "V.EE", "invalid path: ends with 'E', want [AFMO]"},
		{b, "F..O", "invalid path: unexpected '.' in type context"},
		{b, "T.OO", "invalid path: code 'O' in object context"},
		{b, "T.EO", "cannot apply 'E' to b.T (got *types.Named, want pointer, slice, array, chan or map)"},
		{b, "A.O", "cannot apply 'O' to struct{x int} (got struct{x int}, want named)"},
		{b, "A.UF0", "cannot apply 'U' to struct{x int} (got struct{x int}, want named)"},
		{b, "M.UPO", "cannot apply 'P' to map[struct{x int}]struct{y int} (got *types.Map, want signature)"},
		{b, "C.O", "path denotes type a.Int int, which belongs to a different package"},
	} {
		obj, err := objectpath.Object(test.pkg, test.path)
		if err == nil {
			t.Errorf("Object(%s, %q) = %s, want error",
				test.pkg.Path(), test.path, obj)
			continue
		}
		if err.Error() != test.wantErr {
			t.Errorf("Object(%s, %q) error was %q, want %q",
				test.pkg.Path(), test.path, err, test.wantErr)
			continue
		}
	}

	// bad objects
	bInfo := prog.Imported["b"]
	for _, test := range []struct {
		obj     types.Object
		wantErr string
	}{
		{types.Universe.Lookup("nil"), "predeclared nil has no path"},
		{types.Universe.Lookup("len"), "predeclared builtin len has no path"},
		{types.Universe.Lookup("int"), "predeclared type int has no path"},
		{bInfo.Info.Implicits[bInfo.Files[0].Imports[0]], "no path for package a"}, // import "a"
		{b.Scope().Lookup("unexportedFunc"), "no path for non-exported func b.unexportedFunc()"},
	} {
		path, err := objectpath.For(test.obj)
		if err == nil {
			t.Errorf("Object(%s) = %q, want error", test.obj, path)
			continue
		}
		if err.Error() != test.wantErr {
			t.Errorf("Object(%s) error was %q, want %q", test.obj, err, test.wantErr)
			continue
		}
	}
}

// TestSourceAndExportData uses objectpath to compute a correspondence
// of objects between two versions of the same package, one loaded from
// source, the other from export data.
func TestSourceAndExportData(t *testing.T) {
	const src = `
package p

type I int

func (I) F() *struct{ X, Y int } {
	return nil
}

type Foo interface {
	Method() (string, func(int) struct{ X int })
}

var X chan struct{ Z int }
var Z map[string]struct{ A int }
`

	// Parse source file and type-check it as a package, "src".
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "src.go", src, 0)
	if err != nil {
		t.Fatal(err)
	}
	conf := types.Config{Importer: importer.For("source", nil)}
	info := &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
	}
	srcpkg, err := conf.Check("src/p", fset, []*ast.File{f}, info)
	if err != nil {
		t.Fatal(err)
	}

	// Export binary export data then reload it as a new package, "bin".
	var buf bytes.Buffer
	if err := gcexportdata.Write(&buf, fset, srcpkg); err != nil {
		t.Fatal(err)
	}

	imports := make(map[string]*types.Package)
	binpkg, err := gcexportdata.Read(&buf, fset, imports, "bin/p")
	if err != nil {
		t.Fatal(err)
	}

	// Now find the correspondences between them.
	for _, srcobj := range info.Defs {
		if srcobj == nil {
			continue // e.g. package declaration
		}
		if _, ok := srcobj.(*types.PkgName); ok {
			continue // PkgName has no objectpath
		}

		path, err := objectpath.For(srcobj)
		if err != nil {
			t.Errorf("For(%v): %v", srcobj, err)
			continue
		}
		binobj, err := objectpath.Object(binpkg, path)
		if err != nil {
			t.Errorf("Object(%s, %q): %v", binpkg.Path(), path, err)
			continue
		}

		// Check the object strings match.
		// (We can't check that types are identical because the
		// objects belong to different type-checker realms.)
		srcstr := types.ObjectString(srcobj, (*types.Package).Name)
		binstr := types.ObjectString(binobj, (*types.Package).Name)
		if srcstr != binstr {
			t.Errorf("ObjectStrings do not match: Object(For(%q)) = %s, want %s",
				path, srcstr, binstr)
			continue
		}
	}
}
<commit_msg>go/types/objectpath: fix tests for pre-go1.11<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 objectpath_test

import (
	"bytes"
	"go/ast"
	"go/importer"
	"go/parser"
	"go/token"
	"go/types"
	"strings"
	"testing"

	"golang.org/x/tools/go/buildutil"
	"golang.org/x/tools/go/gcexportdata"
	"golang.org/x/tools/go/loader"
	"golang.org/x/tools/go/types/objectpath"
)

func TestPaths(t *testing.T) {
	pkgs := map[string]map[string]string{
		"b": {"b.go": `
package b

import "a"

const C = a.Int(0)

func F(a, b, c int, d a.T)

type T struct{ A int; b int; a.T }

func (T) M() *interface{ f() }

type U T

type A = struct{ x int }

var V []*a.T

type M map[struct{x int}]struct{y int}

func unexportedFunc()
type unexportedType struct{}
`},
		"a": {"a.go": `
package a

type Int int

type T struct{x, y int}

`},
	}
	conf := loader.Config{Build: buildutil.FakeContext(pkgs)}
	conf.Import("a")
	conf.Import("b")
	prog, err := conf.Load()
	if err != nil {
		t.Fatal(err)
	}
	a := prog.Imported["a"].Pkg
	b := prog.Imported["b"].Pkg

	// We test objectpath by enumerating a set of paths
	// and ensuring that Path(pkg, Object(pkg, path)) == path.
	//
	// It might seem more natural to invert the test:
	// identify a set of objects and for each one,
	// ensure that Object(pkg, Path(pkg, obj)) == obj.
	// However, for most interesting test cases there is no
	// easy way to identify the object short of applying
	// a series of destructuring operations to pkg---which
	// is essentially what objectpath.Object does.
	// (We do a little of that when testing bad paths, below.)
	//
	// The downside is that the test depends on the path encoding.
	// The upside is that the test exercises the encoding.

	// good paths
	for _, test := range []struct {
		pkg     *types.Package
		path    objectpath.Path
		wantobj string
	}{
		{b, "C", "const b.C a.Int"},
		{b, "F", "func b.F(a int, b int, c int, d a.T)"},
		{b, "F.PA0", "var a int"},
		{b, "F.PA1", "var b int"},
		{b, "F.PA2", "var c int"},
		{b, "F.PA3", "var d a.T"},
		{b, "T", "type b.T struct{A int; b int; a.T}"},
		{b, "T.O", "type b.T struct{A int; b int; a.T}"},
		{b, "T.UF0", "field A int"},
		{b, "T.UF1", "field b int"},
		{b, "T.UF2", "field T a.T"},
		{b, "U.UF2", "field T a.T"}, // U.U... are aliases for T.U...
		{b, "A", "type b.A = struct{x int}"},
		{b, "A.F0", "field x int"},
		{b, "V", "var b.V []*a.T"},
		{b, "M", "type b.M map[struct{x int}]struct{y int}"},
		{b, "M.UKF0", "field x int"},
		{b, "M.UEF0", "field y int"},
		{b, "T.M0", "func (b.T).M() *interface{f()}"}, // concrete method
		{b, "T.M0.RA0", "var  *interface{f()}"},       // parameter
		{b, "T.M0.RA0.EM0", "func (interface).f()"},   // interface method
		{b, "unexportedType", "type b.unexportedType struct{}"},
		{a, "T", "type a.T struct{x int; y int}"},
		{a, "T.UF0", "field x int"},
	} {
		// check path -> object
		obj, err := objectpath.Object(test.pkg, test.path)
		if err != nil {
			t.Errorf("Object(%s, %q) failed: %v",
				test.pkg.Path(), test.path, err)
			continue
		}
		if obj.String() != test.wantobj {
			t.Errorf("Object(%s, %q) = %v, want %s",
				test.pkg.Path(), test.path, obj, test.wantobj)
			continue
		}
		if obj.Pkg() != test.pkg {
			t.Errorf("Object(%s, %q) = %v, which belongs to package %s",
				test.pkg.Path(), test.path, obj, obj.Pkg().Path())
			continue
		}

		// check object -> path
		path2, err := objectpath.For(obj)
		if err != nil {
			t.Errorf("For(%v) failed: %v, want %q", obj, err, test.path)
			continue
		}
		// We do not require that test.path == path2. Aliases are legal.
		// But we do require that Object(path2) finds the same object.
		obj2, err := objectpath.Object(test.pkg, path2)
		if err != nil {
			t.Errorf("Object(%s, %q) failed: %v (roundtrip from %q)",
				test.pkg.Path(), path2, err, test.path)
			continue
		}
		if obj2 != obj {
			t.Errorf("Object(%s, For(obj)) != obj: got %s, obj is %s (path1=%q, path2=%q)",
				test.pkg.Path(), obj2, obj, test.path, path2)
			continue
		}
	}

	// bad paths (all relative to package b)
	for _, test := range []struct {
		pkg     *types.Package
		path    objectpath.Path
		wantErr string
	}{
		{b, "", "empty path"},
		{b, "missing", `package b does not contain "missing"`},
		{b, "F.U", "invalid path: ends with 'U', want [AFMO]"},
		{b, "F.PA3.O", "path denotes type a.T struct{x int; y int}, which belongs to a different package"},
		{b, "F.PA!", `invalid path: bad numeric operand "" for code 'A'`},
		{b, "F.PA3.UF0", "path denotes field x int, which belongs to a different package"},
		{b, "F.PA3.UF5", "field index 5 out of range [0-2)"},
		{b, "V.EE", "invalid path: ends with 'E', want [AFMO]"},
		{b, "F..O", "invalid path: unexpected '.' in type context"},
		{b, "T.OO", "invalid path: code 'O' in object context"},
		{b, "T.EO", "cannot apply 'E' to b.T (got *types.Named, want pointer, slice, array, chan or map)"},
		{b, "A.O", "cannot apply 'O' to struct{x int} (got struct{x int}, want named)"},
		{b, "A.UF0", "cannot apply 'U' to struct{x int} (got struct{x int}, want named)"},
		{b, "M.UPO", "cannot apply 'P' to map[struct{x int}]struct{y int} (got *types.Map, want signature)"},
		{b, "C.O", "path denotes type a.Int int, which belongs to a different package"},
	} {
		obj, err := objectpath.Object(test.pkg, test.path)
		if err == nil {
			t.Errorf("Object(%s, %q) = %s, want error",
				test.pkg.Path(), test.path, obj)
			continue
		}
		if err.Error() != test.wantErr {
			t.Errorf("Object(%s, %q) error was %q, want %q",
				test.pkg.Path(), test.path, err, test.wantErr)
			continue
		}
	}

	// bad objects
	bInfo := prog.Imported["b"]
	for _, test := range []struct {
		obj     types.Object
		wantErr string
	}{
		{types.Universe.Lookup("nil"), "predeclared nil has no path"},
		{types.Universe.Lookup("len"), "predeclared builtin len has no path"},
		{types.Universe.Lookup("int"), "predeclared type int has no path"},
		{bInfo.Info.Implicits[bInfo.Files[0].Imports[0]], "no path for package a"}, // import "a"
		{b.Scope().Lookup("unexportedFunc"), "no path for non-exported func b.unexportedFunc()"},
	} {
		path, err := objectpath.For(test.obj)
		if err == nil {
			t.Errorf("Object(%s) = %q, want error", test.obj, path)
			continue
		}
		if err.Error() != test.wantErr {
			t.Errorf("Object(%s) error was %q, want %q", test.obj, err, test.wantErr)
			continue
		}
	}
}

// TestSourceAndExportData uses objectpath to compute a correspondence
// of objects between two versions of the same package, one loaded from
// source, the other from export data.
func TestSourceAndExportData(t *testing.T) {
	const src = `
package p

type I int

func (I) F() *struct{ X, Y int } {
	return nil
}

type Foo interface {
	Method() (string, func(int) struct{ X int })
}

var X chan struct{ Z int }
var Z map[string]struct{ A int }
`

	// Parse source file and type-check it as a package, "src".
	fset := token.NewFileSet()
	f, err := parser.ParseFile(fset, "src.go", src, 0)
	if err != nil {
		t.Fatal(err)
	}
	conf := types.Config{Importer: importer.For("source", nil)}
	info := &types.Info{
		Defs: make(map[*ast.Ident]types.Object),
	}
	srcpkg, err := conf.Check("src/p", fset, []*ast.File{f}, info)
	if err != nil {
		t.Fatal(err)
	}

	// Export binary export data then reload it as a new package, "bin".
	var buf bytes.Buffer
	if err := gcexportdata.Write(&buf, fset, srcpkg); err != nil {
		t.Fatal(err)
	}

	imports := make(map[string]*types.Package)
	binpkg, err := gcexportdata.Read(&buf, fset, imports, "bin/p")
	if err != nil {
		t.Fatal(err)
	}

	// Now find the correspondences between them.
	for _, srcobj := range info.Defs {
		if srcobj == nil {
			continue // e.g. package declaration
		}
		if _, ok := srcobj.(*types.PkgName); ok {
			continue // PkgName has no objectpath
		}

		path, err := objectpath.For(srcobj)
		if err != nil {
			t.Errorf("For(%v): %v", srcobj, err)
			continue
		}
		binobj, err := objectpath.Object(binpkg, path)
		if err != nil {
			t.Errorf("Object(%s, %q): %v", binpkg.Path(), path, err)
			continue
		}

		// Check the object strings match.
		// (We can't check that types are identical because the
		// objects belong to different type-checker realms.)
		srcstr := objectString(srcobj)
		binstr := objectString(binobj)
		if srcstr != binstr {
			t.Errorf("ObjectStrings do not match: Object(For(%q)) = %s, want %s",
				path, srcstr, binstr)
			continue
		}
	}
}

func objectString(obj types.Object) string {
	s := types.ObjectString(obj, (*types.Package).Name)

	// The printing of interface methods changed in go1.11.
	// This work-around makes the specific test pass with earlier versions.
	s = strings.Replace(s, "func (interface).Method", "func (p.Foo).Method", -1)

	return s
}
<|endoftext|>


<commit_before><commit_msg>gocolor -> gxlog<commit_after><|endoftext|>


<commit_before>package kontrolhelper

import (
	"encoding/json"
	"github.com/streadway/amqp"
	"io/ioutil"
	"koding/tools/config"
	"log"
	"os"
	"strings"
)

type Producer struct {
	Conn    *amqp.Connection
	Channel *amqp.Channel
	Name    string
	Done    chan error
}

func NewProducer(name string) *Producer {
	return &Producer{
		Conn:    nil,
		Channel: nil,
		Name:    name,
		Done:    make(chan error),
	}
}

func CreateAmqpConnection() *amqp.Connection {
	user := config.Current.Kontrold.RabbitMq.Login
	password := config.Current.Kontrold.RabbitMq.Password
	host := config.Current.Kontrold.RabbitMq.Host
	port := config.Current.Kontrold.RabbitMq.Port

	if port == "" {
		port = "5672" // default RABBITMQ_NODE_PORT
	}

	url := "amqp://" + user + ":" + password + "@" + host + ":" + port
	conn, err := amqp.Dial(url)
	if err != nil {
		log.Fatalln("AMQP dial: ", err)
	}

	go func() {
		for err := range conn.NotifyClose(make(chan *amqp.Error)) {
			log.Fatalf("AMQP connection: %s" + err.Error())
		}
	}()

	return conn
}

func CreateChannel(conn *amqp.Connection) *amqp.Channel {
	channel, err := conn.Channel()
	if err != nil {
		panic(err)
	}
	go func() {
		for err := range channel.NotifyClose(make(chan *amqp.Error)) {
			log.Fatalf("AMQP channel: %s" + err.Error())
		}
	}()
	return channel
}

func CreateStream(channel *amqp.Channel, kind, exchange, queue, key string, durable, autoDelete bool) <-chan amqp.Delivery {
	if err := channel.ExchangeDeclare(exchange, kind, durable, autoDelete, false, false, nil); err != nil {
		panic(err)
	}

	if _, err := channel.QueueDeclare(queue, true, false, false, false, nil); err != nil {
		panic(err)
	}

	if err := channel.QueueBind(queue, key, exchange, false, nil); err != nil {
		panic(err)
	}

	stream, err := channel.Consume(queue, "", true, true, false, false, nil)
	if err != nil {
		panic(err)
	}

	return stream
}

func CustomHostname() string {
	hostname, err := os.Hostname()
	if err != nil {
		log.Println(err)
	}

	return hostname
}

func ReadVersion() string {
	file, err := ioutil.ReadFile("VERSION")
	if err != nil {
		log.Println(err)
	}

	return strings.TrimSpace(string(file))
}

func ReadFile(config string) string {
	file, err := ioutil.ReadFile(config)
	if err != nil {
		log.Println(err)
	}

	return strings.TrimSpace(string(file))
}

func CreateProducer(name string) (*Producer, error) {
	p := NewProducer(name)
	log.Printf("creating connection for sending %s messages", p.Name)
	p.Conn = CreateAmqpConnection()
	p.Channel = CreateChannel(p.Conn)

	return p, nil
}

func RegisterToKontrol(name, uuid string, port int) error {
	connection := CreateAmqpConnection()
	channel := CreateChannel(connection)

	type workerMessage struct {
		Command string
		Option  string
		Result  string
	}

	type workerMain struct {
		Name     string
		Uuid     string
		Hostname string
		Message  workerMessage
		Port     int
	}

	cmd := workerMain{
		Name:     name,
		Uuid:     uuid,
		Hostname: CustomHostname(),
		Message: workerMessage{
			Command: "addWithProxy",
			Option:  "many",
		},
		Port: port,
	}

	type Wrap struct{ Worker workerMain }

	data, err := json.Marshal(&Wrap{cmd})
	if err != nil {
		return err
	}

	msg := amqp.Publishing{
		Headers:         amqp.Table{},
		ContentType:     "text/plain",
		ContentEncoding: "",
		Body:            data,
		DeliveryMode:    1, // 1=non-persistent, 2=persistent
		AppId:           uuid,
	}

	err = channel.Publish("workerExchange", "input.worker", false, false, msg)
	if err != nil {
		return err
	}

	err = channel.Close()
	if err != nil {
		log.Println("could not close kontrold publisher amqp channel", err)
	}

	return nil
}
<commit_msg>kontrolhelper: create exchange if it doesn't exist<commit_after>package kontrolhelper

import (
	"encoding/json"
	"github.com/streadway/amqp"
	"io/ioutil"
	"koding/tools/config"
	"log"
	"os"
	"strings"
)

type Producer struct {
	Conn    *amqp.Connection
	Channel *amqp.Channel
	Name    string
	Done    chan error
}

func NewProducer(name string) *Producer {
	return &Producer{
		Conn:    nil,
		Channel: nil,
		Name:    name,
		Done:    make(chan error),
	}
}

func CreateAmqpConnection() *amqp.Connection {
	user := config.Current.Kontrold.RabbitMq.Login
	password := config.Current.Kontrold.RabbitMq.Password
	host := config.Current.Kontrold.RabbitMq.Host
	port := config.Current.Kontrold.RabbitMq.Port

	if port == "" {
		port = "5672" // default RABBITMQ_NODE_PORT
	}

	url := "amqp://" + user + ":" + password + "@" + host + ":" + port
	conn, err := amqp.Dial(url)
	if err != nil {
		log.Fatalln("AMQP dial: ", err)
	}

	go func() {
		for err := range conn.NotifyClose(make(chan *amqp.Error)) {
			log.Fatalf("AMQP connection: %s" + err.Error())
		}
	}()

	return conn
}

func CreateChannel(conn *amqp.Connection) *amqp.Channel {
	channel, err := conn.Channel()
	if err != nil {
		panic(err)
	}
	go func() {
		for err := range channel.NotifyClose(make(chan *amqp.Error)) {
			log.Fatalf("AMQP channel: %s" + err.Error())
		}
	}()
	return channel
}

func CreateStream(channel *amqp.Channel, kind, exchange, queue, key string, durable, autoDelete bool) <-chan amqp.Delivery {
	if err := channel.ExchangeDeclare(exchange, kind, durable, autoDelete, false, false, nil); err != nil {
		panic(err)
	}

	if _, err := channel.QueueDeclare(queue, true, false, false, false, nil); err != nil {
		panic(err)
	}

	if err := channel.QueueBind(queue, key, exchange, false, nil); err != nil {
		panic(err)
	}

	stream, err := channel.Consume(queue, "", true, true, false, false, nil)
	if err != nil {
		panic(err)
	}

	return stream
}

func CustomHostname() string {
	hostname, err := os.Hostname()
	if err != nil {
		log.Println(err)
	}

	return hostname
}

func ReadVersion() string {
	file, err := ioutil.ReadFile("VERSION")
	if err != nil {
		log.Println(err)
	}

	return strings.TrimSpace(string(file))
}

func ReadFile(config string) string {
	file, err := ioutil.ReadFile(config)
	if err != nil {
		log.Println(err)
	}

	return strings.TrimSpace(string(file))
}

func CreateProducer(name string) (*Producer, error) {
	p := NewProducer(name)
	log.Printf("creating connection for sending %s messages", p.Name)
	p.Conn = CreateAmqpConnection()
	p.Channel = CreateChannel(p.Conn)

	return p, nil
}

func RegisterToKontrol(name, uuid string, port int) error {
	connection := CreateAmqpConnection()
	channel := CreateChannel(connection)

	type workerMessage struct {
		Command string
		Option  string
		Result  string
	}

	type workerMain struct {
		Name     string
		Uuid     string
		Hostname string
		Message  workerMessage
		Port     int
	}

	cmd := workerMain{
		Name:     name,
		Uuid:     uuid,
		Hostname: CustomHostname(),
		Message: workerMessage{
			Command: "addWithProxy",
			Option:  "many",
		},
		Port: port,
	}

	type Wrap struct{ Worker workerMain }

	data, err := json.Marshal(&Wrap{cmd})
	if err != nil {
		return err
	}

	msg := amqp.Publishing{
		Headers:         amqp.Table{},
		ContentType:     "text/plain",
		ContentEncoding: "",
		Body:            data,
		DeliveryMode:    1, // 1=non-persistent, 2=persistent
		AppId:           uuid,
	}

	if err := channel.ExchangeDeclare("workerExchange", "topic", true, false, false, false, nil); err != nil {
		return err
	}

	err = channel.Publish("workerExchange", "input.worker", false, false, msg)
	if err != nil {
		return err
	}

	err = channel.Close()
	if err != nil {
		log.Println("could not close kontrold publisher amqp channel", err)
	}

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/maxmind/geoipupdate/v4/pkg/geoipupdate"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestMultipleDatabaseDownload(t *testing.T) {
	databaseContent := "database content goes here"

	server := httptest.NewServer(
		http.HandlerFunc(
			func(rw http.ResponseWriter, r *http.Request) {
				err := r.ParseForm()
				require.NoError(t, err, "parse form")

				if strings.HasPrefix(r.URL.Path, "/geoip/databases") {
					buf := &bytes.Buffer{}
					gzWriter := gzip.NewWriter(buf)
					md5Writer := md5.New()
					multiWriter := io.MultiWriter(gzWriter, md5Writer)
					_, err := multiWriter.Write([]byte(
						databaseContent + " " + r.URL.Path,
					))
					require.NoError(t, err)
					err = gzWriter.Close()
					require.NoError(t, err)

					rw.Header().Set(
						"X-Database-MD5",
						fmt.Sprintf("%x", md5Writer.Sum(nil)),
					)
					rw.Header().Set("Last-Modified", time.Now().Format(time.RFC1123))

					_, err = rw.Write(buf.Bytes())
					require.NoError(t, err)

					return
				}

				if r.URL.Path == "/app/update_getfilename" {
					_, err := rw.Write([]byte(r.Form.Get("product_id") + ".mmdb"))
					require.NoError(t, err)
					return
				}

				rw.WriteHeader(http.StatusBadRequest)
			},
		),
	)
	defer server.Close()

	client := server.Client()

	tempDir, err := ioutil.TempDir("", "gutest-")
	require.NoError(t, err)
	defer func() {
		err := os.RemoveAll(tempDir)
		require.NoError(t, err)
	}()

	config := &geoipupdate.Config{
		AccountID:         123,
		DatabaseDirectory: tempDir,
		EditionIDs:        []string{"GeoLite2-City", "GeoLite2-Country"},
		LicenseKey:        "testing",
		LockFile:          filepath.Join(tempDir, ".geoipupdate.lock"),
		URL:               server.URL,
	}

	logOutput := &bytes.Buffer{}
	log.SetOutput(logOutput)

	err = run(client, config)
	assert.NoError(t, err, "run successfully")

	assert.Equal(t, "", logOutput.String(), "no logged output")

	for _, editionID := range config.EditionIDs {
		path := filepath.Join(config.DatabaseDirectory, editionID+".mmdb")
		buf, err := ioutil.ReadFile(path) //nolint:gosec
		require.NoError(t, err, "read file")
		assert.Equal(
			t,
			databaseContent+" /geoip/databases/"+editionID+"/update",
			string(buf),
			"correct database",
		)
	}
}
<commit_msg>Use filepath.Clean instance of nolint<commit_after>package main

import (
	"bytes"
	"compress/gzip"
	"crypto/md5"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

	"github.com/maxmind/geoipupdate/v4/pkg/geoipupdate"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestMultipleDatabaseDownload(t *testing.T) {
	databaseContent := "database content goes here"

	server := httptest.NewServer(
		http.HandlerFunc(
			func(rw http.ResponseWriter, r *http.Request) {
				err := r.ParseForm()
				require.NoError(t, err, "parse form")

				if strings.HasPrefix(r.URL.Path, "/geoip/databases") {
					buf := &bytes.Buffer{}
					gzWriter := gzip.NewWriter(buf)
					md5Writer := md5.New()
					multiWriter := io.MultiWriter(gzWriter, md5Writer)
					_, err := multiWriter.Write([]byte(
						databaseContent + " " + r.URL.Path,
					))
					require.NoError(t, err)
					err = gzWriter.Close()
					require.NoError(t, err)

					rw.Header().Set(
						"X-Database-MD5",
						fmt.Sprintf("%x", md5Writer.Sum(nil)),
					)
					rw.Header().Set("Last-Modified", time.Now().Format(time.RFC1123))

					_, err = rw.Write(buf.Bytes())
					require.NoError(t, err)

					return
				}

				if r.URL.Path == "/app/update_getfilename" {
					_, err := rw.Write([]byte(r.Form.Get("product_id") + ".mmdb"))
					require.NoError(t, err)
					return
				}

				rw.WriteHeader(http.StatusBadRequest)
			},
		),
	)
	defer server.Close()

	client := server.Client()

	tempDir, err := ioutil.TempDir("", "gutest-")
	require.NoError(t, err)
	defer func() {
		err := os.RemoveAll(tempDir)
		require.NoError(t, err)
	}()

	config := &geoipupdate.Config{
		AccountID:         123,
		DatabaseDirectory: tempDir,
		EditionIDs:        []string{"GeoLite2-City", "GeoLite2-Country"},
		LicenseKey:        "testing",
		LockFile:          filepath.Join(tempDir, ".geoipupdate.lock"),
		URL:               server.URL,
	}

	logOutput := &bytes.Buffer{}
	log.SetOutput(logOutput)

	err = run(client, config)
	assert.NoError(t, err, "run successfully")

	assert.Equal(t, "", logOutput.String(), "no logged output")

	for _, editionID := range config.EditionIDs {
		path := filepath.Join(config.DatabaseDirectory, editionID+".mmdb")
		buf, err := ioutil.ReadFile(filepath.Clean(path))
		require.NoError(t, err, "read file")
		assert.Equal(
			t,
			databaseContent+" /geoip/databases/"+editionID+"/update",
			string(buf),
			"correct database",
		)
	}
}
<|endoftext|>


<commit_before>package oskite

import (
	"fmt"
	"koding/virt"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	redigo "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

var (
	oskites   = make(map[string]*OskiteInfo)
	oskitesMu sync.Mutex
)

func (o *Oskite) loadBalancer(correlationName, username, deadService string) string {
	blog := func(v interface{}) {
		log.Info("oskite loadbalancer for [correlationName: '%s' user: '%s' deadService: '%s'] results in --> %v.", correlationName, username, deadService, v)
	}

	resultOskite := o.ServiceUniquename
	lowestOskite := lowestOskiteLoad()
	if lowestOskite != "" {
		if deadService == lowestOskite {
			resultOskite = o.ServiceUniquename
		} else {
			resultOskite = lowestOskite
		}
	}

	var vm *virt.VM
	if bson.IsObjectIdHex(correlationName) {
		mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.FindId(bson.ObjectIdHex(correlationName)).One(&vm)
		})
	}

	if vm == nil {
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Find(bson.M{"hostnameAlias": correlationName}).One(&vm)
		}); err != nil {
			blog(fmt.Sprintf("no hostnameAlias found, returning %s", resultOskite))
			return resultOskite // no vm was found, return this oskite
		}
	}

	if vm.PinnedToHost != "" {
		blog(fmt.Sprintf("returning pinnedHost '%s'", vm.PinnedToHost))
		return vm.PinnedToHost
	}

	if vm.HostKite == "" {
		// also set hoskite to prevent race condition between terminal and
		// oskite. Because if we set it now, the "kite.who" method of terminal
		// will not reply with an empty response.
		err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id, "hostKite": nil}, bson.M{"$set": bson.M{"hostKite": resultOskite}})
		})

		blog(fmt.Sprintf("hostkite is empty returning '%s. (update err: %s)", resultOskite, err))
		return resultOskite
	}

	// maintenance and banned will be handled again in valideVM() function,
	// which will return a permission error.
	if vm.HostKite == "(maintenance)" || vm.HostKite == "(banned)" {
		blog(fmt.Sprintf("hostkite is %s returning '%s'", vm.HostKite, resultOskite))
		return resultOskite
	}

	// Set hostkite to nil if we detect a dead service. On the next call,
	// Oskite will point to an health service in validateVM function()
	// because it will detect that the hostkite is nil and change it to the
	// healthy service given by the client, which is the returned
	// k.ServiceUniqueName.
	if vm.HostKite == deadService {
		blog(fmt.Sprintf("dead service detected %s returning '%s'", vm.HostKite, o.ServiceUniquename))
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id}, bson.M{"$set": bson.M{"hostKite": nil}})
		}); err != nil {
			log.LogError(err, 0, vm.Id.Hex())
		}

		return resultOskite
	}

	blog(fmt.Sprintf("returning existing hostkite '%s'", vm.HostKite))
	return vm.HostKite
}

func (o *Oskite) setupRedis() {
	if o.RedisSession != nil {
		return
	}

	session, err := redis.NewRedisSession(conf.Redis)
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	o.RedisSession = session

	// oskite:production:sj:
	o.RedisPrefix = "oskite:" + conf.Environment + ":" + o.Region + ":"

	// kontainers-production-sj
	o.RedisKontainerSet = "kontainers-" + conf.Environment + "-" + o.Region

	// oskite:production:sj:kite-os-sj|kontainer3_sj_koding_com
	o.RedisKontainerKey = o.RedisPrefix + o.ServiceUniquename

	o.RedisSession.SetPrefix("oskite")
}

func (o *Oskite) redisBalancer() {
	o.setupRedis()

	log.Info("Connected to Redis with %s", o.RedisKontainerKey)
	_, err := redigo.Int(o.RedisSession.Do("SADD", o.RedisKontainerSet, o.RedisKontainerKey))
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	// update regularly our VMS info
	go func() {
		expireDuration := time.Second * 5
		for _ = range time.Tick(2 * time.Second) {
			oskiteInfo := o.GetOskiteInfo()

			if _, err := o.RedisSession.Do("HMSET", redigo.Args{o.RedisKontainerKey}.AddFlat(oskiteInfo)...); err != nil {
				log.Error("redis HMSET err: %v", err.Error())
			}

			reply, err := redigo.Int(o.RedisSession.Do("EXPIRE", o.RedisKontainerKey, expireDuration.Seconds()))
			if err != nil {
				log.Error("redis SET Expire %v. reply: %v err: %v", o.RedisKontainerKey, reply, err.Error())
			}
		}
	}()

	// get oskite statuses from others every 2 seconds
	for _ = range time.Tick(2 * time.Second) {
		kontainers, err := redigo.Strings(o.RedisSession.Do("SMEMBERS", o.RedisKontainerSet))
		if err != nil {
			log.Error("redis SMEMBER kontainers. err: %v", err.Error())
		}

		for _, kontainerHostname := range kontainers {
			// convert to o.ServiceUniquename format
			remoteOskite := strings.TrimPrefix(kontainerHostname, o.RedisPrefix)

			values, err := redigo.Values(o.RedisSession.Do("HGETALL", kontainerHostname))
			if err != nil {
				log.Error("redis HTGETALL %s. err: %v", kontainerHostname, err.Error())

				// kontainer might be dead, key gets than expired, continue with the next one

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			// cleanup members from the set if the key expires. Usually that
			// might due a host who added itself to the kontainer set but then
			// just died or added a wrong set name.
			if len(values) == 0 && kontainerHostname != o.RedisKontainerKey {
				_, err := redigo.Int(o.RedisSession.Do("SREM", o.RedisKontainerSet, kontainerHostname))
				if err != nil {
					log.Error("redis SREM kontainers. err: %v", err.Error())
				}

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			oskiteInfo := new(OskiteInfo)
			if err := redigo.ScanStruct(values, oskiteInfo); err != nil {
				log.Error("redis ScanStruct err: %v", err.Error())
			}

			oskitesMu.Lock()
			oskites[remoteOskite] = oskiteInfo
			oskitesMu.Unlock()
		}
	}
}

func lowestOskiteLoad() (serviceUniquename string) {
	oskitesMu.Lock()
	defer oskitesMu.Unlock()

	oskitesSlice := make([]*OskiteInfo, 0, len(oskites))

	for s, v := range oskites {
		v.ServiceUniquename = s
		oskitesSlice = append(oskitesSlice, v)
	}

	sort.Sort(ByVM(oskitesSlice))

	middle := len(oskitesSlice) / 2
	if middle == 0 {
		return ""
	}

	// return randomly one of the lowest
	l := oskitesSlice[rand.Intn(middle)]

	// also pick up the highest to log information
	h := oskitesSlice[len(oskitesSlice)-1]

	log.Info("oskite picked up as lowest load %s with %d VMs (highest was: %d / %s)",
		l.ServiceUniquename, l.ActiveVMs, h.ActiveVMs, h.ServiceUniquename)

	return l.ServiceUniquename
}
<commit_msg>oskite: only remove from internal map, it get's removed when oskite shuts down<commit_after>package oskite

import (
	"fmt"
	"koding/virt"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"

	redigo "github.com/garyburd/redigo/redis"
	"github.com/koding/redis"
)

var (
	oskites   = make(map[string]*OskiteInfo)
	oskitesMu sync.Mutex
)

func (o *Oskite) loadBalancer(correlationName, username, deadService string) string {
	blog := func(v interface{}) {
		log.Info("oskite loadbalancer for [correlationName: '%s' user: '%s' deadService: '%s'] results in --> %v.", correlationName, username, deadService, v)
	}

	resultOskite := o.ServiceUniquename
	lowestOskite := lowestOskiteLoad()
	if lowestOskite != "" {
		if deadService == lowestOskite {
			resultOskite = o.ServiceUniquename
		} else {
			resultOskite = lowestOskite
		}
	}

	var vm *virt.VM
	if bson.IsObjectIdHex(correlationName) {
		mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.FindId(bson.ObjectIdHex(correlationName)).One(&vm)
		})
	}

	if vm == nil {
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Find(bson.M{"hostnameAlias": correlationName}).One(&vm)
		}); err != nil {
			blog(fmt.Sprintf("no hostnameAlias found, returning %s", resultOskite))
			return resultOskite // no vm was found, return this oskite
		}
	}

	if vm.PinnedToHost != "" {
		blog(fmt.Sprintf("returning pinnedHost '%s'", vm.PinnedToHost))
		return vm.PinnedToHost
	}

	if vm.HostKite == "" {
		// also set hoskite to prevent race condition between terminal and
		// oskite. Because if we set it now, the "kite.who" method of terminal
		// will not reply with an empty response.
		err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id, "hostKite": nil}, bson.M{"$set": bson.M{"hostKite": resultOskite}})
		})

		blog(fmt.Sprintf("hostkite is empty returning '%s. (update err: %s)", resultOskite, err))
		return resultOskite
	}

	// maintenance and banned will be handled again in valideVM() function,
	// which will return a permission error.
	if vm.HostKite == "(maintenance)" || vm.HostKite == "(banned)" {
		blog(fmt.Sprintf("hostkite is %s returning '%s'", vm.HostKite, resultOskite))
		return resultOskite
	}

	// Set hostkite to nil if we detect a dead service. On the next call,
	// Oskite will point to an health service in validateVM function()
	// because it will detect that the hostkite is nil and change it to the
	// healthy service given by the client, which is the returned
	// k.ServiceUniqueName.
	if vm.HostKite == deadService {
		blog(fmt.Sprintf("dead service detected %s returning '%s'", vm.HostKite, o.ServiceUniquename))
		if err := mongodbConn.Run("jVMs", func(c *mgo.Collection) error {
			return c.Update(bson.M{"_id": vm.Id}, bson.M{"$set": bson.M{"hostKite": nil}})
		}); err != nil {
			log.LogError(err, 0, vm.Id.Hex())
		}

		return resultOskite
	}

	blog(fmt.Sprintf("returning existing hostkite '%s'", vm.HostKite))
	return vm.HostKite
}

func (o *Oskite) setupRedis() {
	if o.RedisSession != nil {
		return
	}

	session, err := redis.NewRedisSession(conf.Redis)
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	o.RedisSession = session

	// oskite:production:sj:
	o.RedisPrefix = "oskite:" + conf.Environment + ":" + o.Region + ":"

	// kontainers-production-sj
	o.RedisKontainerSet = "kontainers-" + conf.Environment + "-" + o.Region

	// oskite:production:sj:kite-os-sj|kontainer3_sj_koding_com
	o.RedisKontainerKey = o.RedisPrefix + o.ServiceUniquename

	o.RedisSession.SetPrefix("oskite")
}

func (o *Oskite) redisBalancer() {
	o.setupRedis()

	log.Info("Connected to Redis with %s", o.RedisKontainerKey)
	_, err := redigo.Int(o.RedisSession.Do("SADD", o.RedisKontainerSet, o.RedisKontainerKey))
	if err != nil {
		log.Error("redis SADD kontainers. err: %v", err.Error())
	}

	// update regularly our VMS info
	go func() {
		expireDuration := time.Second * 5
		for _ = range time.Tick(2 * time.Second) {
			oskiteInfo := o.GetOskiteInfo()

			if _, err := o.RedisSession.Do("HMSET", redigo.Args{o.RedisKontainerKey}.AddFlat(oskiteInfo)...); err != nil {
				log.Error("redis HMSET err: %v", err.Error())
			}

			reply, err := redigo.Int(o.RedisSession.Do("EXPIRE", o.RedisKontainerKey, expireDuration.Seconds()))
			if err != nil {
				log.Error("redis SET Expire %v. reply: %v err: %v", o.RedisKontainerKey, reply, err.Error())
			}
		}
	}()

	// get oskite statuses from others every 2 seconds
	for _ = range time.Tick(2 * time.Second) {
		kontainers, err := redigo.Strings(o.RedisSession.Do("SMEMBERS", o.RedisKontainerSet))
		if err != nil {
			log.Error("redis SMEMBER kontainers. err: %v", err.Error())
		}

		for _, kontainerHostname := range kontainers {
			// convert to o.ServiceUniquename format
			remoteOskite := strings.TrimPrefix(kontainerHostname, o.RedisPrefix)

			values, err := redigo.Values(o.RedisSession.Do("HGETALL", kontainerHostname))
			if err != nil {
				log.Error("redis HTGETALL %s. err: %v", kontainerHostname, err.Error())

				// kontainer might be dead, key gets than expired, continue with the next one

				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			// cleanup members from the map if the key expires. Usually that
			// might due a host who added itself to the kontainer set but then
			// just died or added a wrong set name. If the kontainer comes up
			// again we do add it below again.
			if len(values) == 0 {
				oskitesMu.Lock()
				delete(oskites, remoteOskite)
				oskitesMu.Unlock()
				continue
			}

			oskiteInfo := new(OskiteInfo)
			if err := redigo.ScanStruct(values, oskiteInfo); err != nil {
				log.Error("redis ScanStruct err: %v", err.Error())
			}

			oskitesMu.Lock()
			oskites[remoteOskite] = oskiteInfo
			oskitesMu.Unlock()
		}
	}
}

func lowestOskiteLoad() (serviceUniquename string) {
	oskitesMu.Lock()
	defer oskitesMu.Unlock()

	oskitesSlice := make([]*OskiteInfo, 0, len(oskites))

	for s, v := range oskites {
		v.ServiceUniquename = s
		oskitesSlice = append(oskitesSlice, v)
	}

	sort.Sort(ByVM(oskitesSlice))

	middle := len(oskitesSlice) / 2
	if middle == 0 {
		return ""
	}

	// return randomly one of the lowest
	l := oskitesSlice[rand.Intn(middle)]

	// also pick up the highest to log information
	h := oskitesSlice[len(oskitesSlice)-1]

	log.Info("oskite picked up as lowest load %s with %d VMs (highest was: %d / %s)",
		l.ServiceUniquename, l.ActiveVMs, h.ActiveVMs, h.ServiceUniquename)

	return l.ServiceUniquename
}
<|endoftext|>


<commit_before>package main

import (
	"errors"
	"image"
	"image/color"
	"image/draw"
	"os"
	"path"

	"github.com/MJKWoolnough/byteio"
	"github.com/MJKWoolnough/minecraft"
	"github.com/MJKWoolnough/minecraft/nbt"
	"github.com/MJKWoolnough/ora"
)

func (t Transfer) generate(name string, _ *byteio.StickyReader, w *byteio.StickyWriter, f *os.File, size int64) error {
	o, err := ora.Open(f, size)
	if err != nil {
		return err
	}
	terrain := o.Layer("terrain")
	if terrain == nil {
		return layerError{"terrain"}
	}
	height := o.Layer("height")
	if height == nil {
		return layerError{"height"}
	}
	mp := t.c.NewMap()
	if mp == nil {
		return errors.New("failed to create map")
	}

	done := false
	defer func() {
		if !done {
			t.c.RemoveMap(mp.ID)
		}
		go t.c.Save()
	}()

	mp.Lock()
	mp.Name = name
	mapPath := mp.Path
	mp.Server = -2
	mp.Unlock()

	ms := DefaultMapSettings()
	ms["level-type"] = minecraft.FlatGenerator
	ms["generator-settings"] = "0"
	ms["motd"] = name

	pf, err := os.Create(path.Join(mapPath, "properties.map"))
	if err != nil {
		return err
	}

	if err = ms.WriteTo(pf); err != nil {
		return err
	}
	pf.Close()

	b := o.Bounds()
	w.WriteUint8(2)
	w.WriteInt32(int32(b.Max.X) >> 4)
	w.WriteInt32(int32(b.Max.Y) >> 4)
	c := make(chan paint, 1024)
	m := make(chan string, 4)
	e := make(chan struct{}, 0)
	defer close(e)
	go func() {
		defer close(c)
		defer close(m)
		for {
			select {
			case message := <-m:
				w.WriteUint8(3)
				writeString(w, message)
			case p := <-c:
				w.WriteUint8(4)
				w.WriteInt32(p.X)
				w.WriteInt32(p.Y)
				r, g, b, a := p.RGBA()
				w.WriteUint8(uint8(r >> 8))
				w.WriteUint8(uint8(g >> 8))
				w.WriteUint8(uint8(b >> 8))
				w.WriteUint8(uint8(a >> 8))
			case <-e:
				return
			}
		}
	}()

	sTerrain := image.NewPaletted(o.Bounds(), terrainColours)
	terrainI, err := terrain.Image()
	if err != nil {
		return err
	}
	draw.Draw(sTerrain, image.Rect(terrain.X, terrain.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), terrainI, image.Point{}, draw.Src)
	terrainI = nil
	sHeight := image.NewGray(o.Bounds())
	heightI, err := height.Image()
	if err != nil {
		return err
	}
	draw.Draw(sHeight, image.Rect(height.X, height.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), heightI, image.Point{}, draw.Src)
	heightI = nil

	p, err := minecraft.NewFilePath(mapPath)
	if err != nil {
		return err
	}

	level, err := minecraft.NewLevel(p)
	if err != nil {
		return err
	}

	level.LevelName(name)

	m <- "Building Terrain"
	if err := buildTerrain(p, level, sTerrain, sHeight, c); err != nil {
		return err
	}

	level.LevelName(name)
	level.MobSpawning(false)
	level.KeepInventory(true)
	level.FireTick(false)
	level.DayLightCycle(false)
	level.MobGriefing(false)
	level.Spawn(10, 250, 10)
	level.Generator(minecraft.FlatGenerator)
	level.GeneratorOptions("0")
	level.GameMode(minecraft.Creative)
	level.AllowCommands(true)

	m <- "Exporting"
	level.Save()
	level.Close()
	done = true
	mp.Lock()
	mp.Server = -1
	mp.Unlock()

	return nil
}

type paint struct {
	color.Color
	X, Y int32
}

type layerError struct {
	name string
}

func (l layerError) Error() string {
	return "missing layer: " + l.name
}

type terrain struct {
	Base, Top minecraft.Block
	TopLevel  uint8
}

var (
	terrainColours = color.Palette{
		color.RGBA{},
		color.RGBA{255, 255, 0, 255},   // Yellow - Sand
		color.RGBA{0, 255, 0, 255},     // Green - Grass
		color.RGBA{87, 59, 12, 255},    // Brown - Dirt
		color.RGBA{255, 128, 0, 255},   // Orange - Farm
		color.RGBA{128, 128, 128, 255}, // Grey - Stone
		color.RGBA{255, 255, 255, 255}, // White - Snow
	}
	terrainBlocks = []terrain{
		{},
		{minecraft.Block{ID: 24, Data: 2}, minecraft.Block{ID: 12}, 5}, // Sandstone - Sand
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 2}, 1},            // Dirt - Grass
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 3}, 0},            // Dirt - Dirt
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 60, Data: 7}, 1},  // Dirt - Farmland
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 1}, 0},            // Stone - Stone
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 80}, 3},           // Stone - Snow
	}
)

func modeTerrain(p *image.Paletted) uint8 {
	b := p.Bounds()
	modeMap := make([]uint16, len(terrainColours))
	most := uint16(0)
	mode := uint8(0)
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			pos := p.ColorIndexAt(i, j)
			modeMap[pos]++
			if m := modeMap[pos]; m > most {
				most = m
				mode = pos
			}
		}
	}
	return mode
}

func meanHeight(g *image.Gray) uint8 {
	b := g.Bounds()
	var total int64
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			total += int64(g.GrayAt(i, j).Y)
		}
	}
	return uint8(total / int64((b.Dx() * b.Dy())))
}

type chunkCache struct {
	mem   *minecraft.MemPath
	level *minecraft.Level
	clear nbt.Tag
	cache map[uint16]nbt.Tag
}

func newCache() *chunkCache {
	mem := minecraft.NewMemPath()
	l, _ := minecraft.NewLevel(mem)

	bedrock := minecraft.Block{ID: 7}

	l.SetBlock(0, 0, 0, minecraft.Block{})
	l.Save()
	l.Close()
	clearChunk, _ := mem.GetChunk(0, 0)

	for i := int32(-16); i < 16; i++ {
		for j := int32(0); j < 255; j++ {
			for k := int32(-16); j < 16; j++ {
				l.SetBlock(i, j, k, bedrock)
			}
		}
	}
	l.Save()
	l.Close()
	mem.SetChunk(clearChunk)
	return &chunkCache{
		mem,
		l,
		clearChunk,
		make(map[uint16]nbt.Tag),
	}
}

func (c *chunkCache) getFromCache(x, z int32, terrain uint8, height int32) nbt.Tag {
	cacheID := uint16(terrain)<<8 | uint16(height)
	chunk, ok := c.cache[cacheID]
	if !ok {
		b := terrainBlocks[terrain].Base
		for j := int32(0); j < height; j++ {
			for i := int32(0); i < 16; i++ {
				for k := int32(0); k < 16; k++ {
					c.level.SetBlock(i, j, k, b)
				}
			}
		}
		c.level.Save()
		c.level.Close()
		chunk, _ = c.mem.GetChunk(0, 0)
		c.mem.SetChunk(c.clear)
		c.cache[cacheID] = chunk
	}
	ld := chunk.Data().(nbt.Compound).Get("Level").Data().(nbt.Compound)
	ld.Set(nbt.NewTag("xPos", nbt.Int(x)))
	ld.Set(nbt.NewTag("zPos", nbt.Int(z)))
	return chunk
}

func buildTerrain(mpath minecraft.Path, level *minecraft.Level, terrain *image.Paletted, height *image.Gray, c chan paint) error {
	b := terrain.Bounds()
	proceed := make(chan struct{}, 10)
	errChan := make(chan error, 1)
	go func() {
		defer close(proceed)
		cc := newCache()
		for j := 0; j < b.Max.Y; j += 16 {
			chunkZ := int32(j >> 4)
			for i := 0; i < b.Max.X; i += 16 {
				chunkX := int32(i >> 4)
				p := terrain.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Paletted)
				g := height.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Gray)
				t := modeTerrain(p)
				h := int32(meanHeight(g))
				err := mpath.SetChunk(cc.getFromCache(chunkX, chunkZ, t, h))
				if err != nil {
					errChan <- err
					return
				}
				c <- paint{
					terrainColours[t],
					chunkX, chunkZ,
				}
				proceed <- struct{}{}
			}
		}
	}()
	for i := 0; i < (b.Max.X>>4)+2; i++ {
		<-proceed // get far enough ahead so all chunks are surrounded before shaping, to get correct lighting
	}
	select {
	case err := <-errChan:
		return err
	default:
	}
	for j := int32(0); j < int32(b.Max.Y); j += 16 {
		chunkZ := j >> 4
		for i := int32(0); i < int32(b.Max.X); i += 16 {
			chunkX := i >> 4
			var totalHeight int32
			for x := i; x < i+16; x++ {
				for z := j; z < j+16; z++ {
					h := int32(height.GrayAt(int(x), int(z)).Y)
					min := h
					for mz := int(z) - 1; mz <= int(z)+1; mz++ {
						if mz < 0 || mz >= b.Max.Y {
							continue
						}
						for mx := int(x) - 1; mx <= int(x)+1; mx++ {
							if mx < 0 || mx >= b.Max.X {
								continue
							}
							if mh := int32(height.GrayAt(mx, mz).Y); mh < min {
								min = mh
							}
						}
					}

					totalHeight += h
					y, err := level.GetHeight(x, z)
					if err != nil {
						return err
					}
					for ; y > h; y-- {
						level.SetBlock(x, y, z, minecraft.Block{})
					}
					t := terrainBlocks[terrain.ColorIndexAt(int(x), int(z))]
					for ; y > h-int32(t.TopLevel); y-- {
						level.SetBlock(x, y, z, t.Top)
					}
					for ; y >= min; y-- {
						level.SetBlock(x, y, z, t.Base)
					}
				}
			}
			c <- paint{
				color.Alpha{uint8(totalHeight >> 8)},
				chunkX, chunkZ,
			}
			select {
			case <-proceed:
			case err := <-errChan:
				return err
			}
		}
	}
	return nil
}
<commit_msg>Added initial code to support biomes in generation<commit_after>package main

import (
	"errors"
	"image"
	"image/color"
	"image/draw"
	"os"
	"path"

	"github.com/MJKWoolnough/byteio"
	"github.com/MJKWoolnough/minecraft"
	"github.com/MJKWoolnough/minecraft/nbt"
	"github.com/MJKWoolnough/ora"
)

func (t Transfer) generate(name string, _ *byteio.StickyReader, w *byteio.StickyWriter, f *os.File, size int64) error {
	o, err := ora.Open(f, size)
	if err != nil {
		return err
	}
	terrain := o.Layer("terrain")
	if terrain == nil {
		return layerError{"terrain"}
	}
	height := o.Layer("height")
	if height == nil {
		return layerError{"height"}
	}
	mp := t.c.NewMap()
	if mp == nil {
		return errors.New("failed to create map")
	}

	done := false
	defer func() {
		if !done {
			t.c.RemoveMap(mp.ID)
		}
		go t.c.Save()
	}()

	mp.Lock()
	mp.Name = name
	mapPath := mp.Path
	mp.Server = -2
	mp.Unlock()

	ms := DefaultMapSettings()
	ms["level-type"] = minecraft.FlatGenerator
	ms["generator-settings"] = "0"
	ms["motd"] = name

	pf, err := os.Create(path.Join(mapPath, "properties.map"))
	if err != nil {
		return err
	}

	if err = ms.WriteTo(pf); err != nil {
		return err
	}
	pf.Close()

	b := o.Bounds()
	w.WriteUint8(2)
	w.WriteInt32(int32(b.Max.X) >> 4)
	w.WriteInt32(int32(b.Max.Y) >> 4)
	c := make(chan paint, 1024)
	m := make(chan string, 4)
	e := make(chan struct{}, 0)
	defer close(e)
	go func() {
		defer close(c)
		defer close(m)
		for {
			select {
			case message := <-m:
				w.WriteUint8(3)
				writeString(w, message)
			case p := <-c:
				w.WriteUint8(4)
				w.WriteInt32(p.X)
				w.WriteInt32(p.Y)
				r, g, b, a := p.RGBA()
				w.WriteUint8(uint8(r >> 8))
				w.WriteUint8(uint8(g >> 8))
				w.WriteUint8(uint8(b >> 8))
				w.WriteUint8(uint8(a >> 8))
			case <-e:
				return
			}
		}
	}()

	sTerrain := image.NewPaletted(o.Bounds(), terrainColours)
	terrainI, err := terrain.Image()
	if err != nil {
		return err
	}
	draw.Draw(sTerrain, image.Rect(terrain.X, terrain.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), terrainI, image.Point{}, draw.Src)
	terrainI = nil
	sHeight := image.NewGray(o.Bounds())
	heightI, err := height.Image()
	if err != nil {
		return err
	}
	draw.Draw(sHeight, image.Rect(height.X, height.Y, sTerrain.Bounds().Max.X, sTerrain.Bounds().Max.Y), heightI, image.Point{}, draw.Src)
	heightI = nil

	p, err := minecraft.NewFilePath(mapPath)
	if err != nil {
		return err
	}

	level, err := minecraft.NewLevel(p)
	if err != nil {
		return err
	}

	level.LevelName(name)

	m <- "Building Terrain"
	if err := buildTerrain(p, level, sTerrain, nil, sHeight, c); err != nil {
		return err
	}

	level.LevelName(name)
	level.MobSpawning(false)
	level.KeepInventory(true)
	level.FireTick(false)
	level.DayLightCycle(false)
	level.MobGriefing(false)
	level.Spawn(10, 250, 10)
	level.Generator(minecraft.FlatGenerator)
	level.GeneratorOptions("0")
	level.GameMode(minecraft.Creative)
	level.AllowCommands(true)

	m <- "Exporting"
	level.Save()
	level.Close()
	done = true
	mp.Lock()
	mp.Server = -1
	mp.Unlock()

	return nil
}

type paint struct {
	color.Color
	X, Y int32
}

type layerError struct {
	name string
}

func (l layerError) Error() string {
	return "missing layer: " + l.name
}

type terrain struct {
	Base, Top minecraft.Block
	TopLevel  uint8
}

var (
	terrainColours = color.Palette{
		color.RGBA{},
		color.RGBA{255, 255, 0, 255},   // Yellow - Sand
		color.RGBA{0, 255, 0, 255},     // Green - Grass
		color.RGBA{87, 59, 12, 255},    // Brown - Dirt
		color.RGBA{255, 128, 0, 255},   // Orange - Farm
		color.RGBA{128, 128, 128, 255}, // Grey - Stone
		color.RGBA{255, 255, 255, 255}, // White - Snow
	}
	terrainBlocks = []terrain{
		{},
		{minecraft.Block{ID: 24, Data: 2}, minecraft.Block{ID: 12}, 5}, // Sandstone - Sand
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 2}, 1},            // Dirt - Grass
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 3}, 0},            // Dirt - Dirt
		{minecraft.Block{ID: 3}, minecraft.Block{ID: 60, Data: 7}, 1},  // Dirt - Farmland
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 1}, 0},            // Stone - Stone
		{minecraft.Block{ID: 1}, minecraft.Block{ID: 80}, 3},           // Stone - Snow
	}
	biomePalette = color.Palette{}
	biomeList    = []minecraft.Biome{}
)

func modeTerrain(p *image.Paletted) uint8 {
	b := p.Bounds()
	modeMap := make([]uint16, len(terrainColours))
	most := uint16(0)
	mode := uint8(0)
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			pos := p.ColorIndexAt(i, j)
			modeMap[pos]++
			if m := modeMap[pos]; m > most {
				most = m
				mode = pos
			}
		}
	}
	return mode
}

func meanHeight(g *image.Gray) uint8 {
	b := g.Bounds()
	var total int64
	for i := b.Min.X; i < b.Max.X; i++ {
		for j := b.Min.Y; j < b.Max.Y; j++ {
			total += int64(g.GrayAt(i, j).Y)
		}
	}
	return uint8(total / int64((b.Dx() * b.Dy())))
}

type chunkCache struct {
	mem   *minecraft.MemPath
	level *minecraft.Level
	clear nbt.Tag
	cache map[uint16]nbt.Tag
}

func newCache() *chunkCache {
	mem := minecraft.NewMemPath()
	l, _ := minecraft.NewLevel(mem)

	bedrock := minecraft.Block{ID: 7}

	l.SetBlock(0, 0, 0, minecraft.Block{})
	l.Save()
	l.Close()
	clearChunk, _ := mem.GetChunk(0, 0)

	for i := int32(-16); i < 16; i++ {
		for j := int32(0); j < 255; j++ {
			for k := int32(-16); j < 16; j++ {
				l.SetBlock(i, j, k, bedrock)
			}
		}
	}
	l.Save()
	l.Close()
	mem.SetChunk(clearChunk)
	return &chunkCache{
		mem,
		l,
		clearChunk,
		make(map[uint16]nbt.Tag),
	}
}

func (c *chunkCache) getFromCache(x, z int32, terrain uint8, height int32) nbt.Tag {
	cacheID := uint16(terrain)<<8 | uint16(height)
	chunk, ok := c.cache[cacheID]
	if !ok {
		b := terrainBlocks[terrain].Base
		for j := int32(0); j < height; j++ {
			for i := int32(0); i < 16; i++ {
				for k := int32(0); k < 16; k++ {
					c.level.SetBlock(i, j, k, b)
				}
			}
		}
		c.level.Save()
		c.level.Close()
		chunk, _ = c.mem.GetChunk(0, 0)
		c.mem.SetChunk(c.clear)
		c.cache[cacheID] = chunk
	}
	ld := chunk.Data().(nbt.Compound).Get("Level").Data().(nbt.Compound)
	ld.Set(nbt.NewTag("xPos", nbt.Int(x)))
	ld.Set(nbt.NewTag("zPos", nbt.Int(z)))
	return chunk
}

func buildTerrain(mpath minecraft.Path, level *minecraft.Level, terrain, biomes *image.Paletted, height *image.Gray, c chan paint) error {
	b := terrain.Bounds()
	proceed := make(chan struct{}, 10)
	errChan := make(chan error, 1)
	go func() {
		defer close(proceed)
		cc := newCache()
		for j := 0; j < b.Max.Y; j += 16 {
			chunkZ := int32(j >> 4)
			for i := 0; i < b.Max.X; i += 16 {
				chunkX := int32(i >> 4)
				p := terrain.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Paletted)
				g := height.SubImage(image.Rect(i, j, i+16, j+16)).(*image.Gray)
				t := modeTerrain(p)
				h := int32(meanHeight(g))
				err := mpath.SetChunk(cc.getFromCache(chunkX, chunkZ, t, h))
				if err != nil {
					errChan <- err
					return
				}
				c <- paint{
					terrainColours[t],
					chunkX, chunkZ,
				}
				proceed <- struct{}{}
			}
		}
	}()
	for i := 0; i < (b.Max.X>>4)+2; i++ {
		<-proceed // get far enough ahead so all chunks are surrounded before shaping, to get correct lighting
	}
	select {
	case err := <-errChan:
		return err
	default:
	}
	for j := int32(0); j < int32(b.Max.Y); j += 16 {
		chunkZ := j >> 4
		for i := int32(0); i < int32(b.Max.X); i += 16 {
			chunkX := i >> 4
			var totalHeight int32
			for x := i; x < i+16; x++ {
				for z := j; z < j+16; z++ {
					if biomes != nil {
						level.SetBiome(x, z, biomeList[biomePalette.ColorIndexAt(int(x), int(z))])
					}
					h := int32(height.GrayAt(int(x), int(z)).Y)
					min := h
					for mz := int(z) - 1; mz <= int(z)+1; mz++ {
						if mz < 0 || mz >= b.Max.Y {
							continue
						}
						for mx := int(x) - 1; mx <= int(x)+1; mx++ {
							if mx < 0 || mx >= b.Max.X {
								continue
							}
							if mh := int32(height.GrayAt(mx, mz).Y); mh < min {
								min = mh
							}
						}
					}

					totalHeight += h
					y, err := level.GetHeight(x, z)
					if err != nil {
						return err
					}
					for ; y > h; y-- {
						level.SetBlock(x, y, z, minecraft.Block{})
					}
					t := terrainBlocks[terrain.ColorIndexAt(int(x), int(z))]
					for ; y > h-int32(t.TopLevel); y-- {
						level.SetBlock(x, y, z, t.Top)
					}
					for ; y >= min; y-- {
						level.SetBlock(x, y, z, t.Base)
					}
				}
			}
			c <- paint{
				color.Alpha{uint8(totalHeight >> 8)},
				chunkX, chunkZ,
			}
			select {
			case <-proceed:
			case err := <-errChan:
				return err
			}
		}
	}
	return nil
}
<|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/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]
package main

import (
	"encoding/json"
	"fmt"
	"mig.ninja/mig"
	"mig.ninja/mig/client"
	"os"
	"path"
	"time"
)

func getResultsStoragePath(nm string) (rdir string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getResultsStoragePath() -> %v", e)
		}
	}()
	tstamp := time.Now().UTC().Format("20060102")
	rdir = path.Join(ctx.Runner.RunDirectory, nm, "results", tstamp)

	_, err = os.Stat(rdir)
	if err != nil {
		if !os.IsNotExist(err) {
			panic(err)
		}
		err = os.MkdirAll(rdir, 0755)
		if err != nil {
			panic(err)
		}
	}
	return rdir, nil
}

func getResults(r mig.RunnerResult) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getResults() -> %v", e)
		}
	}()

	mlog("fetching results for %v/%.0f", r.EntityName, r.Action.ID)

	cli, err := client.NewClient(ctx.ClientConf, "mig-runner-results")
	if err != nil {
		panic(err)
	}
	results, err := cli.FetchActionResults(r.Action)
	if err != nil {
		panic(err)
	}

	// Store the results
	outpath, err := getResultsStoragePath(r.EntityName)
	if err != nil {
		panic(err)
	}
	outpath = path.Join(outpath, fmt.Sprintf("%.0f", r.Action.ID))
	fd, err := os.Create(outpath)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	r.Commands = results
	buf, err := json.Marshal(r)
	if err != nil {
		panic(err)
	}
	_, err = fd.Write(buf)
	if err != nil {
		panic(err)
	}
	mlog("wrote results for %.0f to %v", r.Action.ID, outpath)

	// If a plugin has been configured on the result set, call the
	// plugin on the data.
	if r.UsePlugin != "" {
		err = runPlugin(r)
		if err != nil {
			panic(err)
		}
	} else {
		mlog("no output plugin for %.0f, skipping plugin processing", r.Action.ID)
	}

	return nil
}

func processResults() {
	mlog("results processing routine started")

	reslist := make([]mig.RunnerResult, 0)

	for {
		timeout := false
		select {
		case nr := <-ctx.Channels.Results:
			mlog("monitoring result for %v/%.0f", nr.EntityName, nr.Action.ID)
			reslist = append(reslist, nr)
		case <-time.After(time.Duration(5) * time.Second):
			timeout = true
		}
		if !timeout {
			// Only attempt action results fetch if we are idle
			continue
		}

		// See if any actions have expired, if so grab the results
		oldres := reslist
		reslist = reslist[:0]
		for _, x := range oldres {
			if time.Now().After(x.Action.ExpireAfter) {
				err := getResults(x)
				if err != nil {
					mlog("results error for %v: %v", x.EntityName, err)
				}
				continue
			}
			reslist = append(reslist, x)
		}
	}
}
<commit_msg>[minor] add additional logging on results fetch<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/.
//
// Contributor: Aaron Meihm ameihm@mozilla.com [:alm]
package main

import (
	"encoding/json"
	"fmt"
	"mig.ninja/mig"
	"mig.ninja/mig/client"
	"os"
	"path"
	"time"
)

func getResultsStoragePath(nm string) (rdir string, err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getResultsStoragePath() -> %v", e)
		}
	}()
	tstamp := time.Now().UTC().Format("20060102")
	rdir = path.Join(ctx.Runner.RunDirectory, nm, "results", tstamp)

	_, err = os.Stat(rdir)
	if err != nil {
		if !os.IsNotExist(err) {
			panic(err)
		}
		err = os.MkdirAll(rdir, 0755)
		if err != nil {
			panic(err)
		}
	}
	return rdir, nil
}

func getResults(r mig.RunnerResult) (err error) {
	defer func() {
		if e := recover(); e != nil {
			err = fmt.Errorf("getResults() -> %v", e)
		}
	}()

	mlog("fetching results for %v/%.0f", r.EntityName, r.Action.ID)

	cli, err := client.NewClient(ctx.ClientConf, "mig-runner-results")
	if err != nil {
		panic(err)
	}
	results, err := cli.FetchActionResults(r.Action)
	if err != nil {
		panic(err)
	}
	mlog("%v/%.0f: fetched results, %v commands returned", r.EntityName, r.Action.ID, len(results))

	// Store the results
	outpath, err := getResultsStoragePath(r.EntityName)
	if err != nil {
		panic(err)
	}
	outpath = path.Join(outpath, fmt.Sprintf("%.0f", r.Action.ID))
	fd, err := os.Create(outpath)
	if err != nil {
		panic(err)
	}
	defer fd.Close()

	r.Commands = results
	buf, err := json.Marshal(r)
	if err != nil {
		panic(err)
	}
	_, err = fd.Write(buf)
	if err != nil {
		panic(err)
	}
	mlog("wrote results for %.0f to %v", r.Action.ID, outpath)

	// If a plugin has been configured on the result set, call the
	// plugin on the data.
	if r.UsePlugin != "" {
		err = runPlugin(r)
		if err != nil {
			panic(err)
		}
	} else {
		mlog("no output plugin for %.0f, skipping plugin processing", r.Action.ID)
	}

	return nil
}

func processResults() {
	mlog("results processing routine started")

	reslist := make([]mig.RunnerResult, 0)

	for {
		timeout := false
		select {
		case nr := <-ctx.Channels.Results:
			mlog("monitoring result for %v/%.0f", nr.EntityName, nr.Action.ID)
			reslist = append(reslist, nr)
		case <-time.After(time.Duration(5) * time.Second):
			timeout = true
		}
		if !timeout {
			// Only attempt action results fetch if we are idle
			continue
		}

		// See if any actions have expired, if so grab the results
		oldres := reslist
		reslist = reslist[:0]
		for _, x := range oldres {
			if time.Now().After(x.Action.ExpireAfter) {
				err := getResults(x)
				if err != nil {
					mlog("results error for %v: %v", x.EntityName, err)
				}
				continue
			}
			reslist = append(reslist, x)
		}
	}
}
<|endoftext|>


<commit_before>// Package usl provides functionality to build Universal Scalability Law models
// from sets of observed measurements.
package usl

import (
	"errors"
	"fmt"
	"math"

	"github.com/maorshutman/lm"
)

var (
	// ErrInsufficientMeasurements is returned when fewer than 6 measurements were provided.
	ErrInsufficientMeasurements = errors.New("usl: need at least 6 measurements")
)

// Model is a Universal Scalability Law model.
type Model struct {
	Sigma  float64 // The model's coefficient of contention, σ.
	Kappa  float64 // The model's coefficient of crosstalk/coherency, κ.
	Lambda float64 // The model's coefficient of performance, λ.
}

func (m *Model) String() string {
	return fmt.Sprintf("Model{σ=%v, κ=%v, λ=%v}", m.Sigma, m.Kappa, m.Lambda)
}

// ThroughputAtConcurrency returns the expected throughput given a number of concurrent workers,
// Concurrency(N).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 3".
func (m *Model) ThroughputAtConcurrency(n float64) float64 {
	return (m.Lambda * n) / (1 + (m.Sigma * (n - 1)) + (m.Kappa * n * (n - 1)))
}

// LatencyAtConcurrency returns the expected mean latency given a number of concurrent workers,
// R(N).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 6".
func (m *Model) LatencyAtConcurrency(n float64) float64 {
	return (1 + (m.Sigma * (n - 1)) + (m.Kappa * n * (n - 1))) / m.Lambda
}

// MaxConcurrency returns the maximum expected number of concurrent workers the system can handle,
// Nmax.
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 4".
func (m *Model) MaxConcurrency() float64 {
	return math.Floor(math.Sqrt((1 - m.Sigma) / m.Kappa))
}

// MaxThroughput returns the maximum expected throughput the system can handle, Xmax.
func (m Model) MaxThroughput() float64 {
	return m.ThroughputAtConcurrency(m.MaxConcurrency())
}

// LatencyAtThroughput returns the expected mean latency given a throughput, R(X).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 8".
func (m *Model) LatencyAtThroughput(x float64) float64 {
	return (m.Sigma - 1) / (m.Sigma*x - m.Lambda)
}

// ThroughputAtLatency returns the expected throughput given a mean latency, X(R).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 9".
func (m *Model) ThroughputAtLatency(r float64) float64 {
	a := 2 * m.Kappa * (2*m.Lambda*r + m.Sigma - 2)
	b := math.Sqrt(math.Pow(m.Sigma, 2) + math.Pow(m.Kappa, 2) + a)

	return (b - m.Kappa + m.Sigma) / (2.0 * m.Kappa * r)
}

// ConcurrencyAtLatency returns the expected number of concurrent workers given a mean latency,
// N(R).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 10".
func (m *Model) ConcurrencyAtLatency(r float64) float64 {
	a := 2 * m.Kappa * ((2 * m.Lambda * r) + m.Sigma - 2)
	b := math.Sqrt(math.Pow(m.Sigma, 2) + math.Pow(m.Kappa, 2) + a)

	return (m.Kappa - m.Sigma + b) / (2 * m.Kappa)
}

// ConcurrencyAtThroughput returns the expected number of concurrent workers at a particular
// throughput, N(X).
func (m *Model) ConcurrencyAtThroughput(x float64) float64 {
	return m.LatencyAtThroughput(x) * x
}

// ContentionConstrained returns true if the system is constrained by contention.
func (m *Model) ContentionConstrained() bool {
	return m.Sigma > m.Kappa
}

// CoherencyConstrained returns true if the system is constrained by coherency costs.
func (m *Model) CoherencyConstrained() bool {
	return m.Sigma < m.Kappa
}

// Limitless returns true if the system is linearly scalable.
func (m *Model) Limitless() bool {
	return m.Kappa == 0
}

const minMeasurements = 6

// Build returns a model whose parameters are generated from the given measurements.
//
// Finds a set of coefficients for the equation y = λx/(1+σ(x-1)+κx(x-1)) which best fit the
// observed values using unconstrained least-squares regression. The resulting values for λ, κ, and
// σ are the parameters of the returned model.
func Build(measurements []Measurement) (m *Model, err error) {
	if len(measurements) < minMeasurements {
		return nil, ErrInsufficientMeasurements
	}

	// Calculate an initial guess at the model parameters.
	init := []float64{0.1, 0.01, 0}

	// Use max(x/n) as initial lambda.
	for _, m := range measurements {
		v := m.Throughput / m.Concurrency
		if v > init[2] {
			init[2] = v
		}
	}

	// Calculate the residuals of a possible model.
	f := func(dst, x []float64) {
		model := Model{Sigma: x[0], Kappa: x[1], Lambda: x[2]}

		for i, v := range measurements {
			dst[i] = v.Throughput - model.ThroughputAtConcurrency(v.Concurrency)
		}
	}
	j := lm.NumJac{Func: f}

	// Formulate an LM problem.
	p := lm.LMProblem{
		Dim:        3,                 // Three parameters in the model.
		Size:       len(measurements), // Use all measurements to calculate residuals.
		Func:       f,                 // Reduce the residuals of model predictions to observations.
		Jac:        j.Jac,             // Use the default Jaccardian formula.
		InitParams: init,              // Use our initial guesses at parameters.
		Tau:        1e-6,              // Need a non-zero initial damping factor.
		Eps1:       1e-8,              // Small but non-zero values here prevent singular matrices.
		Eps2:       1e-8,
	}

	// Calculate the model parameters.
	results, err := lm.LM(p, nil)
	if err != nil {
		return nil, err
	}

	// Return the model.
	return &Model{
		Sigma:  results.X[0],
		Kappa:  results.X[1],
		Lambda: results.X[2],
	}, nil
}
<commit_msg>docfix<commit_after>// Package usl provides functionality to build Universal Scalability Law models
// from sets of observed measurements.
package usl

import (
	"errors"
	"fmt"
	"math"

	"github.com/maorshutman/lm"
)

var (
	// ErrInsufficientMeasurements is returned when fewer than 6 measurements were provided.
	ErrInsufficientMeasurements = errors.New("usl: need at least 6 measurements")
)

// Model is a Universal Scalability Law model.
type Model struct {
	Sigma  float64 // The model's coefficient of contention, σ.
	Kappa  float64 // The model's coefficient of crosstalk/coherency, κ.
	Lambda float64 // The model's coefficient of performance, λ.
}

func (m *Model) String() string {
	return fmt.Sprintf("Model{σ=%v, κ=%v, λ=%v}", m.Sigma, m.Kappa, m.Lambda)
}

// ThroughputAtConcurrency returns the expected throughput given a number of concurrent workers,
// Concurrency(N).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 3".
func (m *Model) ThroughputAtConcurrency(n float64) float64 {
	return (m.Lambda * n) / (1 + (m.Sigma * (n - 1)) + (m.Kappa * n * (n - 1)))
}

// LatencyAtConcurrency returns the expected mean latency given a number of concurrent workers,
// R(N).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 6".
func (m *Model) LatencyAtConcurrency(n float64) float64 {
	return (1 + (m.Sigma * (n - 1)) + (m.Kappa * n * (n - 1))) / m.Lambda
}

// MaxConcurrency returns the maximum expected number of concurrent workers the system can handle,
// Nmax.
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 4".
func (m *Model) MaxConcurrency() float64 {
	return math.Floor(math.Sqrt((1 - m.Sigma) / m.Kappa))
}

// MaxThroughput returns the maximum expected throughput the system can handle, Xmax.
func (m Model) MaxThroughput() float64 {
	return m.ThroughputAtConcurrency(m.MaxConcurrency())
}

// LatencyAtThroughput returns the expected mean latency given a throughput, R(X).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 8".
func (m *Model) LatencyAtThroughput(x float64) float64 {
	return (m.Sigma - 1) / (m.Sigma*x - m.Lambda)
}

// ThroughputAtLatency returns the expected throughput given a mean latency, X(R).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 9".
func (m *Model) ThroughputAtLatency(r float64) float64 {
	a := 2 * m.Kappa * (2*m.Lambda*r + m.Sigma - 2)
	b := math.Sqrt(math.Pow(m.Sigma, 2) + math.Pow(m.Kappa, 2) + a)

	return (b - m.Kappa + m.Sigma) / (2.0 * m.Kappa * r)
}

// ConcurrencyAtLatency returns the expected number of concurrent workers given a mean latency,
// N(R).
//
// See "Practical Scalability Analysis with the Universal Scalability Law, Equation 10".
func (m *Model) ConcurrencyAtLatency(r float64) float64 {
	a := 2 * m.Kappa * ((2 * m.Lambda * r) + m.Sigma - 2)
	b := math.Sqrt(math.Pow(m.Sigma, 2) + math.Pow(m.Kappa, 2) + a)

	return (m.Kappa - m.Sigma + b) / (2 * m.Kappa)
}

// ConcurrencyAtThroughput returns the expected number of concurrent workers at a particular
// throughput, N(X).
func (m *Model) ConcurrencyAtThroughput(x float64) float64 {
	return m.LatencyAtThroughput(x) * x
}

// ContentionConstrained returns true if the system is constrained by contention.
func (m *Model) ContentionConstrained() bool {
	return m.Sigma > m.Kappa
}

// CoherencyConstrained returns true if the system is constrained by coherency costs.
func (m *Model) CoherencyConstrained() bool {
	return m.Sigma < m.Kappa
}

// Limitless returns true if the system is linearly scalable.
func (m *Model) Limitless() bool {
	return m.Kappa == 0
}

const minMeasurements = 6

// Build returns a model whose parameters are generated from the given measurements.
//
// Finds a set of coefficients for the equation y = λx/(1+σ(x-1)+κx(x-1)) which best fit the
// observed values using unconstrained least-squares regression. The resulting values for λ, κ, and
// σ are the parameters of the returned model.
func Build(measurements []Measurement) (m *Model, err error) {
	if len(measurements) < minMeasurements {
		return nil, ErrInsufficientMeasurements
	}

	// Calculate an initial guess at the model parameters.
	init := []float64{0.1, 0.01, 0}

	// Use max(x/n) as initial lambda.
	for _, m := range measurements {
		v := m.Throughput / m.Concurrency
		if v > init[2] {
			init[2] = v
		}
	}

	// Calculate the residuals of a possible model.
	f := func(dst, x []float64) {
		model := Model{Sigma: x[0], Kappa: x[1], Lambda: x[2]}

		for i, v := range measurements {
			dst[i] = v.Throughput - model.ThroughputAtConcurrency(v.Concurrency)
		}
	}
	j := lm.NumJac{Func: f}

	// Formulate an LM problem.
	p := lm.LMProblem{
		Dim:        3,                 // Three parameters in the model.
		Size:       len(measurements), // Use all measurements to calculate residuals.
		Func:       f,                 // Reduce the residuals of model predictions to observations.
		Jac:        j.Jac,             // Approximate the Jacobian by finite differences.
		InitParams: init,              // Use our initial guesses at parameters.
		Tau:        1e-6,              // Need a non-zero initial damping factor.
		Eps1:       1e-8,              // Small but non-zero values here prevent singular matrices.
		Eps2:       1e-8,
	}

	// Calculate the model parameters.
	results, err := lm.LM(p, nil)
	if err != nil {
		return nil, err
	}

	// Return the model.
	return &Model{
		Sigma:  results.X[0],
		Kappa:  results.X[1],
		Lambda: results.X[2],
	}, nil
}
<|endoftext|>


<commit_before>package util

import "testing"

func Test_DoRequest(t *testing.T) {
	req := &AliyunRequest{
		"http://weibo.com",
	}
	_, err := req.DoGetRequest()
	if err != nil {
		t.Error("Request failed.")
	}
}

func Test_GetQueryFromMap(t *testing.T) {
	params := map[string]string{
		"user": "root",
		"pass": "test",
	}
	query := GetQueryFromMap(params)
	if query != "user=root&pass=test" && query != "pass=test&user=root" {
		t.Error(query + "Query string is incorrect.")
	}
}
<commit_msg>	modified:   util/request_test.go<commit_after>package util

import "testing"

func Test_DoRequest(t *testing.T) {
	req := &AliyunRequest{
		"http://weibo.com",
	}
	_, err := req.DoGetRequest()
	if err != nil {
		t.Error("Request failed.")
	}
}

func Test_GetQueryFromMap(t *testing.T) {
	params := map[string]string{
		"user": "root",
		"pass": "test",
	}
	query := GetQueryFromMap(params)
	if query != "user=root&pass=test" && query != "pass=test&user=root" {
		t.Error("Query string is incorrect.")
	}
}
<|endoftext|>


<commit_before>package idxfile

import (
	"bytes"
	"io"
	"sort"

	encbin "encoding/binary"

	"gopkg.in/src-d/go-git.v4/plumbing"
)

const (
	// VersionSupported is the only idx version supported.
	VersionSupported = 2

	noMapping = -1
)

var (
	idxHeader = []byte{255, 't', 'O', 'c'}
)

// Index represents an index of a packfile.
type Index interface {
	// Contains checks whether the given hash is in the index.
	Contains(h plumbing.Hash) (bool, error)
	// FindOffset finds the offset in the packfile for the object with
	// the given hash.
	FindOffset(h plumbing.Hash) (int64, error)
	// FindCRC32 finds the CRC32 of the object with the given hash.
	FindCRC32(h plumbing.Hash) (uint32, error)
	// FindHash finds the hash for the object with the given offset.
	FindHash(o int64) (plumbing.Hash, error)
	// Count returns the number of entries in the index.
	Count() (int64, error)
	// Entries returns an iterator to retrieve all index entries.
	Entries() (EntryIter, error)
	// EntriesByOffset returns an iterator to retrieve all index entries ordered
	// by offset.
	EntriesByOffset() (EntryIter, error)
}

// MemoryIndex is the in memory representation of an idx file.
type MemoryIndex struct {
	Version uint32
	Fanout  [256]uint32
	// FanoutMapping maps the position in the fanout table to the position
	// in the Names, Offset32 and CRC32 slices. This improves the memory
	// usage by not needing an array with unnecessary empty slots.
	FanoutMapping    [256]int
	Names            [][]byte
	Offset32         [][]byte
	CRC32            [][]byte
	Offset64         []byte
	PackfileChecksum [20]byte
	IdxChecksum      [20]byte

	offsetHash       map[int64]plumbing.Hash
	offsetHashIsFull bool
}

var _ Index = (*MemoryIndex)(nil)

// NewMemoryIndex returns an instance of a new MemoryIndex.
func NewMemoryIndex() *MemoryIndex {
	return &MemoryIndex{}
}

func (idx *MemoryIndex) findHashIndex(h plumbing.Hash) (int, bool) {
	k := idx.FanoutMapping[h[0]]
	if k == noMapping {
		return 0, false
	}

	if len(idx.Names) <= k {
		return 0, false
	}

	data := idx.Names[k]
	high := uint64(len(idx.Offset32[k])) >> 2
	if high == 0 {
		return 0, false
	}

	low := uint64(0)
	for {
		mid := (low + high) >> 1
		offset := mid * objectIDLength

		cmp := bytes.Compare(h[:], data[offset:offset+objectIDLength])
		if cmp < 0 {
			high = mid
		} else if cmp == 0 {
			return int(mid), true
		} else {
			low = mid + 1
		}

		if low >= high {
			break
		}
	}

	return 0, false
}

// Contains implements the Index interface.
func (idx *MemoryIndex) Contains(h plumbing.Hash) (bool, error) {
	_, ok := idx.findHashIndex(h)
	return ok, nil
}

// FindOffset implements the Index interface.
func (idx *MemoryIndex) FindOffset(h plumbing.Hash) (int64, error) {
	if len(idx.FanoutMapping) <= int(h[0]) {
		return 0, plumbing.ErrObjectNotFound
	}

	k := idx.FanoutMapping[h[0]]
	i, ok := idx.findHashIndex(h)
	if !ok {
		return 0, plumbing.ErrObjectNotFound
	}

	offset := idx.getOffset(k, i)

	if !idx.offsetHashIsFull {
		// Save the offset for reverse lookup
		if idx.offsetHash == nil {
			idx.offsetHash = make(map[int64]plumbing.Hash)
		}
		idx.offsetHash[int64(offset)] = h
	}

	return int64(offset), nil
}

const isO64Mask = uint64(1) << 31

func (idx *MemoryIndex) getOffset(firstLevel, secondLevel int) uint64 {
	offset := secondLevel << 2
	ofs := encbin.BigEndian.Uint32(idx.Offset32[firstLevel][offset : offset+4])

	if (uint64(ofs) & isO64Mask) != 0 {
		offset := 8 * (uint64(ofs) & ^isO64Mask)
		n := encbin.BigEndian.Uint64(idx.Offset64[offset : offset+8])
		return n
	}

	return uint64(ofs)
}

// FindCRC32 implements the Index interface.
func (idx *MemoryIndex) FindCRC32(h plumbing.Hash) (uint32, error) {
	k := idx.FanoutMapping[h[0]]
	i, ok := idx.findHashIndex(h)
	if !ok {
		return 0, plumbing.ErrObjectNotFound
	}

	return idx.getCRC32(k, i), nil
}

func (idx *MemoryIndex) getCRC32(firstLevel, secondLevel int) uint32 {
	offset := secondLevel << 2
	return encbin.BigEndian.Uint32(idx.CRC32[firstLevel][offset : offset+4])
}

// FindHash implements the Index interface.
func (idx *MemoryIndex) FindHash(o int64) (plumbing.Hash, error) {
	var hash plumbing.Hash
	var ok bool

	if idx.offsetHash != nil {
		if hash, ok = idx.offsetHash[o]; ok {
			return hash, nil
		}
	}

	// Lazily generate the reverse offset/hash map if required.
	if !idx.offsetHashIsFull || idx.offsetHash == nil {
		if err := idx.genOffsetHash(); err != nil {
			return plumbing.ZeroHash, err
		}

		hash, ok = idx.offsetHash[o]
	}

	if !ok {
		return plumbing.ZeroHash, plumbing.ErrObjectNotFound
	}

	return hash, nil
}

// genOffsetHash generates the offset/hash mapping for reverse search.
func (idx *MemoryIndex) genOffsetHash() error {
	count, err := idx.Count()
	if err != nil {
		return err
	}

	idx.offsetHash = make(map[int64]plumbing.Hash, count)
	idx.offsetHashIsFull = true

	iter, err := idx.Entries()
	if err != nil {
		return err
	}

	for {
		entry, err := iter.Next()
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return err
		}

		idx.offsetHash[int64(entry.Offset)] = entry.Hash
	}
}

// Count implements the Index interface.
func (idx *MemoryIndex) Count() (int64, error) {
	return int64(idx.Fanout[fanout-1]), nil
}

// Entries implements the Index interface.
func (idx *MemoryIndex) Entries() (EntryIter, error) {
	return &idxfileEntryIter{idx, 0, 0, 0}, nil
}

// EntriesByOffset implements the Index interface.
func (idx *MemoryIndex) EntriesByOffset() (EntryIter, error) {
	count, err := idx.Count()
	if err != nil {
		return nil, err
	}

	iter := &idxfileEntryOffsetIter{
		entries: make(entriesByOffset, count),
	}

	entries, err := idx.Entries()
	if err != nil {
		return nil, err
	}

	for pos := 0; int64(pos) < count; pos++ {
		entry, err := entries.Next()
		if err != nil {
			return nil, err
		}

		iter.entries[pos] = entry
	}

	sort.Sort(iter.entries)

	return iter, nil
}

// EntryIter is an iterator that will return the entries in a packfile index.
type EntryIter interface {
	// Next returns the next entry in the packfile index.
	Next() (*Entry, error)
	// Close closes the iterator.
	Close() error
}

type idxfileEntryIter struct {
	idx                     *MemoryIndex
	total                   int
	firstLevel, secondLevel int
}

func (i *idxfileEntryIter) Next() (*Entry, error) {
	for {
		if i.firstLevel >= fanout {
			return nil, io.EOF
		}

		if i.total >= int(i.idx.Fanout[i.firstLevel]) {
			i.firstLevel++
			i.secondLevel = 0
			continue
		}

		entry := new(Entry)
		ofs := i.secondLevel * objectIDLength
		copy(entry.Hash[:], i.idx.Names[i.idx.FanoutMapping[i.firstLevel]][ofs:])

		pos := i.idx.FanoutMapping[entry.Hash[0]]

		entry.Offset = i.idx.getOffset(pos, i.secondLevel)
		entry.CRC32 = i.idx.getCRC32(pos, i.secondLevel)

		i.secondLevel++
		i.total++

		return entry, nil
	}
}

func (i *idxfileEntryIter) Close() error {
	i.firstLevel = fanout
	return nil
}

// Entry is the in memory representation of an object entry in the idx file.
type Entry struct {
	Hash   plumbing.Hash
	CRC32  uint32
	Offset uint64
}

type idxfileEntryOffsetIter struct {
	entries entriesByOffset
	pos     int
}

func (i *idxfileEntryOffsetIter) Next() (*Entry, error) {
	if i.pos >= len(i.entries) {
		return nil, io.EOF
	}

	entry := i.entries[i.pos]
	i.pos++

	return entry, nil
}

func (i *idxfileEntryOffsetIter) Close() error {
	i.pos = len(i.entries) + 1
	return nil
}

type entriesByOffset []*Entry

func (o entriesByOffset) Len() int {
	return len(o)
}

func (o entriesByOffset) Less(i int, j int) bool {
	return o[i].Offset < o[j].Offset
}

func (o entriesByOffset) Swap(i int, j int) {
	o[i], o[j] = o[j], o[i]
}
<commit_msg>plumbing: format/idxfile, save another 18% of time in genOffsetHash by not using iterator and not loading CRC<commit_after>package idxfile

import (
	"bytes"
	"io"
	"sort"

	encbin "encoding/binary"

	"gopkg.in/src-d/go-git.v4/plumbing"
)

const (
	// VersionSupported is the only idx version supported.
	VersionSupported = 2

	noMapping = -1
)

var (
	idxHeader = []byte{255, 't', 'O', 'c'}
)

// Index represents an index of a packfile.
type Index interface {
	// Contains checks whether the given hash is in the index.
	Contains(h plumbing.Hash) (bool, error)
	// FindOffset finds the offset in the packfile for the object with
	// the given hash.
	FindOffset(h plumbing.Hash) (int64, error)
	// FindCRC32 finds the CRC32 of the object with the given hash.
	FindCRC32(h plumbing.Hash) (uint32, error)
	// FindHash finds the hash for the object with the given offset.
	FindHash(o int64) (plumbing.Hash, error)
	// Count returns the number of entries in the index.
	Count() (int64, error)
	// Entries returns an iterator to retrieve all index entries.
	Entries() (EntryIter, error)
	// EntriesByOffset returns an iterator to retrieve all index entries ordered
	// by offset.
	EntriesByOffset() (EntryIter, error)
}

// MemoryIndex is the in memory representation of an idx file.
type MemoryIndex struct {
	Version uint32
	Fanout  [256]uint32
	// FanoutMapping maps the position in the fanout table to the position
	// in the Names, Offset32 and CRC32 slices. This improves the memory
	// usage by not needing an array with unnecessary empty slots.
	FanoutMapping    [256]int
	Names            [][]byte
	Offset32         [][]byte
	CRC32            [][]byte
	Offset64         []byte
	PackfileChecksum [20]byte
	IdxChecksum      [20]byte

	offsetHash       map[int64]plumbing.Hash
	offsetHashIsFull bool
}

var _ Index = (*MemoryIndex)(nil)

// NewMemoryIndex returns an instance of a new MemoryIndex.
func NewMemoryIndex() *MemoryIndex {
	return &MemoryIndex{}
}

func (idx *MemoryIndex) findHashIndex(h plumbing.Hash) (int, bool) {
	k := idx.FanoutMapping[h[0]]
	if k == noMapping {
		return 0, false
	}

	if len(idx.Names) <= k {
		return 0, false
	}

	data := idx.Names[k]
	high := uint64(len(idx.Offset32[k])) >> 2
	if high == 0 {
		return 0, false
	}

	low := uint64(0)
	for {
		mid := (low + high) >> 1
		offset := mid * objectIDLength

		cmp := bytes.Compare(h[:], data[offset:offset+objectIDLength])
		if cmp < 0 {
			high = mid
		} else if cmp == 0 {
			return int(mid), true
		} else {
			low = mid + 1
		}

		if low >= high {
			break
		}
	}

	return 0, false
}

// Contains implements the Index interface.
func (idx *MemoryIndex) Contains(h plumbing.Hash) (bool, error) {
	_, ok := idx.findHashIndex(h)
	return ok, nil
}

// FindOffset implements the Index interface.
func (idx *MemoryIndex) FindOffset(h plumbing.Hash) (int64, error) {
	if len(idx.FanoutMapping) <= int(h[0]) {
		return 0, plumbing.ErrObjectNotFound
	}

	k := idx.FanoutMapping[h[0]]
	i, ok := idx.findHashIndex(h)
	if !ok {
		return 0, plumbing.ErrObjectNotFound
	}

	offset := idx.getOffset(k, i)

	if !idx.offsetHashIsFull {
		// Save the offset for reverse lookup
		if idx.offsetHash == nil {
			idx.offsetHash = make(map[int64]plumbing.Hash)
		}
		idx.offsetHash[int64(offset)] = h
	}

	return int64(offset), nil
}

const isO64Mask = uint64(1) << 31

func (idx *MemoryIndex) getOffset(firstLevel, secondLevel int) uint64 {
	offset := secondLevel << 2
	ofs := encbin.BigEndian.Uint32(idx.Offset32[firstLevel][offset : offset+4])

	if (uint64(ofs) & isO64Mask) != 0 {
		offset := 8 * (uint64(ofs) & ^isO64Mask)
		n := encbin.BigEndian.Uint64(idx.Offset64[offset : offset+8])
		return n
	}

	return uint64(ofs)
}

// FindCRC32 implements the Index interface.
func (idx *MemoryIndex) FindCRC32(h plumbing.Hash) (uint32, error) {
	k := idx.FanoutMapping[h[0]]
	i, ok := idx.findHashIndex(h)
	if !ok {
		return 0, plumbing.ErrObjectNotFound
	}

	return idx.getCRC32(k, i), nil
}

func (idx *MemoryIndex) getCRC32(firstLevel, secondLevel int) uint32 {
	offset := secondLevel << 2
	return encbin.BigEndian.Uint32(idx.CRC32[firstLevel][offset : offset+4])
}

// FindHash implements the Index interface.
func (idx *MemoryIndex) FindHash(o int64) (plumbing.Hash, error) {
	var hash plumbing.Hash
	var ok bool

	if idx.offsetHash != nil {
		if hash, ok = idx.offsetHash[o]; ok {
			return hash, nil
		}
	}

	// Lazily generate the reverse offset/hash map if required.
	if !idx.offsetHashIsFull || idx.offsetHash == nil {
		if err := idx.genOffsetHash(); err != nil {
			return plumbing.ZeroHash, err
		}

		hash, ok = idx.offsetHash[o]
	}

	if !ok {
		return plumbing.ZeroHash, plumbing.ErrObjectNotFound
	}

	return hash, nil
}

// genOffsetHash generates the offset/hash mapping for reverse search.
func (idx *MemoryIndex) genOffsetHash() error {
	count, err := idx.Count()
	if err != nil {
		return err
	}

	idx.offsetHash = make(map[int64]plumbing.Hash, count)
	idx.offsetHashIsFull = true

	var hash plumbing.Hash
	i := uint32(0)
	for firstLevel, fanoutValue := range idx.Fanout {
		pos := idx.FanoutMapping[firstLevel]
		for secondLevel := uint32(0); i < fanoutValue; i++ {
			copy(hash[:], idx.Names[pos][secondLevel*objectIDLength:])
			offset := int64(idx.getOffset(pos, int(secondLevel)))
			idx.offsetHash[offset] = hash
			secondLevel++
		}
	}

	return nil
}

// Count implements the Index interface.
func (idx *MemoryIndex) Count() (int64, error) {
	return int64(idx.Fanout[fanout-1]), nil
}

// Entries implements the Index interface.
func (idx *MemoryIndex) Entries() (EntryIter, error) {
	return &idxfileEntryIter{idx, 0, 0, 0}, nil
}

// EntriesByOffset implements the Index interface.
func (idx *MemoryIndex) EntriesByOffset() (EntryIter, error) {
	count, err := idx.Count()
	if err != nil {
		return nil, err
	}

	iter := &idxfileEntryOffsetIter{
		entries: make(entriesByOffset, count),
	}

	entries, err := idx.Entries()
	if err != nil {
		return nil, err
	}

	for pos := 0; int64(pos) < count; pos++ {
		entry, err := entries.Next()
		if err != nil {
			return nil, err
		}

		iter.entries[pos] = entry
	}

	sort.Sort(iter.entries)

	return iter, nil
}

// EntryIter is an iterator that will return the entries in a packfile index.
type EntryIter interface {
	// Next returns the next entry in the packfile index.
	Next() (*Entry, error)
	// Close closes the iterator.
	Close() error
}

type idxfileEntryIter struct {
	idx                     *MemoryIndex
	total                   int
	firstLevel, secondLevel int
}

func (i *idxfileEntryIter) Next() (*Entry, error) {
	for {
		if i.firstLevel >= fanout {
			return nil, io.EOF
		}

		if i.total >= int(i.idx.Fanout[i.firstLevel]) {
			i.firstLevel++
			i.secondLevel = 0
			continue
		}

		entry := new(Entry)
		ofs := i.secondLevel * objectIDLength
		copy(entry.Hash[:], i.idx.Names[i.idx.FanoutMapping[i.firstLevel]][ofs:])

		pos := i.idx.FanoutMapping[entry.Hash[0]]

		entry.Offset = i.idx.getOffset(pos, i.secondLevel)
		entry.CRC32 = i.idx.getCRC32(pos, i.secondLevel)

		i.secondLevel++
		i.total++

		return entry, nil
	}
}

func (i *idxfileEntryIter) Close() error {
	i.firstLevel = fanout
	return nil
}

// Entry is the in memory representation of an object entry in the idx file.
type Entry struct {
	Hash   plumbing.Hash
	CRC32  uint32
	Offset uint64
}

type idxfileEntryOffsetIter struct {
	entries entriesByOffset
	pos     int
}

func (i *idxfileEntryOffsetIter) Next() (*Entry, error) {
	if i.pos >= len(i.entries) {
		return nil, io.EOF
	}

	entry := i.entries[i.pos]
	i.pos++

	return entry, nil
}

func (i *idxfileEntryOffsetIter) Close() error {
	i.pos = len(i.entries) + 1
	return nil
}

type entriesByOffset []*Entry

func (o entriesByOffset) Len() int {
	return len(o)
}

func (o entriesByOffset) Less(i int, j int) bool {
	return o[i].Offset < o[j].Offset
}

func (o entriesByOffset) Swap(i int, j int) {
	o[i], o[j] = o[j], o[i]
}
<|endoftext|>


<commit_before>package queue

import (
	"code.google.com/p/go-uuid/uuid"

	"os"
	"path"
	"testing"
)

func TestStorage(t *testing.T) {
	var (
		directory   = path.Join(os.TempDir(), uuid.New())
		data        = []byte("test")
		numMessages = 5
		messageIDs  = make([]string, numMessages)
		m           = &Message{}
	)
	if err := os.MkdirAll(directory, 0755); err != nil {
		t.Fatal(err)
	}
	defer os.RemoveAll(directory)
	if s, messages, err := NewStorage(directory); err == nil {
		if len(messages) != 0 {
			t.Fatalf("%d != 0", len(messages))
		}
		if w, id, err := s.NewBody(); err == nil {
			if _, err := w.Write(data); err != nil {
				t.Fatal(err)
			}
			if err := w.Close(); err != nil {
				t.Fatal(err)
			}
			m.Body = id
		} else {
			t.Fatal(err)
		}
		for i := 0; i < numMessages; i++ {
			if id, err := s.NewMessage(m); err == nil {
				messageIDs[i] = id
			} else {
				t.Fatal(err)
			}
		}
	} else {
		t.Fatal(err)
	}
	if s, messages, err := NewStorage(directory); err == nil {
		if len(messages) != numMessages {
			t.Fatalf("%d != %d", len(messages), numMessages)
		}
		for _, id := range messageIDs {
			if err := s.DeleteMessage(id); err != nil {
				t.Fatal(err)
			}
		}
		if _, err := s.GetBody(m.Body); err != InvalidID {
			t.Fatalf("%v != %v", err, InvalidID)
		}
	} else {
		t.Fatal(err)
	}
}
<commit_msg>Switched Storage tests to using TempDir.<commit_after>package queue

import (
	"github.com/nathan-osman/go-cannon/util"

	"testing"
)

func TestStorage(t *testing.T) {
	var (
		data        = []byte("test")
		numMessages = 5
		messageIDs  = make([]string, numMessages)
		m           = &Message{}
	)
	if d, err := util.NewTempDir(); err == nil {
		defer d.Delete()
		if s, messages, err := NewStorage(d.Path); err == nil {
			if len(messages) != 0 {
				t.Fatalf("%d != 0", len(messages))
			}
			if w, id, err := s.NewBody(); err == nil {
				if _, err := w.Write(data); err != nil {
					t.Fatal(err)
				}
				if err := w.Close(); err != nil {
					t.Fatal(err)
				}
				m.Body = id
			} else {
				t.Fatal(err)
			}
			for i := 0; i < numMessages; i++ {
				if id, err := s.NewMessage(m); err == nil {
					messageIDs[i] = id
				} else {
					t.Fatal(err)
				}
			}
		} else {
			t.Fatal(err)
		}
		if s, messages, err := NewStorage(d.Path); err == nil {
			if len(messages) != numMessages {
				t.Fatalf("%d != %d", len(messages), numMessages)
			}
			for _, id := range messageIDs {
				if err := s.DeleteMessage(id); err != nil {
					t.Fatal(err)
				}
			}
			if _, err := s.GetBody(m.Body); err != InvalidID {
				t.Fatalf("%v != %v", err, InvalidID)
			}
		} else {
			t.Fatal(err)
		}
	} else {
		t.Fatal(err)
	}
}
<|endoftext|>