<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package persistence

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/resource"
)

const (
	resourcesC = "resources"
)

const stagedIDSuffix = "#staged"

// resourceID converts an external resource ID into an internal one.
func (p Persistence) resourceID(id, serviceID string) string {
	return fmt.Sprintf("resource#%s#%s", serviceID, id)
}

// stagedID converts an external resource ID into an internal staged one.
func (p Persistence) stagedID(id, serviceID string) string {
	return p.resourceID(id, serviceID) + stagedIDSuffix
}

func (p Persistence) newStagedResourceOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newStagedDoc(id, serviceID, res)

	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
}

func (p Persistence) newEnsureStagedSameOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newStagedDoc(id, serviceID, res)

	// Other than cause the txn to abort, we don't do anything here.
	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: doc, // TODO(ericsnow) Is this okay?
	}}
}

func (p Persistence) newRemoveStagedOps(id, serviceID string) []txn.Op {
	fullID := p.stagedID(id, serviceID)

	// We don't assert that it exists. We want "missing" to be a noop.
	return []txn.Op{{
		C:      resourcesC,
		Id:     fullID,
		Remove: true,
	}}
}

func (p Persistence) newInsertResourceOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newResourceDoc(id, serviceID, res)

	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
}

func (p Persistence) newUpdateResourceOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newResourceDoc(id, serviceID, res)

	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: txn.DocExists,
		Update: doc,
	}}
}

// newResourceDoc generates a doc that represents the given resource.
func (p Persistence) newResourceDoc(id, serviceID string, res resource.Resource) *resourceDoc {
	fullID := p.resourceID(id, serviceID)
	return resource2doc(fullID, serviceID, res)
}

// newStagedDoc generates a staging doc that represents the given resource.
func (p Persistence) newStagedDoc(id, serviceID string, res resource.Resource) *resourceDoc {
	stagedID := p.stagedID(id, serviceID)
	return resource2doc(stagedID, serviceID, res)
}

// resources returns the resource docs for the given service.
func (p Persistence) resources(serviceID string) ([]resourceDoc, error) {
	var docs []resourceDoc
	query := bson.D{{"service-id", serviceID}}
	if err := p.base.All(resourcesC, query, &docs); err != nil {
		return nil, errors.Trace(err)
	}
	return docs, nil
}

// resourceDoc is the top-level document for resources.
type resourceDoc struct {
	DocID     string `bson:"_id"`
	EnvUUID   string `bson:"env-uuid"`
	ServiceID string `bson:"service-id"`

	Name    string `bson:"name"`
	Type    string `bson:"type"`
	Path    string `bson:"path"`
	Comment string `bson:"comment"`

	Origin      string `bson:"origin"`
	Revision    int    `bson:"revision"`
	Fingerprint []byte `bson:"fingerprint"`
	Size        int64  `bson:"size"`

	Username  string    `bson:"username"`
	Timestamp time.Time `bson:"timestamp-when-added"`
}

// resource2doc converts the resource into a DB doc.
func resource2doc(id, serviceID string, res resource.Resource) *resourceDoc {
	// TODO(ericsnow) We may need to limit the resolution of timestamps
	// in order to avoid some conversion problems from Mongo.
	return &resourceDoc{
		DocID:     id,
		ServiceID: serviceID,

		Name:    res.Name,
		Type:    res.Type.String(),
		Path:    res.Path,
		Comment: res.Comment,

		Origin:      res.Origin.String(),
		Revision:    res.Revision,
		Fingerprint: res.Fingerprint.Bytes(),
		Size:        res.Size,

		Username:  res.Username,
		Timestamp: res.Timestamp,
	}
}

// doc2resource returns the resource.Resource represented by the doc.
func doc2resource(doc resourceDoc) (resource.Resource, error) {
	var res resource.Resource

	resType, err := charmresource.ParseType(doc.Type)
	if err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}

	origin, err := charmresource.ParseOrigin(doc.Origin)
	if err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}

	fp, err := charmresource.NewFingerprint(doc.Fingerprint)
	if err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}

	res = resource.Resource{
		Resource: charmresource.Resource{
			Meta: charmresource.Meta{
				Name:    doc.Name,
				Type:    resType,
				Path:    doc.Path,
				Comment: doc.Comment,
			},
			Origin:      origin,
			Revision:    doc.Revision,
			Fingerprint: fp,
			Size:        doc.Size,
		},
		Username:  doc.Username,
		Timestamp: doc.Timestamp,
	}
	if err := res.Validate(); err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}
	return res, nil
}
<commit_msg>Consolidate consts.<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package persistence

import (
	"fmt"
	"time"

	"github.com/juju/errors"
	charmresource "gopkg.in/juju/charm.v6-unstable/resource"
	"gopkg.in/mgo.v2/bson"
	"gopkg.in/mgo.v2/txn"

	"github.com/juju/juju/resource"
)

const (
	resourcesC = "resources"

	stagedIDSuffix = "#staged"
)

// resourceID converts an external resource ID into an internal one.
func (p Persistence) resourceID(id, serviceID string) string {
	return fmt.Sprintf("resource#%s#%s", serviceID, id)
}

// stagedID converts an external resource ID into an internal staged one.
func (p Persistence) stagedID(id, serviceID string) string {
	return p.resourceID(id, serviceID) + stagedIDSuffix
}

func (p Persistence) newStagedResourceOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newStagedDoc(id, serviceID, res)

	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
}

func (p Persistence) newEnsureStagedSameOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newStagedDoc(id, serviceID, res)

	// Other than cause the txn to abort, we don't do anything here.
	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: doc, // TODO(ericsnow) Is this okay?
	}}
}

func (p Persistence) newRemoveStagedOps(id, serviceID string) []txn.Op {
	fullID := p.stagedID(id, serviceID)

	// We don't assert that it exists. We want "missing" to be a noop.
	return []txn.Op{{
		C:      resourcesC,
		Id:     fullID,
		Remove: true,
	}}
}

func (p Persistence) newInsertResourceOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newResourceDoc(id, serviceID, res)

	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: txn.DocMissing,
		Insert: doc,
	}}
}

func (p Persistence) newUpdateResourceOps(id, serviceID string, res resource.Resource) []txn.Op {
	doc := p.newResourceDoc(id, serviceID, res)

	return []txn.Op{{
		C:      resourcesC,
		Id:     doc.DocID,
		Assert: txn.DocExists,
		Update: doc,
	}}
}

// newResourceDoc generates a doc that represents the given resource.
func (p Persistence) newResourceDoc(id, serviceID string, res resource.Resource) *resourceDoc {
	fullID := p.resourceID(id, serviceID)
	return resource2doc(fullID, serviceID, res)
}

// newStagedDoc generates a staging doc that represents the given resource.
func (p Persistence) newStagedDoc(id, serviceID string, res resource.Resource) *resourceDoc {
	stagedID := p.stagedID(id, serviceID)
	return resource2doc(stagedID, serviceID, res)
}

// resources returns the resource docs for the given service.
func (p Persistence) resources(serviceID string) ([]resourceDoc, error) {
	var docs []resourceDoc
	query := bson.D{{"service-id", serviceID}}
	if err := p.base.All(resourcesC, query, &docs); err != nil {
		return nil, errors.Trace(err)
	}
	return docs, nil
}

// resourceDoc is the top-level document for resources.
type resourceDoc struct {
	DocID     string `bson:"_id"`
	EnvUUID   string `bson:"env-uuid"`
	ServiceID string `bson:"service-id"`

	Name    string `bson:"name"`
	Type    string `bson:"type"`
	Path    string `bson:"path"`
	Comment string `bson:"comment"`

	Origin      string `bson:"origin"`
	Revision    int    `bson:"revision"`
	Fingerprint []byte `bson:"fingerprint"`
	Size        int64  `bson:"size"`

	Username  string    `bson:"username"`
	Timestamp time.Time `bson:"timestamp-when-added"`
}

// resource2doc converts the resource into a DB doc.
func resource2doc(id, serviceID string, res resource.Resource) *resourceDoc {
	// TODO(ericsnow) We may need to limit the resolution of timestamps
	// in order to avoid some conversion problems from Mongo.
	return &resourceDoc{
		DocID:     id,
		ServiceID: serviceID,

		Name:    res.Name,
		Type:    res.Type.String(),
		Path:    res.Path,
		Comment: res.Comment,

		Origin:      res.Origin.String(),
		Revision:    res.Revision,
		Fingerprint: res.Fingerprint.Bytes(),
		Size:        res.Size,

		Username:  res.Username,
		Timestamp: res.Timestamp,
	}
}

// doc2resource returns the resource.Resource represented by the doc.
func doc2resource(doc resourceDoc) (resource.Resource, error) {
	var res resource.Resource

	resType, err := charmresource.ParseType(doc.Type)
	if err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}

	origin, err := charmresource.ParseOrigin(doc.Origin)
	if err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}

	fp, err := charmresource.NewFingerprint(doc.Fingerprint)
	if err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}

	res = resource.Resource{
		Resource: charmresource.Resource{
			Meta: charmresource.Meta{
				Name:    doc.Name,
				Type:    resType,
				Path:    doc.Path,
				Comment: doc.Comment,
			},
			Origin:      origin,
			Revision:    doc.Revision,
			Fingerprint: fp,
			Size:        doc.Size,
		},
		Username:  doc.Username,
		Timestamp: doc.Timestamp,
	}
	if err := res.Validate(); err != nil {
		return res, errors.Annotate(err, "got invalid data from DB")
	}
	return res, nil
}
<|endoftext|>


<commit_before>package agent

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/nomad"
	sconfig "github.com/hashicorp/nomad/nomad/structs/config"
)

func getPort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}

func tmpDir(t testing.TB) string {
	dir, err := ioutil.TempDir("", "nomad")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	return dir
}

func makeAgent(t testing.TB, cb func(*Config)) (string, *Agent) {
	dir := tmpDir(t)
	conf := DevConfig()

	// Customize the server configuration
	config := nomad.DefaultConfig()
	conf.NomadConfig = config

	// Bind and set ports
	conf.BindAddr = "127.0.0.1"
	conf.Ports = &Ports{
		HTTP: getPort(),
		RPC:  getPort(),
		Serf: getPort(),
	}
	conf.NodeName = fmt.Sprintf("Node %d", conf.Ports.RPC)
	conf.Consul = sconfig.DefaultConsulConfig()
	conf.Vault.Enabled = new(bool)

	// Tighten the Serf timing
	config.SerfConfig.MemberlistConfig.SuspicionMult = 2
	config.SerfConfig.MemberlistConfig.RetransmitMult = 2
	config.SerfConfig.MemberlistConfig.ProbeTimeout = 50 * time.Millisecond
	config.SerfConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	config.SerfConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	// Tighten the Raft timing
	config.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond
	config.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond
	config.RaftConfig.ElectionTimeout = 40 * time.Millisecond
	config.RaftConfig.StartAsLeader = true
	config.RaftTimeout = 500 * time.Millisecond

	if cb != nil {
		cb(conf)
	}

	agent, err := NewAgent(conf, os.Stderr)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf("err: %v", err)
	}
	return dir, agent
}

func TestAgent_RPCPing(t *testing.T) {
	dir, agent := makeAgent(t, nil)
	defer os.RemoveAll(dir)
	defer agent.Shutdown()

	var out struct{}
	if err := agent.RPC("Status.Ping", struct{}{}, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestAgent_ServerConfig(t *testing.T) {
	conf := DefaultConfig()
	a := &Agent{config: conf}
	testlogger := func(s string) {
		t.Log(s)
	}
	dev := false

	// Returns error on bad serf addr
	conf.AdvertiseAddrs.Serf = "nope"
	_, err := a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "Failed to parse Serf") {
		t.Fatalf("expected serf address error, got: %#v", err)
	}
	conf.AdvertiseAddrs.Serf = "127.0.0.1:4000"

	// Returns error on bad rpc addr
	conf.AdvertiseAddrs.RPC = "nope"
	_, err = a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "Failed to parse RPC") {
		t.Fatalf("expected rpc address error, got: %#v", err)
	}
	conf.AdvertiseAddrs.RPC = "127.0.0.1:4001"
	conf.AdvertiseAddrs.HTTP = "10.10.11.1:4005"

	// Parses the advertise addrs correctly
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err := a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	serfAddr := out.SerfConfig.MemberlistConfig.AdvertiseAddr
	if serfAddr != "127.0.0.1" {
		t.Fatalf("expect 127.0.0.1, got: %s", serfAddr)
	}
	serfPort := out.SerfConfig.MemberlistConfig.AdvertisePort
	if serfPort != 4000 {
		t.Fatalf("expected 4000, got: %d", serfPort)
	}

	// Assert addresses weren't changed
	if addr := conf.AdvertiseAddrs.RPC; addr != "127.0.0.1:4001" {
		t.Fatalf("bad rpc advertise addr: %#v", addr)
	}
	if addr := conf.AdvertiseAddrs.HTTP; addr != "10.10.11.1:4005" {
		t.Fatalf("expect 10.11.11.1:4005, got: %v", addr)
	}
	if addr := conf.Addresses.RPC; addr != "0.0.0.0:4647" {
		t.Fatalf("expect 0.0.0.0:4001, got: %v", addr)
	}

	// Sets up the ports properly
	conf.Addresses.RPC = ""
	conf.Addresses.Serf = ""
	conf.Ports.RPC = 4003
	conf.Ports.Serf = 4004

	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if addr := out.RPCAddr.Port; addr != 4003 {
		t.Fatalf("expect 4003, got: %d", out.RPCAddr.Port)
	}
	if port := out.SerfConfig.MemberlistConfig.BindPort; port != 4004 {
		t.Fatalf("expect 4004, got: %d", port)
	}

	// Prefers advertise over bind addr
	conf.BindAddr = "127.0.0.3"
	conf.Addresses.HTTP = "127.0.0.2"
	conf.Addresses.RPC = "127.0.0.2"
	conf.Addresses.Serf = "127.0.0.2"
	conf.AdvertiseAddrs.HTTP = "10.0.0.10"
	conf.AdvertiseAddrs.RPC = ""
	conf.AdvertiseAddrs.Serf = "10.0.0.12:4004"

	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	fmt.Println(conf.Addresses.RPC)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if addr := out.RPCAddr.IP.String(); addr != "127.0.0.2" {
		t.Fatalf("expect 127.0.0.2, got: %s", addr)
	}
	if port := out.RPCAddr.Port; port != 4003 {
		t.Fatalf("expect 4647, got: %d", port)
	}
	if addr := out.SerfConfig.MemberlistConfig.BindAddr; addr != "127.0.0.2" {
		t.Fatalf("expect 127.0.0.2, got: %s", addr)
	}
	if port := out.SerfConfig.MemberlistConfig.BindPort; port != 4004 {
		t.Fatalf("expect 4648, got: ^d", port)
	}
	if addr := conf.Addresses.HTTP; addr != "127.0.0.2:4646" {
		t.Fatalf("expect 127.0.0.2:4646, got: %s", addr)
	}
	if addr := conf.Addresses.RPC; addr != "127.0.0.2:4003" {
		t.Fatalf("expect 127.0.0.2:4003, got: %s", addr)
	}
	if addr := conf.Addresses.Serf; addr != "127.0.0.2:4004" {
		t.Fatalf("expect 10.0.0.12:4004, got: %s", addr)
	}
	if addr := conf.AdvertiseAddrs.HTTP; addr != "10.0.0.10:4646" {
		t.Fatalf("expect 10.0.0.10:4646, got: %s", addr)
	}
	if addr := conf.AdvertiseAddrs.RPC; addr != "127.0.0.3:4003" {
		t.Fatalf("expect 127.0.0.3:4003, got: %s", addr)
	}
	if addr := conf.AdvertiseAddrs.Serf; addr != "10.0.0.12:4004" {
		t.Fatalf("expect 10.0.0.12:4004, got: %s", addr)
	}

	conf.Server.NodeGCThreshold = "42g"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "unknown unit") {
		t.Fatalf("expected unknown unit error, got: %#v", err)
	}

	conf.Server.NodeGCThreshold = "10s"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if threshold := out.NodeGCThreshold; threshold != time.Second*10 {
		t.Fatalf("expect 10s, got: %s", threshold)
	}

	conf.Server.HeartbeatGrace = "42g"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "unknown unit") {
		t.Fatalf("expected unknown unit error, got: %#v", err)
	}

	conf.Server.HeartbeatGrace = "37s"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if threshold := out.HeartbeatGrace; threshold != time.Second*37 {
		t.Fatalf("expect 37s, got: %s", threshold)
	}

	// Defaults to the global bind addr
	conf.Addresses.RPC = ""
	conf.Addresses.Serf = ""
	conf.Addresses.HTTP = ""
	conf.AdvertiseAddrs.RPC = ""
	conf.AdvertiseAddrs.HTTP = ""
	conf.AdvertiseAddrs.Serf = ""
	conf.Ports.HTTP = 4646
	conf.Ports.RPC = 4647
	conf.Ports.Serf = 4648
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if addr := out.RPCAddr.IP.String(); addr != "127.0.0.3" {
		t.Fatalf("expect 127.0.0.3, got: %s", addr)
	}
	if addr := out.SerfConfig.MemberlistConfig.BindAddr; addr != "127.0.0.3" {
		t.Fatalf("expect 127.0.0.3, got: %s", addr)
	}
	if addr := conf.Addresses.HTTP; addr != "127.0.0.3:4646" {
		t.Fatalf("expect 127.0.0.3:4646, got: %s", addr)
	}
	if addr := conf.Addresses.RPC; addr != "127.0.0.3:4647" {
		t.Fatalf("expect 127.0.0.3:4647, got: %s", addr)
	}
	if addr := conf.Addresses.Serf; addr != "127.0.0.3:4648" {
		t.Fatalf("expect 127.0.0.3:4648, got: %s", addr)
	}

	// Properly handles the bootstrap flags
	conf.Server.BootstrapExpect = 1
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if !out.Bootstrap {
		t.Fatalf("should have set bootstrap mode")
	}
	if out.BootstrapExpect != 0 {
		t.Fatalf("boostrap expect should be 0")
	}

	conf.Server.BootstrapExpect = 3
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if out.Bootstrap {
		t.Fatalf("bootstrap mode should be disabled")
	}
	if out.BootstrapExpect != 3 {
		t.Fatalf("should have bootstrap-expect = 3")
	}
}

func TestAgent_ClientConfig(t *testing.T) {
	conf := DefaultConfig()
	a := &Agent{config: conf}
	conf.Client.Enabled = true
	conf.Addresses.HTTP = "127.0.0.1"
	conf.Ports.HTTP = 5678

	c, err := a.clientConfig()
	if err != nil {
		t.Fatalf("got err: %v", err)
	}

	expectedHttpAddr := "127.0.0.1:5678"
	if c.Node.HTTPAddr != expectedHttpAddr {
		t.Fatalf("Expected http addr: %v, got: %v", expectedHttpAddr, c.Node.HTTPAddr)
	}

	conf = DefaultConfig()
	a = &Agent{config: conf}
	conf.Client.Enabled = true
	conf.Addresses.HTTP = "127.0.0.1"

	c, err = a.clientConfig()
	if err != nil {
		t.Fatalf("got err: %v", err)
	}

	expectedHttpAddr = "127.0.0.1:4646"
	if c.Node.HTTPAddr != expectedHttpAddr {
		t.Fatalf("Expected http addr: %v, got: %v", expectedHttpAddr, c.Node.HTTPAddr)
	}
}
<commit_msg>Normalize configs for agent tests<commit_after>package agent

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/hashicorp/nomad/nomad"
	sconfig "github.com/hashicorp/nomad/nomad/structs/config"
)

// getTestLogger returns a log func appropriate for passing to
// Config.normalize()
func getTestLogger(t testing.TB) func(string) {
	return func(s string) {
		t.Log(s)
	}
}

func getPort() int {
	addr, err := net.ResolveTCPAddr("tcp", "localhost:0")
	if err != nil {
		panic(err)
	}

	l, err := net.ListenTCP("tcp", addr)
	if err != nil {
		panic(err)
	}
	defer l.Close()
	return l.Addr().(*net.TCPAddr).Port
}

func tmpDir(t testing.TB) string {
	dir, err := ioutil.TempDir("", "nomad")
	if err != nil {
		t.Fatalf("err: %v", err)
	}
	return dir
}

func makeAgent(t testing.TB, cb func(*Config)) (string, *Agent) {
	dir := tmpDir(t)
	conf := DevConfig()

	// Customize the server configuration
	config := nomad.DefaultConfig()
	conf.NomadConfig = config

	// Bind and set ports
	conf.BindAddr = "127.0.0.1"
	conf.Ports = &Ports{
		HTTP: getPort(),
		RPC:  getPort(),
		Serf: getPort(),
	}
	conf.NodeName = fmt.Sprintf("Node %d", conf.Ports.RPC)
	conf.Consul = sconfig.DefaultConsulConfig()
	conf.Vault.Enabled = new(bool)

	// Tighten the Serf timing
	config.SerfConfig.MemberlistConfig.SuspicionMult = 2
	config.SerfConfig.MemberlistConfig.RetransmitMult = 2
	config.SerfConfig.MemberlistConfig.ProbeTimeout = 50 * time.Millisecond
	config.SerfConfig.MemberlistConfig.ProbeInterval = 100 * time.Millisecond
	config.SerfConfig.MemberlistConfig.GossipInterval = 100 * time.Millisecond

	// Tighten the Raft timing
	config.RaftConfig.LeaderLeaseTimeout = 20 * time.Millisecond
	config.RaftConfig.HeartbeatTimeout = 40 * time.Millisecond
	config.RaftConfig.ElectionTimeout = 40 * time.Millisecond
	config.RaftConfig.StartAsLeader = true
	config.RaftTimeout = 500 * time.Millisecond

	if cb != nil {
		cb(conf)
	}

	if ok := conf.normalize(getTestLogger(t), config.DevMode); !ok {
		t.Fatalf("error normalizing config")
	}
	agent, err := NewAgent(conf, os.Stderr)
	if err != nil {
		os.RemoveAll(dir)
		t.Fatalf("err: %v", err)
	}
	return dir, agent
}

func TestAgent_RPCPing(t *testing.T) {
	dir, agent := makeAgent(t, nil)
	defer os.RemoveAll(dir)
	defer agent.Shutdown()

	var out struct{}
	if err := agent.RPC("Status.Ping", struct{}{}, &out); err != nil {
		t.Fatalf("err: %v", err)
	}
}

func TestAgent_ServerConfig(t *testing.T) {
	conf := DefaultConfig()
	a := &Agent{config: conf}
	testlogger := getTestLogger(t)
	dev := false

	// Returns error on bad serf addr
	conf.AdvertiseAddrs.Serf = "nope"
	_, err := a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "Failed to parse Serf") {
		t.Fatalf("expected serf address error, got: %#v", err)
	}
	conf.AdvertiseAddrs.Serf = "127.0.0.1:4000"

	// Returns error on bad rpc addr
	conf.AdvertiseAddrs.RPC = "nope"
	_, err = a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "Failed to parse RPC") {
		t.Fatalf("expected rpc address error, got: %#v", err)
	}
	conf.AdvertiseAddrs.RPC = "127.0.0.1:4001"
	conf.AdvertiseAddrs.HTTP = "10.10.11.1:4005"

	// Parses the advertise addrs correctly
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err := a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	serfAddr := out.SerfConfig.MemberlistConfig.AdvertiseAddr
	if serfAddr != "127.0.0.1" {
		t.Fatalf("expect 127.0.0.1, got: %s", serfAddr)
	}
	serfPort := out.SerfConfig.MemberlistConfig.AdvertisePort
	if serfPort != 4000 {
		t.Fatalf("expected 4000, got: %d", serfPort)
	}

	// Assert addresses weren't changed
	if addr := conf.AdvertiseAddrs.RPC; addr != "127.0.0.1:4001" {
		t.Fatalf("bad rpc advertise addr: %#v", addr)
	}
	if addr := conf.AdvertiseAddrs.HTTP; addr != "10.10.11.1:4005" {
		t.Fatalf("expect 10.11.11.1:4005, got: %v", addr)
	}
	if addr := conf.Addresses.RPC; addr != "0.0.0.0:4647" {
		t.Fatalf("expect 0.0.0.0:4001, got: %v", addr)
	}

	// Sets up the ports properly
	conf.Addresses.RPC = ""
	conf.Addresses.Serf = ""
	conf.Ports.RPC = 4003
	conf.Ports.Serf = 4004

	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if addr := out.RPCAddr.Port; addr != 4003 {
		t.Fatalf("expect 4003, got: %d", out.RPCAddr.Port)
	}
	if port := out.SerfConfig.MemberlistConfig.BindPort; port != 4004 {
		t.Fatalf("expect 4004, got: %d", port)
	}

	// Prefers advertise over bind addr
	conf.BindAddr = "127.0.0.3"
	conf.Addresses.HTTP = "127.0.0.2"
	conf.Addresses.RPC = "127.0.0.2"
	conf.Addresses.Serf = "127.0.0.2"
	conf.AdvertiseAddrs.HTTP = "10.0.0.10"
	conf.AdvertiseAddrs.RPC = ""
	conf.AdvertiseAddrs.Serf = "10.0.0.12:4004"

	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	fmt.Println(conf.Addresses.RPC)
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if addr := out.RPCAddr.IP.String(); addr != "127.0.0.2" {
		t.Fatalf("expect 127.0.0.2, got: %s", addr)
	}
	if port := out.RPCAddr.Port; port != 4003 {
		t.Fatalf("expect 4647, got: %d", port)
	}
	if addr := out.SerfConfig.MemberlistConfig.BindAddr; addr != "127.0.0.2" {
		t.Fatalf("expect 127.0.0.2, got: %s", addr)
	}
	if port := out.SerfConfig.MemberlistConfig.BindPort; port != 4004 {
		t.Fatalf("expect 4648, got: ^d", port)
	}
	if addr := conf.Addresses.HTTP; addr != "127.0.0.2:4646" {
		t.Fatalf("expect 127.0.0.2:4646, got: %s", addr)
	}
	if addr := conf.Addresses.RPC; addr != "127.0.0.2:4003" {
		t.Fatalf("expect 127.0.0.2:4003, got: %s", addr)
	}
	if addr := conf.Addresses.Serf; addr != "127.0.0.2:4004" {
		t.Fatalf("expect 10.0.0.12:4004, got: %s", addr)
	}
	if addr := conf.AdvertiseAddrs.HTTP; addr != "10.0.0.10:4646" {
		t.Fatalf("expect 10.0.0.10:4646, got: %s", addr)
	}
	if addr := conf.AdvertiseAddrs.RPC; addr != "127.0.0.3:4003" {
		t.Fatalf("expect 127.0.0.3:4003, got: %s", addr)
	}
	if addr := conf.AdvertiseAddrs.Serf; addr != "10.0.0.12:4004" {
		t.Fatalf("expect 10.0.0.12:4004, got: %s", addr)
	}

	conf.Server.NodeGCThreshold = "42g"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "unknown unit") {
		t.Fatalf("expected unknown unit error, got: %#v", err)
	}

	conf.Server.NodeGCThreshold = "10s"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if threshold := out.NodeGCThreshold; threshold != time.Second*10 {
		t.Fatalf("expect 10s, got: %s", threshold)
	}

	conf.Server.HeartbeatGrace = "42g"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err == nil || !strings.Contains(err.Error(), "unknown unit") {
		t.Fatalf("expected unknown unit error, got: %#v", err)
	}

	conf.Server.HeartbeatGrace = "37s"
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if threshold := out.HeartbeatGrace; threshold != time.Second*37 {
		t.Fatalf("expect 37s, got: %s", threshold)
	}

	// Defaults to the global bind addr
	conf.Addresses.RPC = ""
	conf.Addresses.Serf = ""
	conf.Addresses.HTTP = ""
	conf.AdvertiseAddrs.RPC = ""
	conf.AdvertiseAddrs.HTTP = ""
	conf.AdvertiseAddrs.Serf = ""
	conf.Ports.HTTP = 4646
	conf.Ports.RPC = 4647
	conf.Ports.Serf = 4648
	if ok := conf.normalize(testlogger, dev); !ok {
		t.Fatalf("failed to normalize config")
	}
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if addr := out.RPCAddr.IP.String(); addr != "127.0.0.3" {
		t.Fatalf("expect 127.0.0.3, got: %s", addr)
	}
	if addr := out.SerfConfig.MemberlistConfig.BindAddr; addr != "127.0.0.3" {
		t.Fatalf("expect 127.0.0.3, got: %s", addr)
	}
	if addr := conf.Addresses.HTTP; addr != "127.0.0.3:4646" {
		t.Fatalf("expect 127.0.0.3:4646, got: %s", addr)
	}
	if addr := conf.Addresses.RPC; addr != "127.0.0.3:4647" {
		t.Fatalf("expect 127.0.0.3:4647, got: %s", addr)
	}
	if addr := conf.Addresses.Serf; addr != "127.0.0.3:4648" {
		t.Fatalf("expect 127.0.0.3:4648, got: %s", addr)
	}

	// Properly handles the bootstrap flags
	conf.Server.BootstrapExpect = 1
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if !out.Bootstrap {
		t.Fatalf("should have set bootstrap mode")
	}
	if out.BootstrapExpect != 0 {
		t.Fatalf("boostrap expect should be 0")
	}

	conf.Server.BootstrapExpect = 3
	out, err = a.serverConfig()
	if err != nil {
		t.Fatalf("err: %s", err)
	}
	if out.Bootstrap {
		t.Fatalf("bootstrap mode should be disabled")
	}
	if out.BootstrapExpect != 3 {
		t.Fatalf("should have bootstrap-expect = 3")
	}
}

func TestAgent_ClientConfig(t *testing.T) {
	conf := DefaultConfig()
	a := &Agent{config: conf}
	conf.Client.Enabled = true
	conf.Addresses.HTTP = "127.0.0.1"
	conf.Ports.HTTP = 5678

	c, err := a.clientConfig()
	if err != nil {
		t.Fatalf("got err: %v", err)
	}

	expectedHttpAddr := "127.0.0.1:5678"
	if c.Node.HTTPAddr != expectedHttpAddr {
		t.Fatalf("Expected http addr: %v, got: %v", expectedHttpAddr, c.Node.HTTPAddr)
	}

	conf = DefaultConfig()
	a = &Agent{config: conf}
	conf.Client.Enabled = true
	conf.Addresses.HTTP = "127.0.0.1"

	c, err = a.clientConfig()
	if err != nil {
		t.Fatalf("got err: %v", err)
	}

	expectedHttpAddr = "127.0.0.1:4646"
	if c.Node.HTTPAddr != expectedHttpAddr {
		t.Fatalf("Expected http addr: %v, got: %v", expectedHttpAddr, c.Node.HTTPAddr)
	}
}
<|endoftext|>


<commit_before>package dnssd

/*
#include <stdlib.h>
#include "dns_sd.h"
DNSServiceErrorType Browse(DNSServiceRef *, DNSServiceFlags, const char *, void *);
DNSServiceErrorType Resolve(DNSServiceRef *, DNSServiceFlags, uint32_t, const char *, const char *, const char *, void *);
DNSServiceErrorType QueryRecord(DNSServiceRef *, DNSServiceFlags, uint32_t, const char *, uint16_t, uint16_t, void *);
DNSServiceErrorType ServiceRegister(DNSServiceRef *, DNSServiceFlags, uint32_t, const char *, const char *, const char *, const char *, uint16_t, uint16_t, const void *, void*);
DNSServiceErrorType GetAddrInfo(DNSServiceRef *, DNSServiceFlags, uint32_t, DNSServiceProtocol, const char *, void *);
*/
import "C"

import (
	"bytes"
	"encoding/binary"
	"reflect"
	"strings"
	"unsafe"
)

type Context struct {
	ref *C.DNSServiceRef
	c   reflect.Value
}

type BrowseReply struct {
	InterfaceIndex uint32
	ServiceName    string
	RegType        string
	ReplyDomain    string
}

type ResolveReply struct {
	InterfaceIndex uint32
	FullName       string
	HostTarget     string
	Port           uint16
	TxtRecordMap   map[string]string
}

type QueryRecordReply struct {
	InterfaceIndex uint32
	FullName       string
	Rrtype         uint16
	Rrclass        uint16
	Rdlen          uint16
	Rdata          []byte
	Ttl            uint32
}

type RegisterReply struct {
	Name    string
	RegType string
	Domain  string
}

type RecordSRV struct {
	Priority uint16
	Weight   uint16
	Port     uint16
	Host     string
}

type GetAddrInfoReply struct {
	InterfaceIndex uint32
	HostName       string
	Ip             string
	Ttl            uint32
}

func init() {
}

func Process(ctx *Context) {
	C.DNSServiceProcessResult(*(ctx.ref))
	ctx.c.Close()
}

func (ctx *Context) Release() {
	C.DNSServiceRefDeallocate(*(ctx.ref))
}

//export goBrowseReply
func goBrowseReply(interfaceIndex uint32, serviceName *C.char, regType *C.char, replyDomain *C.char, contextpt unsafe.Pointer) {
	br := &BrowseReply{interfaceIndex, C.GoString(serviceName), C.GoString(regType), C.GoString(replyDomain)}
	fn := (*(*func(*BrowseReply))(unsafe.Pointer(&contextpt)))
	fn(br)
}

func Browse(flags C.DNSServiceFlags, regType string, c chan *BrowseReply) (*Context, error) {

	var fn = func(browseReply *BrowseReply) {
		c <- browseReply
	}
	var fnptr = unsafe.Pointer(&fn)

	cregType := C.CString(regType)
	defer C.free(unsafe.Pointer(cregType))

	var ref C.DNSServiceRef
	cerr := C.Browse(&ref, flags, cregType, (*(*unsafe.Pointer)(fnptr)))

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c)}, nil
	}

	return nil, createErr(cerr)
}

//export goResolveReply
func goResolveReply(interfaceIndex uint32, fullName *C.char, hostTarget *C.char, port uint16, txtLen uint16, txtRecords unsafe.Pointer, contextpt unsafe.Pointer) {

	txtRecordMap := make(map[string]string)
	count := int(C.TXTRecordGetCount((C.uint16_t)(txtLen), txtRecords))

	for i := 0; i < count; i++ {
		key := make([]byte, txtLen)
		var valueLen C.uint8_t
		var value unsafe.Pointer

		// [cgo] How to convert Go type []byte to C type *_Ctype_char?
		// https://groups.google.com/forum/#!topic/golang-nuts/f1sWP3-YbuI
		C.TXTRecordGetItemAtIndex(
			(C.uint16_t)(txtLen),
			txtRecords,
			(C.uint16_t)(i),
			256,
			(*C.char)(unsafe.Pointer(&key[0])),
			&valueLen,
			&value,
		)

		byteValue := C.GoBytes(value, C.int(valueLen))
		// How Can I Remove the ZEROS in a String?
		// https://groups.google.com/forum/#!topic/golang-nuts/JuLPI4nCsOo
		keyString := string(bytes.Trim(key, string(0)))
		txtRecordMap[keyString] = string(byteValue)
	}

	rr := &ResolveReply{interfaceIndex, C.GoString(fullName), C.GoString(hostTarget), port, txtRecordMap}
	fn := (*(*func(*ResolveReply))(unsafe.Pointer(&contextpt)))
	fn(rr)
}

func Resolve(flags C.DNSServiceFlags, interfaceIndex uint32, serviceName string, regType string, replyDomain string, c chan *ResolveReply) (*Context, error) {
	var fn = func(resolveReply *ResolveReply) {
		c <- resolveReply
	}
	var fnptr = unsafe.Pointer(&fn)

	var ref C.DNSServiceRef
	cserviceName := C.CString(serviceName)
	cregType := C.CString(regType)
	creplyDomain := C.CString(replyDomain)

	defer C.free(unsafe.Pointer(cserviceName))
	defer C.free(unsafe.Pointer(cregType))
	defer C.free(unsafe.Pointer(creplyDomain))

	cerr := C.Resolve(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		cserviceName,
		cregType,
		creplyDomain,
		(*(*unsafe.Pointer)(fnptr)),
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c)}, nil
	}

	return nil, createErr(cerr)
}

//export goQueryRecordReply
func goQueryRecordReply(interfaceIndex uint32, fullName *C.char, rrtype uint16, rrclass uint16, rdlen uint16, rdata unsafe.Pointer, ttl uint32, contextpt unsafe.Pointer) {
	qrr := &QueryRecordReply{interfaceIndex, C.GoString(fullName), rrtype, rrclass, rdlen, C.GoBytes(rdata, C.int(rdlen)), ttl}
	fn := (*(*func(*QueryRecordReply))(unsafe.Pointer(&contextpt)))
	fn(qrr)
}

//export goServiceRegisterReply
func goServiceRegisterReply(name *C.char, regType *C.char, domain *C.char, contextpt unsafe.Pointer) {
	srr := &RegisterReply{C.GoString(name), C.GoString(regType), C.GoString(domain)}
	fn := (*(*func(*RegisterReply))(unsafe.Pointer(&contextpt)))
	fn(srr)
}

func QueryRecord(flags C.DNSServiceFlags, interfaceIndex uint32, fullName string, rrtype C.uint16_t, rrclass C.uint16_t, c chan *QueryRecordReply) (*Context, error) {

	var fn = func(queryRecordReply *QueryRecordReply) {
		c <- queryRecordReply
	}
	var fnptr = unsafe.Pointer(&fn)

	var ref C.DNSServiceRef
	cfullName := C.CString(fullName)

	defer C.free(unsafe.Pointer(cfullName))

	cerr := C.QueryRecord(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		cfullName,
		rrtype,
		rrclass,
		(*(*unsafe.Pointer)(fnptr)),
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c)}, nil
	}

	return nil, createErr(cerr)
}

func ServiceRegister(flags C.DNSServiceFlags, interfaceIndex uint32, name string, regType string, domain string, host string, port uint16, txtRecords map[string]string, c chan *RegisterReply) (*Context, error) {

	var fn = func(registerReply *RegisterReply) {
		c <- registerReply
	}
	var fnptr = unsafe.Pointer(&fn)

	var ref C.DNSServiceRef
	c_name := C.CString(name)
	c_regType := C.CString(regType)
	c_domain := C.CString(domain)
	c_host := C.CString(host)

	defer C.free(unsafe.Pointer(c_name))
	defer C.free(unsafe.Pointer(c_regType))
	defer C.free(unsafe.Pointer(c_domain))
	defer C.free(unsafe.Pointer(c_host))

	// Setup text record
	var txtRecordRef C.TXTRecordRef
	C.TXTRecordCreate(
		&txtRecordRef,
		0, // let the c library manage it's own buffer
		nil,
	)

	// Add text records
	for key, value := range txtRecords {
		c_key := C.CString(key)
		c_value := C.CString(value)

		defer C.free(unsafe.Pointer(c_key))
		defer C.free(unsafe.Pointer(c_value))

		txtadderr := C.TXTRecordSetValue(
			&txtRecordRef,
			c_key,
			(C.uint8_t)(len(value)),
			unsafe.Pointer(c_value),
		)

		if txtadderr != DNSServiceErr_NoError {
			return nil, createErr(txtadderr)
		}
	}

	cerr := C.ServiceRegister(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		c_name,
		c_regType,
		c_domain,
		c_host,
		(C.uint16_t)(port),
		C.TXTRecordGetLength(&txtRecordRef),
		C.TXTRecordGetBytesPtr(&txtRecordRef),
		(*(*unsafe.Pointer)(fnptr)),
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c)}, nil
	}

	return nil, createErr(cerr)
}

//export goGetAddrInfoReply
func goGetAddrInfoReply(interfaceIndex uint32, hostName *C.char, ip *C.char, ttl uint32, contextpt unsafe.Pointer) {
	gair := &GetAddrInfoReply{interfaceIndex, C.GoString(hostName), C.GoString(ip), ttl}
	fn := (*(*func(*GetAddrInfoReply))(unsafe.Pointer(&contextpt)))
	fn(gair)
}

func GetAddrInfo(flags C.DNSServiceFlags, interfaceIndex uint32, protocol C.DNSServiceProtocol, hostName string, c chan *GetAddrInfoReply) (*Context, error) {
	var fn = func(getAddrInfoReply *GetAddrInfoReply) {
		c <- getAddrInfoReply
	}
	var fnptr = unsafe.Pointer(&fn)

	var ref C.DNSServiceRef
	chostName := C.CString(hostName)

	defer C.free(unsafe.Pointer(chostName))

	cerr := C.GetAddrInfo(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		protocol,
		chostName,
		(*(*unsafe.Pointer)(fnptr)),
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c)}, nil
	}

	return nil, createErr(cerr)
}

func (qrr *QueryRecordReply) SRV() *RecordSRV {
	reader := bytes.NewReader(qrr.Rdata)
	var priority, weight, port uint16
	var host string

	binary.Read(reader, binary.BigEndian, &priority)
	binary.Read(reader, binary.BigEndian, &weight)
	binary.Read(reader, binary.BigEndian, &port)

	hosts := make([]string, 0)

	for reader.Len() > 0 {
		var c byte
		binary.Read(reader, binary.BigEndian, &c)
		line := make([]byte, c)
		binary.Read(reader, binary.BigEndian, line)
		hosts = append(hosts, string(line))
	}

	host = strings.Join(hosts, ".")

	return &RecordSRV{priority, weight, port, host}
}
<commit_msg>store callback and pass int to c world<commit_after>package dnssd

/*
#include <stdlib.h>
#include "dns_sd.h"
DNSServiceErrorType Browse(DNSServiceRef *, DNSServiceFlags, const char *, void *);
DNSServiceErrorType Resolve(DNSServiceRef *, DNSServiceFlags, uint32_t, const char *, const char *, const char *, void *);
DNSServiceErrorType QueryRecord(DNSServiceRef *, DNSServiceFlags, uint32_t, const char *, uint16_t, uint16_t, void *);
DNSServiceErrorType ServiceRegister(DNSServiceRef *, DNSServiceFlags, uint32_t, const char *, const char *, const char *, const char *, uint16_t, uint16_t, const void *, void*);
DNSServiceErrorType GetAddrInfo(DNSServiceRef *, DNSServiceFlags, uint32_t, DNSServiceProtocol, const char *, void *);
*/
import "C"

import (
	"bytes"
	"encoding/binary"
	"reflect"
	"strings"
	"sync"
	"unsafe"
)

var (
	callbacks      = map[uintptr]interface{}{}
	mx             sync.Mutex
	callbackOffset uintptr = 0
)

func storeCallback(cb interface{}) uintptr {
	mx.Lock()
	defer mx.Unlock()
	callbackOffset += 1
	callbacks[callbackOffset] = cb
	return callbackOffset
}

func restoreCallback(i uintptr) interface{} {
	mx.Lock()
	defer mx.Unlock()
	return callbacks[i]
}

func deleteCallback(i uintptr) {
	mx.Lock()
	defer mx.Unlock()
	delete(callbacks, i)
}

type Context struct {
	ref           *C.DNSServiceRef
	c             reflect.Value
	callbackIndex uintptr
}

type BrowseReply struct {
	InterfaceIndex uint32
	ServiceName    string
	RegType        string
	ReplyDomain    string
}

type ResolveReply struct {
	InterfaceIndex uint32
	FullName       string
	HostTarget     string
	Port           uint16
	TxtRecordMap   map[string]string
}

type QueryRecordReply struct {
	InterfaceIndex uint32
	FullName       string
	Rrtype         uint16
	Rrclass        uint16
	Rdlen          uint16
	Rdata          []byte
	Ttl            uint32
}

type RegisterReply struct {
	Name    string
	RegType string
	Domain  string
}

type RecordSRV struct {
	Priority uint16
	Weight   uint16
	Port     uint16
	Host     string
}

type GetAddrInfoReply struct {
	InterfaceIndex uint32
	HostName       string
	Ip             string
	Ttl            uint32
}

func init() {
}

func Process(ctx *Context) {
	C.DNSServiceProcessResult(*(ctx.ref))
	ctx.c.Close()
}

func (ctx *Context) Release() {
	C.DNSServiceRefDeallocate(*(ctx.ref))
	deleteCallback(ctx.callbackIndex)
}

//export goBrowseReply
func goBrowseReply(interfaceIndex uint32, serviceName *C.char, regType *C.char, replyDomain *C.char, contextpt unsafe.Pointer) {
	br := &BrowseReply{interfaceIndex, C.GoString(serviceName), C.GoString(regType), C.GoString(replyDomain)}
	fn := restoreCallback(uintptr(contextpt)).(func(*BrowseReply))
	fn(br)
}

func Browse(flags C.DNSServiceFlags, regType string, c chan *BrowseReply) (*Context, error) {

	var fn = func(browseReply *BrowseReply) {
		c <- browseReply
	}
	callbackIndex := storeCallback(fn)
	var fnptr = unsafe.Pointer(callbackIndex)

	cregType := C.CString(regType)
	defer C.free(unsafe.Pointer(cregType))

	var ref C.DNSServiceRef
	cerr := C.Browse(&ref, flags, cregType, fnptr)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c), callbackIndex}, nil
	}

	return nil, createErr(cerr)
}

//export goResolveReply
func goResolveReply(interfaceIndex uint32, fullName *C.char, hostTarget *C.char, port uint16, txtLen uint16, txtRecords unsafe.Pointer, contextpt unsafe.Pointer) {

	txtRecordMap := make(map[string]string)
	count := int(C.TXTRecordGetCount((C.uint16_t)(txtLen), txtRecords))

	for i := 0; i < count; i++ {
		key := make([]byte, txtLen)
		var valueLen C.uint8_t
		var value unsafe.Pointer

		// [cgo] How to convert Go type []byte to C type *_Ctype_char?
		// https://groups.google.com/forum/#!topic/golang-nuts/f1sWP3-YbuI
		C.TXTRecordGetItemAtIndex(
			(C.uint16_t)(txtLen),
			txtRecords,
			(C.uint16_t)(i),
			256,
			(*C.char)(unsafe.Pointer(&key[0])),
			&valueLen,
			&value,
		)

		byteValue := C.GoBytes(value, C.int(valueLen))
		// How Can I Remove the ZEROS in a String?
		// https://groups.google.com/forum/#!topic/golang-nuts/JuLPI4nCsOo
		keyString := string(bytes.Trim(key, string(0)))
		txtRecordMap[keyString] = string(byteValue)
	}

	rr := &ResolveReply{interfaceIndex, C.GoString(fullName), C.GoString(hostTarget), port, txtRecordMap}
	fn := restoreCallback(uintptr(contextpt)).(func(*ResolveReply))
	fn(rr)
}

func Resolve(flags C.DNSServiceFlags, interfaceIndex uint32, serviceName string, regType string, replyDomain string, c chan *ResolveReply) (*Context, error) {
	var fn = func(resolveReply *ResolveReply) {
		c <- resolveReply
	}
	callbackIndex := storeCallback(fn)
	var fnptr = unsafe.Pointer(callbackIndex)

	var ref C.DNSServiceRef
	cserviceName := C.CString(serviceName)
	cregType := C.CString(regType)
	creplyDomain := C.CString(replyDomain)

	defer C.free(unsafe.Pointer(cserviceName))
	defer C.free(unsafe.Pointer(cregType))
	defer C.free(unsafe.Pointer(creplyDomain))

	cerr := C.Resolve(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		cserviceName,
		cregType,
		creplyDomain,
		fnptr,
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c), callbackIndex}, nil
	}

	return nil, createErr(cerr)
}

//export goQueryRecordReply
func goQueryRecordReply(interfaceIndex uint32, fullName *C.char, rrtype uint16, rrclass uint16, rdlen uint16, rdata unsafe.Pointer, ttl uint32, contextpt unsafe.Pointer) {
	qrr := &QueryRecordReply{interfaceIndex, C.GoString(fullName), rrtype, rrclass, rdlen, C.GoBytes(rdata, C.int(rdlen)), ttl}
	fn := restoreCallback(uintptr(contextpt)).(func(*QueryRecordReply))
	fn(qrr)
}

//export goServiceRegisterReply
func goServiceRegisterReply(name *C.char, regType *C.char, domain *C.char, contextpt unsafe.Pointer) {
	srr := &RegisterReply{C.GoString(name), C.GoString(regType), C.GoString(domain)}
	fn := restoreCallback(uintptr(contextpt)).(func(*RegisterReply))
	fn(srr)
}

func QueryRecord(flags C.DNSServiceFlags, interfaceIndex uint32, fullName string, rrtype C.uint16_t, rrclass C.uint16_t, c chan *QueryRecordReply) (*Context, error) {

	var fn = func(queryRecordReply *QueryRecordReply) {
		c <- queryRecordReply
	}
	callbackIndex := storeCallback(fn)
	var fnptr = unsafe.Pointer(callbackIndex)

	var ref C.DNSServiceRef
	cfullName := C.CString(fullName)

	defer C.free(unsafe.Pointer(cfullName))

	cerr := C.QueryRecord(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		cfullName,
		rrtype,
		rrclass,
		fnptr,
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c), callbackIndex}, nil
	}

	return nil, createErr(cerr)
}

func ServiceRegister(flags C.DNSServiceFlags, interfaceIndex uint32, name string, regType string, domain string, host string, port uint16, txtRecords map[string]string, c chan *RegisterReply) (*Context, error) {

	var fn = func(registerReply *RegisterReply) {
		c <- registerReply
	}
	callbackIndex := storeCallback(fn)
	var fnptr = unsafe.Pointer(callbackIndex)

	var ref C.DNSServiceRef
	c_name := C.CString(name)
	c_regType := C.CString(regType)
	c_domain := C.CString(domain)
	c_host := C.CString(host)

	defer C.free(unsafe.Pointer(c_name))
	defer C.free(unsafe.Pointer(c_regType))
	defer C.free(unsafe.Pointer(c_domain))
	defer C.free(unsafe.Pointer(c_host))

	// Setup text record
	var txtRecordRef C.TXTRecordRef
	C.TXTRecordCreate(
		&txtRecordRef,
		0, // let the c library manage it's own buffer
		nil,
	)

	// Add text records
	for key, value := range txtRecords {
		c_key := C.CString(key)
		c_value := C.CString(value)

		defer C.free(unsafe.Pointer(c_key))
		defer C.free(unsafe.Pointer(c_value))

		txtadderr := C.TXTRecordSetValue(
			&txtRecordRef,
			c_key,
			(C.uint8_t)(len(value)),
			unsafe.Pointer(c_value),
		)

		if txtadderr != DNSServiceErr_NoError {
			return nil, createErr(txtadderr)
		}
	}

	cerr := C.ServiceRegister(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		c_name,
		c_regType,
		c_domain,
		c_host,
		(C.uint16_t)(port),
		C.TXTRecordGetLength(&txtRecordRef),
		C.TXTRecordGetBytesPtr(&txtRecordRef),
		fnptr,
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c), callbackIndex}, nil
	}

	return nil, createErr(cerr)
}

//export goGetAddrInfoReply
func goGetAddrInfoReply(interfaceIndex uint32, hostName *C.char, ip *C.char, ttl uint32, contextpt unsafe.Pointer) {
	gair := &GetAddrInfoReply{interfaceIndex, C.GoString(hostName), C.GoString(ip), ttl}
	fn := restoreCallback(uintptr(contextpt)).(func(*GetAddrInfoReply))
	fn(gair)
}

func GetAddrInfo(flags C.DNSServiceFlags, interfaceIndex uint32, protocol C.DNSServiceProtocol, hostName string, c chan *GetAddrInfoReply) (*Context, error) {
	var fn = func(getAddrInfoReply *GetAddrInfoReply) {
		c <- getAddrInfoReply
	}
	callbackIndex := storeCallback(fn)
	var fnptr = unsafe.Pointer(callbackIndex)

	var ref C.DNSServiceRef
	chostName := C.CString(hostName)

	defer C.free(unsafe.Pointer(chostName))

	cerr := C.GetAddrInfo(
		&ref,
		flags,
		(C.uint32_t)(interfaceIndex),
		protocol,
		chostName,
		fnptr,
	)

	if cerr == DNSServiceErr_NoError {
		return &Context{&ref, reflect.ValueOf(c), callbackIndex}, nil
	}

	return nil, createErr(cerr)
}

func (qrr *QueryRecordReply) SRV() *RecordSRV {
	reader := bytes.NewReader(qrr.Rdata)
	var priority, weight, port uint16
	var host string

	binary.Read(reader, binary.BigEndian, &priority)
	binary.Read(reader, binary.BigEndian, &weight)
	binary.Read(reader, binary.BigEndian, &port)

	hosts := make([]string, 0)

	for reader.Len() > 0 {
		var c byte
		binary.Read(reader, binary.BigEndian, &c)
		line := make([]byte, c)
		binary.Read(reader, binary.BigEndian, line)
		hosts = append(hosts, string(line))
	}

	host = strings.Join(hosts, ".")

	return &RecordSRV{priority, weight, port, host}
}
<|endoftext|>


<commit_before>package genetics

import (
	"testing"
	"github.com/yaricom/goNEAT/neat"
)

func buildSpeciesWithOrganisms(id int) *Species {
	gen := &Genome{
		GenomeId:1,
	}

	sp := NewSpecies(id)
	sp.addOrganism(NewOrganism(5.0 * float64(id), gen, id))
	sp.addOrganism(NewOrganism(15.0 * float64(id), gen, id))
	sp.addOrganism(NewOrganism(10.0 * float64(id), gen, id))

	return sp
}

// Tests Species ReadGene
func TestGene_adjustFitness(t *testing.T)  {
	sp := buildSpeciesWithOrganisms(1)

	// Configuration
	conf := neat.Neat{
		DropOffAge:5,
		SurvivalThresh:0.5,
		AgeSignificance:0.5,
	}
	sp.adjustFitness(&conf)

	// test results
	if sp.Organisms[0].IsChampion != true {
		t.Error("sp.Organisms[0].IsChampion", true, sp.Organisms[0].IsChampion)
	}
	if sp.AgeOfLastImprovement != 1 {
		t.Error("sp.AgeOfLastImprovement", 1, sp.AgeOfLastImprovement)
	}
	if sp.MaxFitnessEver != 15.0 {
		t.Error("sp.MaxFitnessEver", 15.0, sp.MaxFitnessEver)
	}
	if sp.Organisms[2].ToEliminate != true {
		t.Error("sp.Organisms[2].ToEliminate", true, sp.Organisms[2].ToEliminate)
	}
}

// Tests Species countOffspring
func TestSpecies_countOffspring(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)
	for i, o := range sp.Organisms {
		o.ExpectedOffspring = float64(i) * 1.5
	}

	skim := sp.countOffspring(0.5)

	if sp.ExpectedOffspring != 5 {
		t.Error("sp.ExpectedOffspring", 5, sp.ExpectedOffspring)
	}
	if skim != 0 {
		t.Error("skim", 0, skim)
	}

	sp = buildSpeciesWithOrganisms(2)
	for i, o := range sp.Organisms {
		o.ExpectedOffspring = float64(i) * 1.5
	}
	skim = sp.countOffspring(0.4)
	if sp.ExpectedOffspring != 4 {
		t.Error("sp.ExpectedOffspring", 5, sp.ExpectedOffspring)
	}
	if skim != 0.9 {
		t.Error("skim", 0.9, skim)
	}
}

// Tests Species computeAvgFitness
func TestSpecies_computeAvgFitness(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	avg_check := 0.0
	for _, o := range sp.Organisms{
		avg_check += o.Fitness
	}
	avg_check /= float64(len(sp.Organisms))

	sp.computeAvgFitness()
	if sp.AvgFitness != avg_check {
		t.Error("sp.AvgFitness != avg_check", sp.AvgFitness, avg_check)
	}
}

func TestSpecies_computeMaxFitness(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	sp.computeMaxFitness()
	if sp.MaxFitness != 15.0 {
		t.Error("sp.MaxFitness != 15.0", 15.0, sp.MaxFitness)
	}
}

func TestSpecies_findChampion(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	champ := sp.findChampion()
	if champ.Fitness != 15.0 {
		t.Error("champ.Fitness != 15.0", champ.Fitness)
	}

}

func TestSpecies_removeOrganism(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	// test remove
	size := len(sp.Organisms)
	res, err := sp.removeOrganism(sp.Organisms[0])
	if res != true {
		t.Error("res != true", res, err)
	}
	if err != nil {
		t.Error("err != nil", err)
	}
	if size - 1 != len(sp.Organisms) {
		t.Error("size - 1 != len(sp.Organisms)", size - 1, len(sp.Organisms))
	}

	// test fail to remove
	size = len(sp.Organisms)
	gen := &Genome{
		GenomeId:1,
	}
	org := NewOrganism(6.0, gen, 1)
	res, err = sp.removeOrganism(org)
	if res == true {
		t.Error("res == true", res, err)
	}
	if err == nil {
		t.Error("err == nil", res, err)
	}
	if size != len(sp.Organisms) {
		t.Error("size != len(sp.Organisms)", size, len(sp.Organisms))
	}
}

// Tests Species reproduce failure
func TestSpecies_reproduce_fail(t *testing.T) {
	sp := NewSpecies(1)

	sp.ExpectedOffspring = 1

	res, err := sp.reproduce(1, nil, nil, nil)
	if res != false {
		t.Error("res != false")
	}
	if err == nil {
		t.Error("err == nil")
	}
}

// Tests Species reproduce success
func TestSpecies_reproduce(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)
	sp.ExpectedOffspring = 2

	sorted_species := []*Species {
		buildSpeciesWithOrganisms(4),
		buildSpeciesWithOrganisms(3),
		buildSpeciesWithOrganisms(2),
	}

	// Configuration
	conf := neat.Neat{
		DropOffAge:5,
		SurvivalThresh:0.5,
		AgeSignificance:0.5,
		PopSize:3,
	}

	res, err := sp.reproduce(1, nil, sorted_species, &conf)
	if !res {
		t.Error("!res", err)
	}
	if err != nil {
		t.Error("err != nil", err)
	}
	if len(sp.Organisms) != 4 {
		t.Error("No new baby was created")
	}
}
<commit_msg>Commented out not yet complete test<commit_after>package genetics

import (
	"testing"
	"github.com/yaricom/goNEAT/neat"
)

func buildSpeciesWithOrganisms(id int) *Species {
	gen := &Genome{
		Id:1,
	}

	sp := NewSpecies(id)
	sp.addOrganism(NewOrganism(5.0 * float64(id), gen, id))
	sp.addOrganism(NewOrganism(15.0 * float64(id), gen, id))
	sp.addOrganism(NewOrganism(10.0 * float64(id), gen, id))

	return sp
}

// Tests Species ReadGene
func TestGene_adjustFitness(t *testing.T)  {
	sp := buildSpeciesWithOrganisms(1)

	// Configuration
	conf := neat.Neat{
		DropOffAge:5,
		SurvivalThresh:0.5,
		AgeSignificance:0.5,
	}
	sp.adjustFitness(&conf)

	// test results
	if sp.Organisms[0].IsChampion != true {
		t.Error("sp.Organisms[0].IsChampion", true, sp.Organisms[0].IsChampion)
	}
	if sp.AgeOfLastImprovement != 1 {
		t.Error("sp.AgeOfLastImprovement", 1, sp.AgeOfLastImprovement)
	}
	if sp.MaxFitnessEver != 15.0 {
		t.Error("sp.MaxFitnessEver", 15.0, sp.MaxFitnessEver)
	}
	if sp.Organisms[2].ToEliminate != true {
		t.Error("sp.Organisms[2].ToEliminate", true, sp.Organisms[2].ToEliminate)
	}
}

// Tests Species countOffspring
func TestSpecies_countOffspring(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)
	for i, o := range sp.Organisms {
		o.ExpectedOffspring = float64(i) * 1.5
	}

	skim := sp.countOffspring(0.5)

	if sp.ExpectedOffspring != 5 {
		t.Error("sp.ExpectedOffspring", 5, sp.ExpectedOffspring)
	}
	if skim != 0 {
		t.Error("skim", 0, skim)
	}

	sp = buildSpeciesWithOrganisms(2)
	for i, o := range sp.Organisms {
		o.ExpectedOffspring = float64(i) * 1.5
	}
	skim = sp.countOffspring(0.4)
	if sp.ExpectedOffspring != 4 {
		t.Error("sp.ExpectedOffspring", 5, sp.ExpectedOffspring)
	}
	if skim != 0.9 {
		t.Error("skim", 0.9, skim)
	}
}

// Tests Species computeAvgFitness
func TestSpecies_computeAvgFitness(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	avg_check := 0.0
	for _, o := range sp.Organisms{
		avg_check += o.Fitness
	}
	avg_check /= float64(len(sp.Organisms))

	sp.computeAvgFitness()
	if sp.AvgFitness != avg_check {
		t.Error("sp.AvgFitness != avg_check", sp.AvgFitness, avg_check)
	}
}

func TestSpecies_computeMaxFitness(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	sp.computeMaxFitness()
	if sp.MaxFitness != 15.0 {
		t.Error("sp.MaxFitness != 15.0", 15.0, sp.MaxFitness)
	}
}

func TestSpecies_findChampion(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	champ := sp.findChampion()
	if champ.Fitness != 15.0 {
		t.Error("champ.Fitness != 15.0", champ.Fitness)
	}

}

func TestSpecies_removeOrganism(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)

	// test remove
	size := len(sp.Organisms)
	res, err := sp.removeOrganism(sp.Organisms[0])
	if res != true {
		t.Error("res != true", res, err)
	}
	if err != nil {
		t.Error("err != nil", err)
	}
	if size - 1 != len(sp.Organisms) {
		t.Error("size - 1 != len(sp.Organisms)", size - 1, len(sp.Organisms))
	}

	// test fail to remove
	size = len(sp.Organisms)
	gen := &Genome{
		Id:1,
	}
	org := NewOrganism(6.0, gen, 1)
	res, err = sp.removeOrganism(org)
	if res == true {
		t.Error("res == true", res, err)
	}
	if err == nil {
		t.Error("err == nil", res, err)
	}
	if size != len(sp.Organisms) {
		t.Error("size != len(sp.Organisms)", size, len(sp.Organisms))
	}
}

// Tests Species reproduce failure
func TestSpecies_reproduce_fail(t *testing.T) {
	sp := NewSpecies(1)

	sp.ExpectedOffspring = 1

	res, err := sp.reproduce(1, nil, nil, nil)
	if res != false {
		t.Error("res != false")
	}
	if err == nil {
		t.Error("err == nil")
	}
}

// Tests Species reproduce success
// TODO do this
/*
func TestSpecies_reproduce(t *testing.T) {
	sp := buildSpeciesWithOrganisms(1)
	sp.ExpectedOffspring = 2

	sorted_species := []*Species {
		buildSpeciesWithOrganisms(4),
		buildSpeciesWithOrganisms(3),
		buildSpeciesWithOrganisms(2),
	}

	// Configuration
	conf := neat.Neat{
		DropOffAge:5,
		SurvivalThresh:0.5,
		AgeSignificance:0.5,
		PopSize:3,
	}

	res, err := sp.reproduce(1, nil, sorted_species, &conf)
	if !res {
		t.Error("!res", err)
	}
	if err != nil {
		t.Error("err != nil", err)
	}
	if len(sp.Organisms) != 4 {
		t.Error("No new baby was created")
	}
}*/
<|endoftext|>


<commit_before>package commands

import (
	"bufio"
	"io/ioutil"
	"os"
	"strings"

	"github.com/github/git-lfs/lfs"
	"github.com/github/git-lfs/vendor/_nuts/github.com/spf13/cobra"
)

var (
	untrackCmd = &cobra.Command{
		Use:   "untrack",
		Short: "Remove an entry from .gitattributes",
		Run:   untrackCommand,
	}
)

// untrackCommand takes a list of paths as an argument, and removes each path from the
// default attribtues file (.gitattributes), if it exists.
func untrackCommand(cmd *cobra.Command, args []string) {
	if lfs.LocalGitDir == "" {
		Print("Not a git repository.")
		os.Exit(128)
	}
	if lfs.LocalWorkingDir == "" {
		Print("This operation must be run in a work tree.")
		os.Exit(128)
	}

	lfs.InstallHooks(false)

	if len(args) < 1 {
		Print("git lfs untrack <path> [path]*")
		return
	}

	data, err := ioutil.ReadFile(".gitattributes")
	if err != nil {
		return
	}

	attributes := strings.NewReader(string(data))

	attributesFile, err := os.Create(".gitattributes")
	if err != nil {
		Print("Error opening .gitattributes for writing")
		return
	}

	scanner := bufio.NewScanner(attributes)

	// Iterate through each line of the attributes file and rewrite it,
	// if the path was meant to be untracked, omit it, and print a message instead.
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.Contains(line, "filter=lfs") {
			attributesFile.WriteString(line + "\n")
			continue
		}

		fields := strings.Fields(line)
		removeThisPath := false
		for _, t := range args {
			if t == fields[0] {
				removeThisPath = true
			}
		}

		if !removeThisPath {
			attributesFile.WriteString(line + "\n")
		} else {
			Print("Untracking %s", fields[0])
		}
	}

	attributesFile.Close()
}

func init() {
	RootCmd.AddCommand(untrackCmd)
}
<commit_msg>ンンン ンンン ンン<commit_after>package commands

import (
	"bufio"
	"io/ioutil"
	"os"
	"strings"

	"github.com/github/git-lfs/lfs"
	"github.com/github/git-lfs/vendor/_nuts/github.com/spf13/cobra"
)

var (
	untrackCmd = &cobra.Command{
		Use:   "untrack",
		Short: "Remove an entry from .gitattributes",
		Run:   untrackCommand,
	}
)

// untrackCommand takes a list of paths as an argument, and removes each path from the
// default attribtues file (.gitattributes), if it exists.
func untrackCommand(cmd *cobra.Command, args []string) {
	if lfs.LocalGitDir == "" {
		Print("Not a git repository.")
		os.Exit(128)
	}
	if lfs.LocalWorkingDir == "" {
		Print("This operation must be run in a work tree.")
		os.Exit(128)
	}

	lfs.InstallHooks(false)

	if len(args) < 1 {
		Print("git lfs untrack <path> [path]*")
		return
	}

	data, err := ioutil.ReadFile(".gitattributes")
	if err != nil {
		return
	}

	attributes := strings.NewReader(string(data))

	attributesFile, err := os.Create(".gitattributes")
	if err != nil {
		Print("Error opening .gitattributes for writing")
		return
	}
	defer attributesFile.Close()

	scanner := bufio.NewScanner(attributes)

	// Iterate through each line of the attributes file and rewrite it,
	// if the path was meant to be untracked, omit it, and print a message instead.
	for scanner.Scan() {
		line := scanner.Text()
		if !strings.Contains(line, "filter=lfs") {
			attributesFile.WriteString(line + "\n")
			continue
		}

		fields := strings.Fields(line)
		removeThisPath := false
		for _, t := range args {
			if t == fields[0] {
				removeThisPath = true
			}
		}

		if !removeThisPath {
			attributesFile.WriteString(line + "\n")
		} else {
			Print("Untracking %s", fields[0])
		}
	}
}

func init() {
	RootCmd.AddCommand(untrackCmd)
}
<|endoftext|>


<commit_before>package store

import (
	"context"
	"errors"
	"os"
	"path/filepath"

	"github.com/grafana/grafana/pkg/infra/filestorage"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/registry"
	ac "github.com/grafana/grafana/pkg/services/accesscontrol"
	"github.com/grafana/grafana/pkg/services/featuremgmt"
	"github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"
)

var grafanaStorageLogger = log.New("grafanaStorageLogger")

var ErrUnsupportedStorage = errors.New("storage does not support this operation")
var ErrUploadInternalError = errors.New("upload internal error")
var ErrValidationFailed = errors.New("request validation failed")
var ErrFileAlreadyExists = errors.New("file exists")
var ErrStorageNotFound = errors.New("storage not found")
var ErrAccessDenied = errors.New("access denied")

const RootPublicStatic = "public-static"
const RootResources = "resources"
const RootDevenv = "devenv"

const MAX_UPLOAD_SIZE = 1 * 1024 * 1024 // 3MB

type DeleteFolderCmd struct {
	Path  string `json:"path"`
	Force bool   `json:"force"`
}

type CreateFolderCmd struct {
	Path string `json:"path"`
}

type StorageService interface {
	registry.BackgroundService

	// List folder contents
	List(ctx context.Context, user *models.SignedInUser, path string) (*StorageListFrame, error)

	// Read raw file contents out of the store
	Read(ctx context.Context, user *models.SignedInUser, path string) (*filestorage.File, error)

	Upload(ctx context.Context, user *models.SignedInUser, req *UploadRequest) error

	Delete(ctx context.Context, user *models.SignedInUser, path string) error

	DeleteFolder(ctx context.Context, user *models.SignedInUser, cmd *DeleteFolderCmd) error

	CreateFolder(ctx context.Context, user *models.SignedInUser, cmd *CreateFolderCmd) error

	validateUploadRequest(ctx context.Context, user *models.SignedInUser, req *UploadRequest, storagePath string) validationResult

	// sanitizeUploadRequest sanitizes the upload request and converts it into a command accepted by the FileStorage API
	sanitizeUploadRequest(ctx context.Context, user *models.SignedInUser, req *UploadRequest, storagePath string) (*filestorage.UpsertFileCommand, error)
}

type storageServiceConfig struct {
	allowUnsanitizedSvgUpload bool
}

type standardStorageService struct {
	sql         *sqlstore.SQLStore
	tree        *nestedTree
	cfg         storageServiceConfig
	authService storageAuthService
}

func ProvideService(sql *sqlstore.SQLStore, features featuremgmt.FeatureToggles, cfg *setting.Cfg) StorageService {
	globalRoots := []storageRuntime{
		newDiskStorage(RootPublicStatic, "Public static files", &StorageLocalDiskConfig{
			Path: cfg.StaticRootPath,
			Roots: []string{
				"/testdata/",
				"/img/",
				"/gazetteer/",
				"/maps/",
			},
		}).setReadOnly(true).setBuiltin(true).
			setDescription("Access files from the static public files"),
	}

	// Development dashboards
	if setting.Env != setting.Prod {
		devenv := filepath.Join(cfg.StaticRootPath, "..", "devenv")
		if _, err := os.Stat(devenv); !os.IsNotExist(err) {
			// path/to/whatever exists
			s := newDiskStorage(RootDevenv, "Development Environment", &StorageLocalDiskConfig{
				Path: devenv,
				Roots: []string{
					"/dev-dashboards/",
				},
			}).setReadOnly(false).setDescription("Explore files within the developer environment directly")
			globalRoots = append(globalRoots, s)
		}
	}

	initializeOrgStorages := func(orgId int64) []storageRuntime {
		storages := make([]storageRuntime, 0)
		if features.IsEnabled(featuremgmt.FlagStorageLocalUpload) {
			storages = append(storages,
				newSQLStorage(RootResources,
					"Resources",
					&StorageSQLConfig{orgId: orgId}, sql).
					setBuiltin(true).
					setDescription("Upload custom resource files"))
		}

		return storages
	}

	authService := newStaticStorageAuthService(func(ctx context.Context, user *models.SignedInUser, storageName string) map[string]filestorage.PathFilter {
		if user == nil || !user.IsGrafanaAdmin {
			return nil
		}

		switch storageName {
		case RootPublicStatic:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  denyAllPathFilter,
				ActionFilesDelete: denyAllPathFilter,
			}
		case RootDevenv:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  denyAllPathFilter,
				ActionFilesDelete: denyAllPathFilter,
			}
		case RootResources:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  allowAllPathFilter,
				ActionFilesDelete: allowAllPathFilter,
			}
		default:
			return nil
		}
	})

	return newStandardStorageService(sql, globalRoots, initializeOrgStorages, authService)
}

func newStandardStorageService(sql *sqlstore.SQLStore, globalRoots []storageRuntime, initializeOrgStorages func(orgId int64) []storageRuntime, authService storageAuthService) *standardStorageService {
	rootsByOrgId := make(map[int64][]storageRuntime)
	rootsByOrgId[ac.GlobalOrgID] = globalRoots

	res := &nestedTree{
		initializeOrgStorages: initializeOrgStorages,
		rootsByOrgId:          rootsByOrgId,
	}
	res.init()
	return &standardStorageService{
		sql:         sql,
		tree:        res,
		authService: authService,
		cfg: storageServiceConfig{
			allowUnsanitizedSvgUpload: false,
		},
	}
}

func (s *standardStorageService) Run(ctx context.Context) error {
	grafanaStorageLogger.Info("storage starting")
	return nil
}

func getOrgId(user *models.SignedInUser) int64 {
	if user == nil {
		return ac.GlobalOrgID
	}

	return user.OrgId
}

func (s *standardStorageService) List(ctx context.Context, user *models.SignedInUser, path string) (*StorageListFrame, error) {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(path))
	return s.tree.ListFolder(ctx, getOrgId(user), path, guardian.getPathFilter(ActionFilesRead))
}

func (s *standardStorageService) Read(ctx context.Context, user *models.SignedInUser, path string) (*filestorage.File, error) {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(path))
	if !guardian.canView(path) {
		return nil, ErrAccessDenied
	}
	return s.tree.GetFile(ctx, getOrgId(user), path)
}

type UploadRequest struct {
	Contents           []byte
	MimeType           string // TODO: remove MimeType from the struct once we can infer it from file contents
	Path               string
	CacheControl       string
	ContentDisposition string
	Properties         map[string]string
	EntityType         EntityType

	OverwriteExistingFile bool
}

func (s *standardStorageService) Upload(ctx context.Context, user *models.SignedInUser, req *UploadRequest) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(req.Path))
	if !guardian.canWrite(req.Path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), req.Path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	validationResult := s.validateUploadRequest(ctx, user, req, storagePath)
	if !validationResult.ok {
		grafanaStorageLogger.Warn("file upload validation failed", "filetype", req.MimeType, "path", req.Path, "reason", validationResult.reason)
		return ErrValidationFailed
	}

	upsertCommand, err := s.sanitizeUploadRequest(ctx, user, req, storagePath)
	if err != nil {
		grafanaStorageLogger.Error("failed while sanitizing the upload request", "filetype", req.MimeType, "path", req.Path, "error", err)
		return ErrUploadInternalError
	}

	grafanaStorageLogger.Info("uploading a file", "filetype", req.MimeType, "path", req.Path)

	if !req.OverwriteExistingFile {
		file, err := root.Store().Get(ctx, storagePath)
		if err != nil {
			grafanaStorageLogger.Error("failed while checking file existence", "err", err, "path", req.Path)
			return ErrUploadInternalError
		}

		if file != nil {
			return ErrFileAlreadyExists
		}
	}

	if err := root.Store().Upsert(ctx, upsertCommand); err != nil {
		grafanaStorageLogger.Error("failed while uploading the file", "err", err, "path", req.Path)
		return ErrUploadInternalError
	}

	return nil
}

func (s *standardStorageService) DeleteFolder(ctx context.Context, user *models.SignedInUser, cmd *DeleteFolderCmd) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(cmd.Path))
	if !guardian.canDelete(cmd.Path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), cmd.Path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	if storagePath == "" {
		storagePath = filestorage.Delimiter
	}
	return root.Store().DeleteFolder(ctx, storagePath, &filestorage.DeleteFolderOptions{Force: true, AccessFilter: guardian.getPathFilter(ActionFilesDelete)})
}

func (s *standardStorageService) CreateFolder(ctx context.Context, user *models.SignedInUser, cmd *CreateFolderCmd) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(cmd.Path))
	if !guardian.canWrite(cmd.Path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), cmd.Path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	err := root.Store().CreateFolder(ctx, storagePath)
	if err != nil {
		return err
	}
	return nil
}

func (s *standardStorageService) Delete(ctx context.Context, user *models.SignedInUser, path string) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(path))
	if !guardian.canDelete(path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	err := root.Store().Delete(ctx, storagePath)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Storage: `system/branding` storage (#51987)<commit_after>package store

import (
	"context"
	"errors"
	"os"
	"path/filepath"

	"github.com/grafana/grafana/pkg/infra/filestorage"
	"github.com/grafana/grafana/pkg/infra/log"
	"github.com/grafana/grafana/pkg/models"
	"github.com/grafana/grafana/pkg/registry"
	ac "github.com/grafana/grafana/pkg/services/accesscontrol"
	"github.com/grafana/grafana/pkg/services/featuremgmt"
	"github.com/grafana/grafana/pkg/services/sqlstore"
	"github.com/grafana/grafana/pkg/setting"
)

var grafanaStorageLogger = log.New("grafanaStorageLogger")

var ErrUnsupportedStorage = errors.New("storage does not support this operation")
var ErrUploadInternalError = errors.New("upload internal error")
var ErrValidationFailed = errors.New("request validation failed")
var ErrFileAlreadyExists = errors.New("file exists")
var ErrStorageNotFound = errors.New("storage not found")
var ErrAccessDenied = errors.New("access denied")

const RootPublicStatic = "public-static"
const RootResources = "resources"
const RootDevenv = "devenv"
const RootSystem = "system"

const SystemBrandingStorage = "system/branding"

const MAX_UPLOAD_SIZE = 1 * 1024 * 1024 // 3MB

type DeleteFolderCmd struct {
	Path  string `json:"path"`
	Force bool   `json:"force"`
}

type CreateFolderCmd struct {
	Path string `json:"path"`
}

type StorageService interface {
	registry.BackgroundService

	// List folder contents
	List(ctx context.Context, user *models.SignedInUser, path string) (*StorageListFrame, error)

	// Read raw file contents out of the store
	Read(ctx context.Context, user *models.SignedInUser, path string) (*filestorage.File, error)

	Upload(ctx context.Context, user *models.SignedInUser, req *UploadRequest) error

	Delete(ctx context.Context, user *models.SignedInUser, path string) error

	DeleteFolder(ctx context.Context, user *models.SignedInUser, cmd *DeleteFolderCmd) error

	CreateFolder(ctx context.Context, user *models.SignedInUser, cmd *CreateFolderCmd) error

	validateUploadRequest(ctx context.Context, user *models.SignedInUser, req *UploadRequest, storagePath string) validationResult

	// sanitizeUploadRequest sanitizes the upload request and converts it into a command accepted by the FileStorage API
	sanitizeUploadRequest(ctx context.Context, user *models.SignedInUser, req *UploadRequest, storagePath string) (*filestorage.UpsertFileCommand, error)
}

type storageServiceConfig struct {
	allowUnsanitizedSvgUpload bool
}

type standardStorageService struct {
	sql         *sqlstore.SQLStore
	tree        *nestedTree
	cfg         storageServiceConfig
	authService storageAuthService
}

func ProvideService(sql *sqlstore.SQLStore, features featuremgmt.FeatureToggles, cfg *setting.Cfg) StorageService {
	globalRoots := []storageRuntime{
		newDiskStorage(RootPublicStatic, "Public static files", &StorageLocalDiskConfig{
			Path: cfg.StaticRootPath,
			Roots: []string{
				"/testdata/",
				"/img/",
				"/gazetteer/",
				"/maps/",
			},
		}).setReadOnly(true).setBuiltin(true).
			setDescription("Access files from the static public files"),
	}

	// Development dashboards
	if setting.Env != setting.Prod {
		devenv := filepath.Join(cfg.StaticRootPath, "..", "devenv")
		if _, err := os.Stat(devenv); !os.IsNotExist(err) {
			// path/to/whatever exists
			s := newDiskStorage(RootDevenv, "Development Environment", &StorageLocalDiskConfig{
				Path: devenv,
				Roots: []string{
					"/dev-dashboards/",
				},
			}).setReadOnly(false).setDescription("Explore files within the developer environment directly")
			globalRoots = append(globalRoots, s)
		}
	}

	initializeOrgStorages := func(orgId int64) []storageRuntime {
		storages := make([]storageRuntime, 0)
		if features.IsEnabled(featuremgmt.FlagStorageLocalUpload) {
			storages = append(storages,
				newSQLStorage(RootResources,
					"Resources",
					&StorageSQLConfig{orgId: orgId}, sql).
					setBuiltin(true).
					setDescription("Upload custom resource files"))
		}

		storages = append(storages,
			newSQLStorage(RootSystem,
				"System",
				&StorageSQLConfig{orgId: orgId},
				sql,
			).setBuiltin(true).setDescription("Grafana system storage"))

		return storages
	}

	authService := newStaticStorageAuthService(func(ctx context.Context, user *models.SignedInUser, storageName string) map[string]filestorage.PathFilter {
		if user == nil || !user.IsGrafanaAdmin {
			return nil
		}

		switch storageName {
		case RootPublicStatic:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  denyAllPathFilter,
				ActionFilesDelete: denyAllPathFilter,
			}
		case RootDevenv:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  denyAllPathFilter,
				ActionFilesDelete: denyAllPathFilter,
			}
		case RootResources:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  allowAllPathFilter,
				ActionFilesDelete: allowAllPathFilter,
			}
		case RootSystem:
			return map[string]filestorage.PathFilter{
				ActionFilesRead:   allowAllPathFilter,
				ActionFilesWrite:  allowAllPathFilter,
				ActionFilesDelete: allowAllPathFilter,
			}
		default:
			return nil
		}
	})

	return newStandardStorageService(sql, globalRoots, initializeOrgStorages, authService)
}

func newStandardStorageService(sql *sqlstore.SQLStore, globalRoots []storageRuntime, initializeOrgStorages func(orgId int64) []storageRuntime, authService storageAuthService) *standardStorageService {
	rootsByOrgId := make(map[int64][]storageRuntime)
	rootsByOrgId[ac.GlobalOrgID] = globalRoots

	res := &nestedTree{
		initializeOrgStorages: initializeOrgStorages,
		rootsByOrgId:          rootsByOrgId,
	}
	res.init()
	return &standardStorageService{
		sql:         sql,
		tree:        res,
		authService: authService,
		cfg: storageServiceConfig{
			allowUnsanitizedSvgUpload: false,
		},
	}
}

func (s *standardStorageService) Run(ctx context.Context) error {
	grafanaStorageLogger.Info("storage starting")
	return nil
}

func getOrgId(user *models.SignedInUser) int64 {
	if user == nil {
		return ac.GlobalOrgID
	}

	return user.OrgId
}

func (s *standardStorageService) List(ctx context.Context, user *models.SignedInUser, path string) (*StorageListFrame, error) {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(path))
	return s.tree.ListFolder(ctx, getOrgId(user), path, guardian.getPathFilter(ActionFilesRead))
}

func (s *standardStorageService) Read(ctx context.Context, user *models.SignedInUser, path string) (*filestorage.File, error) {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(path))
	if !guardian.canView(path) {
		return nil, ErrAccessDenied
	}
	return s.tree.GetFile(ctx, getOrgId(user), path)
}

type UploadRequest struct {
	Contents           []byte
	MimeType           string // TODO: remove MimeType from the struct once we can infer it from file contents
	Path               string
	CacheControl       string
	ContentDisposition string
	Properties         map[string]string
	EntityType         EntityType

	OverwriteExistingFile bool
}

func (s *standardStorageService) Upload(ctx context.Context, user *models.SignedInUser, req *UploadRequest) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(req.Path))
	if !guardian.canWrite(req.Path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), req.Path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	validationResult := s.validateUploadRequest(ctx, user, req, storagePath)
	if !validationResult.ok {
		grafanaStorageLogger.Warn("file upload validation failed", "filetype", req.MimeType, "path", req.Path, "reason", validationResult.reason)
		return ErrValidationFailed
	}

	upsertCommand, err := s.sanitizeUploadRequest(ctx, user, req, storagePath)
	if err != nil {
		grafanaStorageLogger.Error("failed while sanitizing the upload request", "filetype", req.MimeType, "path", req.Path, "error", err)
		return ErrUploadInternalError
	}

	grafanaStorageLogger.Info("uploading a file", "filetype", req.MimeType, "path", req.Path)

	if !req.OverwriteExistingFile {
		file, err := root.Store().Get(ctx, storagePath)
		if err != nil {
			grafanaStorageLogger.Error("failed while checking file existence", "err", err, "path", req.Path)
			return ErrUploadInternalError
		}

		if file != nil {
			return ErrFileAlreadyExists
		}
	}

	if err := root.Store().Upsert(ctx, upsertCommand); err != nil {
		grafanaStorageLogger.Error("failed while uploading the file", "err", err, "path", req.Path)
		return ErrUploadInternalError
	}

	return nil
}

func (s *standardStorageService) DeleteFolder(ctx context.Context, user *models.SignedInUser, cmd *DeleteFolderCmd) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(cmd.Path))
	if !guardian.canDelete(cmd.Path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), cmd.Path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	if storagePath == "" {
		storagePath = filestorage.Delimiter
	}
	return root.Store().DeleteFolder(ctx, storagePath, &filestorage.DeleteFolderOptions{Force: true, AccessFilter: guardian.getPathFilter(ActionFilesDelete)})
}

func (s *standardStorageService) CreateFolder(ctx context.Context, user *models.SignedInUser, cmd *CreateFolderCmd) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(cmd.Path))
	if !guardian.canWrite(cmd.Path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), cmd.Path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	err := root.Store().CreateFolder(ctx, storagePath)
	if err != nil {
		return err
	}
	return nil
}

func (s *standardStorageService) Delete(ctx context.Context, user *models.SignedInUser, path string) error {
	guardian := s.authService.newGuardian(ctx, user, getFirstSegment(path))
	if !guardian.canDelete(path) {
		return ErrAccessDenied
	}

	root, storagePath := s.tree.getRoot(getOrgId(user), path)
	if root == nil {
		return ErrStorageNotFound
	}

	if root.Meta().ReadOnly {
		return ErrUnsupportedStorage
	}

	err := root.Store().Delete(ctx, storagePath)
	if err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package parallel

import (
	"bytes"
	"context"
	"os"

	"github.com/werf/logboek"
	"github.com/werf/logboek/pkg/style"
	"github.com/werf/logboek/pkg/types"

	"github.com/werf/werf/pkg/docker"
)

type DoTasksOptions struct {
	InitDockerCLIForEachWorker bool
	MaxNumberOfWorkers         int
	IsLiveOutputOn             bool
}

func DoTasks(ctx context.Context, numberOfTasks int, options DoTasksOptions, taskFunc func(ctx context.Context, taskId int) error) error {
	if numberOfTasks == 0 {
		return nil
	}

	numberOfWorkers := options.MaxNumberOfWorkers
	if numberOfWorkers <= 0 || numberOfWorkers > numberOfTasks {
		numberOfWorkers = numberOfTasks
	}

	var numberOfTasksPerWorker []int
	for i := 0; i < numberOfWorkers; i++ {
		workerNumberOfTasks := numberOfTasks / numberOfWorkers
		rest := numberOfTasks % numberOfWorkers
		if rest > i {
			workerNumberOfTasks += 1
		}

		numberOfTasksPerWorker = append(numberOfTasksPerWorker, workerNumberOfTasks)
	}

	errCh := make(chan interface{})
	doneTaskCh := make(chan interface{})
	doneWorkerCh := make(chan worker)
	quitCh := make(chan bool)
	doneWorkersCounter := numberOfWorkers
	isLiveOutputOnFlag := options.IsLiveOutputOn

	var liveLogger types.LoggerInterface
	var liveContext context.Context
	if options.IsLiveOutputOn {
		liveLogger = logboek.NewLogger(os.Stdout, os.Stderr)
		liveLogger.GetStreamsSettingsFrom(logboek.Context(ctx))
		liveLogger.SetAcceptedLevel(logboek.Context(ctx).AcceptedLevel())

		liveContext = logboek.NewContext(ctx, liveLogger)

		if docker.IsContext(liveContext) {
			if err := docker.SyncContextCliWithLogger(liveContext); err != nil {
				return err
			}
			defer docker.SyncContextCliWithLogger(ctx)
		}
	}

	var workersBuffs []*bytes.Buffer
	var doneTaskDataList [][]byte
	for i := 0; i < numberOfWorkers; i++ {
		var workerContext context.Context
		var worker worker

		workerId := i

		if i == 0 && options.IsLiveOutputOn {
			workerContext = liveContext
			worker = &liveWorker{}
		} else {
			workerBuf := bytes.NewBuffer([]byte{})
			workersBuffs = append(workersBuffs, workerBuf)
			worker = &bufWorker{buf: workerBuf}

			workerContext = logboek.NewContext(ctx, logboek.Context(ctx).NewSubLogger(workerBuf, workerBuf))
			logboek.Context(workerContext).Streams().SetPrefixStyle(style.Highlight())

			if options.InitDockerCLIForEachWorker {
				workerContextWithDockerCli, err := docker.NewContext(workerContext)
				if err != nil {
					return err
				}

				workerContext = workerContextWithDockerCli
			}
		}

		go func() {
			workerNumberOfTasks := numberOfTasksPerWorker[workerId]

			for workerTaskId := 0; workerTaskId < workerNumberOfTasks; workerTaskId++ {
				taskId := calculateTaskId(numberOfTasks, numberOfWorkers, workerId, workerTaskId)
				if debug() {
					logboek.Context(workerContext).LogF("Running worker %d task %d/%d (%d)\n", workerId, workerTaskId+1, workerNumberOfTasks, numberOfTasks)
				}
				err := taskFunc(workerContext, taskId)

				ch := doneTaskCh
				if err != nil {
					ch = errCh
				}

				select {
				case ch <- worker.TaskResult(err):
					if err != nil {
						return
					}
				case <-quitCh:
					return
				}
			}

			doneWorkerCh <- worker
		}()
	}

	for {
		select {
		case res := <-doneTaskCh:
			switch taskResult := res.(type) {
			case *bufWorkerTaskResult:
				if isLiveOutputOnFlag {
					doneTaskDataList = append(doneTaskDataList, taskResult.data)
				} else {
					processTaskResultData(ctx, taskResult.data)
				}
			}
		case res := <-errCh:
			close(quitCh)

			switch taskResult := res.(type) {
			case *bufWorkerTaskResult:
				if isLiveOutputOnFlag {
					liveLogger.Streams().Mute()
				}

				for _, data := range doneTaskDataList {
					processTaskResultData(ctx, data)
				}

				for _, buf := range workersBuffs {
					if buf != taskResult.buf {
						processTaskResultData(ctx, []byte(buf.String()))
					}
				}

				processTaskResultData(ctx, taskResult.data)

				return taskResult.err
			case *lifeWorkerTaskResult:
				if logboek.Context(ctx).Info().IsAccepted() {
					for _, data := range doneTaskDataList {
						processTaskResultData(ctx, data)
					}

					for _, buf := range workersBuffs {
						processTaskResultData(ctx, buf.Bytes())
					}
				}

				return taskResult.err
			}
		case res := <-doneWorkerCh:
			if res.IsLiveWorker() {
				isLiveOutputOnFlag = false

				for _, data := range doneTaskDataList {
					processTaskResultData(ctx, data)
				}
			}

			doneWorkersCounter--
			if doneWorkersCounter == 0 {
				return nil
			}
		}
	}
}

func calculateTaskId(tasksNumber, workersNumber, workerInd, workerTaskId int) int {
	taskId := workerInd*(tasksNumber/workersNumber) + workerTaskId

	rest := tasksNumber % workersNumber
	if rest != 0 {
		if rest > workerInd {
			taskId += workerInd
		} else {
			taskId += rest
		}
	}

	return taskId
}

func processTaskResultData(ctx context.Context, data []byte) {
	logboek.Streams().DoWithoutIndent(func() {
		if logboek.Context(ctx).Streams().IsPrefixWithTimeEnabled() {
			logboek.Context(ctx).Streams().DisablePrefixWithTime()
			defer logboek.Context(ctx).Streams().EnablePrefixWithTime()
		}

		_, _ = logboek.Context(ctx).ProxyOutStream().Write(data)
		logboek.Context(ctx).LogOptionalLn()
	})
}

func debug() bool {
	return os.Getenv("WERF_DEBUG_PARALLEL") == "1"
}
<commit_msg>[logging] Print main time prefix with parallel task time prefix<commit_after>package parallel

import (
	"bytes"
	"context"
	"os"

	"github.com/werf/logboek"
	"github.com/werf/logboek/pkg/style"
	"github.com/werf/logboek/pkg/types"

	"github.com/werf/werf/pkg/docker"
)

type DoTasksOptions struct {
	InitDockerCLIForEachWorker bool
	MaxNumberOfWorkers         int
	IsLiveOutputOn             bool
}

func DoTasks(ctx context.Context, numberOfTasks int, options DoTasksOptions, taskFunc func(ctx context.Context, taskId int) error) error {
	if numberOfTasks == 0 {
		return nil
	}

	numberOfWorkers := options.MaxNumberOfWorkers
	if numberOfWorkers <= 0 || numberOfWorkers > numberOfTasks {
		numberOfWorkers = numberOfTasks
	}

	var numberOfTasksPerWorker []int
	for i := 0; i < numberOfWorkers; i++ {
		workerNumberOfTasks := numberOfTasks / numberOfWorkers
		rest := numberOfTasks % numberOfWorkers
		if rest > i {
			workerNumberOfTasks += 1
		}

		numberOfTasksPerWorker = append(numberOfTasksPerWorker, workerNumberOfTasks)
	}

	errCh := make(chan interface{})
	doneTaskCh := make(chan interface{})
	doneWorkerCh := make(chan worker)
	quitCh := make(chan bool)
	doneWorkersCounter := numberOfWorkers
	isLiveOutputOnFlag := options.IsLiveOutputOn

	var liveLogger types.LoggerInterface
	var liveContext context.Context
	if options.IsLiveOutputOn {
		liveLogger = logboek.NewLogger(os.Stdout, os.Stderr)
		liveLogger.GetStreamsSettingsFrom(logboek.Context(ctx))
		liveLogger.SetAcceptedLevel(logboek.Context(ctx).AcceptedLevel())

		liveContext = logboek.NewContext(ctx, liveLogger)

		if docker.IsContext(liveContext) {
			if err := docker.SyncContextCliWithLogger(liveContext); err != nil {
				return err
			}
			defer docker.SyncContextCliWithLogger(ctx)
		}
	}

	var workersBuffs []*bytes.Buffer
	var doneTaskDataList [][]byte
	for i := 0; i < numberOfWorkers; i++ {
		var workerContext context.Context
		var worker worker

		workerId := i

		if i == 0 && options.IsLiveOutputOn {
			workerContext = liveContext
			worker = &liveWorker{}
		} else {
			workerBuf := bytes.NewBuffer([]byte{})
			workersBuffs = append(workersBuffs, workerBuf)
			worker = &bufWorker{buf: workerBuf}

			workerContext = logboek.NewContext(ctx, logboek.Context(ctx).NewSubLogger(workerBuf, workerBuf))
			logboek.Context(workerContext).Streams().SetPrefixStyle(style.Highlight())

			if options.InitDockerCLIForEachWorker {
				workerContextWithDockerCli, err := docker.NewContext(workerContext)
				if err != nil {
					return err
				}

				workerContext = workerContextWithDockerCli
			}
		}

		go func() {
			workerNumberOfTasks := numberOfTasksPerWorker[workerId]

			for workerTaskId := 0; workerTaskId < workerNumberOfTasks; workerTaskId++ {
				taskId := calculateTaskId(numberOfTasks, numberOfWorkers, workerId, workerTaskId)
				if debug() {
					logboek.Context(workerContext).LogF("Running worker %d task %d/%d (%d)\n", workerId, workerTaskId+1, workerNumberOfTasks, numberOfTasks)
				}
				err := taskFunc(workerContext, taskId)

				ch := doneTaskCh
				if err != nil {
					ch = errCh
				}

				select {
				case ch <- worker.TaskResult(err):
					if err != nil {
						return
					}
				case <-quitCh:
					return
				}
			}

			doneWorkerCh <- worker
		}()
	}

	for {
		select {
		case res := <-doneTaskCh:
			switch taskResult := res.(type) {
			case *bufWorkerTaskResult:
				if isLiveOutputOnFlag {
					doneTaskDataList = append(doneTaskDataList, taskResult.data)
				} else {
					processTaskResultData(ctx, taskResult.data)
				}
			}
		case res := <-errCh:
			close(quitCh)

			switch taskResult := res.(type) {
			case *bufWorkerTaskResult:
				if isLiveOutputOnFlag {
					liveLogger.Streams().Mute()
				}

				for _, data := range doneTaskDataList {
					processTaskResultData(ctx, data)
				}

				for _, buf := range workersBuffs {
					if buf != taskResult.buf {
						processTaskResultData(ctx, []byte(buf.String()))
					}
				}

				processTaskResultData(ctx, taskResult.data)

				return taskResult.err
			case *lifeWorkerTaskResult:
				if len(workersBuffs) != 0 {
					if logboek.Context(ctx).Info().IsAccepted() {
						logboek.Context(liveContext).LogLn()

						for _, data := range doneTaskDataList {
							processTaskResultData(ctx, data)
						}

						for _, buf := range workersBuffs {
							processTaskResultData(ctx, buf.Bytes())
						}
					}
				}

				return taskResult.err
			}
		case res := <-doneWorkerCh:
			if res.IsLiveWorker() {
				isLiveOutputOnFlag = false

				for _, data := range doneTaskDataList {
					processTaskResultData(ctx, data)
				}
			}

			doneWorkersCounter--
			if doneWorkersCounter == 0 {
				return nil
			}
		}
	}
}

func calculateTaskId(tasksNumber, workersNumber, workerInd, workerTaskId int) int {
	taskId := workerInd*(tasksNumber/workersNumber) + workerTaskId

	rest := tasksNumber % workersNumber
	if rest != 0 {
		if rest > workerInd {
			taskId += workerInd
		} else {
			taskId += rest
		}
	}

	return taskId
}

func processTaskResultData(ctx context.Context, data []byte) {
	logboek.Streams().DoWithoutIndent(func() {
		_, _ = logboek.Context(ctx).ProxyOutStream().Write(data)
		logboek.Context(ctx).LogOptionalLn()
	})
}

func debug() bool {
	return os.Getenv("WERF_DEBUG_PARALLEL") == "1"
}
<|endoftext|>


<commit_before>package integration

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"

	. "github.com/onsi/gomega"
	"gopkg.in/yaml.v2"

	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/ziutek/syslog"

	"github.com/concourse/blackbox"
)

type SyslogServer struct {
	Addr string

	server *syslog.Server
}

func NewSyslogServer(inbox *Inbox) *SyslogServer {
	server := syslog.NewServer()
	server.AddHandler(inbox)

	return &SyslogServer{
		server: server,
	}
}

func (s *SyslogServer) Start() {
	l, err := net.Listen("tcp", ":0")
	Expect(err).NotTo(HaveOccurred())
	l.Close()

	_, port, err := net.SplitHostPort(l.Addr().String())
	Expect(err).NotTo(HaveOccurred())

	addr := fmt.Sprintf("127.0.0.1:%s", port)
	s.server.Listen(addr)

	s.Addr = addr
}

func (s *SyslogServer) Stop() {
	s.server.Shutdown()
	s.Addr = ""
}

type Inbox struct {
	Messages chan *syslog.Message
}

func NewInbox() *Inbox {
	return &Inbox{
		Messages: make(chan *syslog.Message),
	}
}

func (i *Inbox) Handle(m *syslog.Message) *syslog.Message {
	if m == nil {
		close(i.Messages)
		return nil
	}

	i.Messages <- m
	return nil
}

type BlackboxRunner struct {
	blackboxPath    string
	blackboxProcess ifrit.Process
}

func NewBlackboxRunner(blackboxPath string) *BlackboxRunner {
	return &BlackboxRunner{
		blackboxPath: blackboxPath,
	}
}

func (runner *BlackboxRunner) StartWithConfig(config blackbox.Config) {
	configPath := runner.createConfigFile(config)

	blackboxCmd := exec.Command(runner.blackboxPath, "-config", configPath)
	blackboxRunner := ginkgomon.New(
		ginkgomon.Config{
			Name:          "blackbox",
			Command:       blackboxCmd,
			AnsiColorCode: "90m",
			StartCheck:    "Seeked",
			Cleanup: func() {
				os.Remove(configPath)
			},
		},
	)

	runner.blackboxProcess = ginkgomon.Invoke(blackboxRunner)
}

func (runner *BlackboxRunner) Stop() {
	ginkgomon.Interrupt(runner.blackboxProcess)
}

func (runner *BlackboxRunner) createConfigFile(config blackbox.Config) string {
	configFile, err := ioutil.TempFile("", "blackbox_config")
	Expect(err).NotTo(HaveOccurred())
	defer configFile.Close()

	yamlToWrite, err := yaml.Marshal(config)
	Expect(err).NotTo(HaveOccurred())

	configFile.Write(yamlToWrite)

	return configFile.Name()
}
<commit_msg>use predefined port for syslog server<commit_after>package integration

import (
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"gopkg.in/yaml.v2"

	"github.com/tedsuo/ifrit"
	"github.com/tedsuo/ifrit/ginkgomon"
	"github.com/ziutek/syslog"

	"github.com/concourse/blackbox"
)

type SyslogServer struct {
	Addr string

	server *syslog.Server
}

func NewSyslogServer(inbox *Inbox) *SyslogServer {
	server := syslog.NewServer()
	server.AddHandler(inbox)

	return &SyslogServer{
		server: server,
	}
}

func (s *SyslogServer) Start() {
	port := fmt.Sprintf("%d", 9090+GinkgoParallelNode())
	l, err := net.Listen("tcp", ":"+port)
	Expect(err).NotTo(HaveOccurred())
	l.Close()

	addr := fmt.Sprintf("127.0.0.1:%s", port)
	err = s.server.Listen(addr)
	Expect(err).NotTo(HaveOccurred())

	s.Addr = addr
}

func (s *SyslogServer) Stop() {
	s.server.Shutdown()
	Eventually(func() error {
		_, err := net.Dial("tcp", s.Addr)
		return err
	}).ShouldNot(BeNil())

	s.Addr = ""
}

type Inbox struct {
	Messages chan *syslog.Message
}

func NewInbox() *Inbox {
	return &Inbox{
		Messages: make(chan *syslog.Message),
	}
}

func (i *Inbox) Handle(m *syslog.Message) *syslog.Message {
	if m == nil {
		close(i.Messages)
		return nil
	}

	i.Messages <- m
	return nil
}

type BlackboxRunner struct {
	blackboxPath    string
	blackboxProcess ifrit.Process
}

func NewBlackboxRunner(blackboxPath string) *BlackboxRunner {
	return &BlackboxRunner{
		blackboxPath: blackboxPath,
	}
}

func (runner *BlackboxRunner) StartWithConfig(config blackbox.Config) {
	configPath := runner.createConfigFile(config)

	blackboxCmd := exec.Command(runner.blackboxPath, "-config", configPath)
	blackboxRunner := ginkgomon.New(
		ginkgomon.Config{
			Name:          "blackbox",
			Command:       blackboxCmd,
			AnsiColorCode: "90m",
			StartCheck:    "Seeked",
			Cleanup: func() {
				os.Remove(configPath)
			},
		},
	)

	runner.blackboxProcess = ginkgomon.Invoke(blackboxRunner)
}

func (runner *BlackboxRunner) Stop() {
	ginkgomon.Interrupt(runner.blackboxProcess)
}

func (runner *BlackboxRunner) createConfigFile(config blackbox.Config) string {
	configFile, err := ioutil.TempFile("", "blackbox_config")
	Expect(err).NotTo(HaveOccurred())
	defer configFile.Close()

	yamlToWrite, err := yaml.Marshal(config)
	Expect(err).NotTo(HaveOccurred())

	configFile.Write(yamlToWrite)

	return configFile.Name()
}
<|endoftext|>