<commit_before>package log

//go:generate go run $GOPATH/src/v2ray.com/core/common/errors/errorgen/main.go -pkg log -path App,Log

import (
	"context"
	"sync"

	"v2ray.com/core/app/log/internal"
	"v2ray.com/core/common"
	"v2ray.com/core/common/log"
)

type Instance struct {
	sync.RWMutex
	config       *Config
	accessLogger internal.LogWriter
	errorLogger  internal.LogWriter
}

func New(ctx context.Context, config *Config) (*Instance, error) {
	return &Instance{
		config: config,
	}, nil
}

func (*Instance) Interface() interface{} {
	return (*Instance)(nil)
}

func (g *Instance) initAccessLogger() error {
	switch g.config.AccessLogType {
	case LogType_File:
		logger, err := internal.NewFileLogWriter(g.config.AccessLogPath)
		if err != nil {
			return err
		}
		g.accessLogger = logger
	case LogType_Console:
		g.accessLogger = internal.NewStdOutLogWriter()
	default:
	}
	return nil
}

func (g *Instance) initErrorLogger() error {
	switch g.config.ErrorLogType {
	case LogType_File:
		logger, err := internal.NewFileLogWriter(g.config.ErrorLogPath)
		if err != nil {
			return err
		}
		g.errorLogger = logger
	case LogType_Console:
		g.errorLogger = internal.NewStdOutLogWriter()
	default:
	}
	return nil
}

func (g *Instance) Start() error {
	if err := g.initAccessLogger(); err != nil {
		return newError("failed to initialize access logger").Base(err).AtWarning()
	}
	if err := g.initErrorLogger(); err != nil {
		return newError("failed to initialize error logger").Base(err).AtWarning()
	}
	log.RegisterHandler(g)
	return nil
}

func (g *Instance) Handle(msg log.Message) {
	switch msg := msg.(type) {
	case *log.AccessMessage:
		g.RLock()
		defer g.RUnlock()
		if g.accessLogger != nil {
			g.accessLogger.Log(msg)
		}
	case *log.GeneralMessage:
		if msg.Severity.SevererThan(g.config.ErrorLogLevel) {
			g.RLock()
			defer g.RUnlock()
			if g.errorLogger != nil {
				g.errorLogger.Log(msg)
			}
		}
	default:
		// Swallow
	}
}

func (g *Instance) Close() {
	g.Lock()
	defer g.Unlock()

	g.accessLogger.Close()
	g.accessLogger = nil

	g.errorLogger.Close()
	g.errorLogger = nil
}

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return New(ctx, config.(*Config))
	}))
}
<commit_msg>fix nil pointer on close().<commit_after>package log

//go:generate go run $GOPATH/src/v2ray.com/core/common/errors/errorgen/main.go -pkg log -path App,Log

import (
	"context"
	"sync"

	"v2ray.com/core/app/log/internal"
	"v2ray.com/core/common"
	"v2ray.com/core/common/log"
)

// Instance is an app.Application that handles logs.
type Instance struct {
	sync.RWMutex
	config       *Config
	accessLogger internal.LogWriter
	errorLogger  internal.LogWriter
}

// New creates a new log.Instance based on the given config.
func New(ctx context.Context, config *Config) (*Instance, error) {
	return &Instance{
		config: config,
	}, nil
}

// Interface implements app.Application.Interface().
func (*Instance) Interface() interface{} {
	return (*Instance)(nil)
}

func (g *Instance) initAccessLogger() error {
	switch g.config.AccessLogType {
	case LogType_File:
		logger, err := internal.NewFileLogWriter(g.config.AccessLogPath)
		if err != nil {
			return err
		}
		g.accessLogger = logger
	case LogType_Console:
		g.accessLogger = internal.NewStdOutLogWriter()
	default:
	}
	return nil
}

func (g *Instance) initErrorLogger() error {
	switch g.config.ErrorLogType {
	case LogType_File:
		logger, err := internal.NewFileLogWriter(g.config.ErrorLogPath)
		if err != nil {
			return err
		}
		g.errorLogger = logger
	case LogType_Console:
		g.errorLogger = internal.NewStdOutLogWriter()
	default:
	}
	return nil
}

// Start implements app.Application.Start().
func (g *Instance) Start() error {
	if err := g.initAccessLogger(); err != nil {
		return newError("failed to initialize access logger").Base(err).AtWarning()
	}
	if err := g.initErrorLogger(); err != nil {
		return newError("failed to initialize error logger").Base(err).AtWarning()
	}
	log.RegisterHandler(g)
	return nil
}

// Handle implements log.Handler.
func (g *Instance) Handle(msg log.Message) {
	switch msg := msg.(type) {
	case *log.AccessMessage:
		g.RLock()
		defer g.RUnlock()
		if g.accessLogger != nil {
			g.accessLogger.Log(msg)
		}
	case *log.GeneralMessage:
		if msg.Severity.SevererThan(g.config.ErrorLogLevel) {
			g.RLock()
			defer g.RUnlock()
			if g.errorLogger != nil {
				g.errorLogger.Log(msg)
			}
		}
	default:
		// Swallow
	}
}

// Close implement app.Application.Close().
func (g *Instance) Close() {
	g.Lock()
	defer g.Unlock()

	if g.accessLogger != nil {
		g.accessLogger.Close()
		g.accessLogger = nil
	}

	if g.errorLogger != nil {
		g.errorLogger.Close()
		g.errorLogger = nil
	}
}

func init() {
	common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, config interface{}) (interface{}, error) {
		return New(ctx, config.(*Config))
	}))
}
<|endoftext|>


<commit_before>package wrp

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"fmt"
	"github.com/stretchr/testify/assert"
	"strconv"
	"testing"
)

func encodeBase64(input string) string {
	var output bytes.Buffer
	encoder := base64.NewEncoder(base64.StdEncoding, &output)
	_, err := encoder.Write([]byte(input))
	if err != nil {
		panic(err)
	}

	if err = encoder.Close(); err != nil {
		panic(err)
	}

	return output.String()
}

var (
	// simpleRequestResponseMsgpack is a hand-coded example of a valid request/response message
	// in msgpack format
	simpleRequestResponseMsgpack = []byte{
		0x85, 0xa8, 0x6d, 0x73, 0x67, 0x5f, 0x74, 0x79,
		0x70, 0x65, 0x03, 0xb0, 0x74, 0x72, 0x61, 0x6e,
		0x73, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f,
		0x75, 0x75, 0x69, 0x64, 0xd9, 0x24, 0x39, 0x34,
		0x34, 0x37, 0x32, 0x34, 0x31, 0x63, 0x2d, 0x35,
		0x32, 0x33, 0x38, 0x2d, 0x34, 0x63, 0x62, 0x39,
		0x2d, 0x39, 0x62, 0x61, 0x61, 0x2d, 0x37, 0x30,
		0x37, 0x36, 0x65, 0x33, 0x32, 0x33, 0x32, 0x38,
		0x39, 0x39, 0xa6, 0x73, 0x6f, 0x75, 0x72, 0x63,
		0x65, 0xd9, 0x26, 0x64, 0x6e, 0x73, 0x3a, 0x77,
		0x65, 0x62, 0x70, 0x61, 0x2e, 0x63, 0x6f, 0x6d,
		0x63, 0x61, 0x73, 0x74, 0x2e, 0x63, 0x6f, 0x6d,
		0x2f, 0x76, 0x32, 0x2d, 0x64, 0x65, 0x76, 0x69,
		0x63, 0x65, 0x2d, 0x63, 0x6f, 0x6e, 0x66, 0x69,
		0x67, 0xa4, 0x64, 0x65, 0x73, 0x74, 0xb2, 0x73,
		0x65, 0x72, 0x69, 0x61, 0x6c, 0x3a, 0x31, 0x32,
		0x33, 0x34, 0x2f, 0x63, 0x6f, 0x6e, 0x66, 0x69,
		0x67, 0xa7, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61,
		0x64, 0xc4, 0x45, 0x7b, 0x20, 0x22, 0x6e, 0x61,
		0x6d, 0x65, 0x73, 0x22, 0x3a, 0x20, 0x5b, 0x20,
		0x22, 0x44, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2e,
		0x58, 0x5f, 0x43, 0x49, 0x53, 0x43, 0x4f, 0x5f,
		0x43, 0x4f, 0x4d, 0x5f, 0x53, 0x65, 0x63, 0x75,
		0x72, 0x69, 0x74, 0x79, 0x2e, 0x46, 0x69, 0x72,
		0x65, 0x77, 0x61, 0x6c, 0x6c, 0x2e, 0x46, 0x69,
		0x72, 0x65, 0x77, 0x61, 0x6c, 0x6c, 0x4c, 0x65,
		0x76, 0x65, 0x6c, 0x22, 0x20, 0x5d, 0x20, 0x7d,
	}

	// expectedPayload is the payload within simpleRequestResponseMsgpack
	expectedPayload       = `{ "names": [ "Device.X_CISCO_COM_Security.Firewall.FirewallLevel" ] }`
	expectedPayloadBase64 = encodeBase64(expectedPayload)

	expectedStatus      = int64(123)
	expectedSource      = "mac:112233445566"
	expectedDestination = "dns:somewhere.com/webhook"

	encoderTestData = []struct {
		original     Message
		expectedJSON string
	}{
		{
			original: Message{
				Type:   AuthMessageType,
				Status: &expectedStatus,
			},
			expectedJSON: `
				{"status": 123}
			`,
		},
		{
			original: Message{
				Type:        SimpleEventMessageType,
				Destination: expectedDestination,
				Payload:     []byte(expectedPayload),
			},
			expectedJSON: fmt.Sprintf(`{
				"dest": "%s",
				"payload": "%s"
			}`, expectedDestination, expectedPayloadBase64),
		},
		{
			original: Message{
				Type:        SimpleRequestResponseMessageType,
				Source:      expectedSource,
				Destination: expectedDestination,
				Payload:     []byte(expectedPayload),
			},
			expectedJSON: fmt.Sprintf(`{
				"source": "%s",
				"dest": "%s",
				"payload": "%s"
			}`, expectedSource, expectedDestination, expectedPayloadBase64),
		},
	}
)

// assertSimpleRequestResponse is an assertion specific to the hand-coded request/response message
func assertSimpleRequestResponse(assert *assert.Assertions, actual *Message) {
	assert.Equal(SimpleRequestResponseMessageType, actual.Type)
	assert.Nil(actual.Status)
	assert.Equal("dns:webpa.comcast.com/v2-device-config", actual.Source)
	assert.Equal("serial:1234/config", actual.Destination)
	assert.Equal("9447241c-5238-4cb9-9baa-7076e3232899", actual.TransactionUUID)
	assert.Equal([]byte(expectedPayload), actual.Payload)
	assert.Nil(actual.Valid())
}

// assertStringValue runs some sanity checks on the String() representation of a message.
// This is important because we want certain items output in logs.
func assertStringValue(assert *assert.Assertions, actual *Message) {
	stringValue := actual.String()
	assert.Contains(stringValue, actual.Type.String())
	assert.Contains(stringValue, actual.Source)
	assert.Contains(stringValue, actual.Destination)
	assert.Contains(stringValue, fmt.Sprintf("%v", actual.Payload))

	if actual.Status != nil {
		assert.Contains(stringValue, strconv.FormatInt(*actual.Status, 10))
	} else {
		assert.Contains(stringValue, "nil")
	}
}

func TestDecoderBytesMsgpackSimpleRequestResponse(t *testing.T) {
	assert := assert.New(t)

	var message Message
	assert.NotNil(message.Valid())

	decoder := NewDecoderBytes(simpleRequestResponseMsgpack, Msgpack)
	err := decoder.Decode(&message)
	assert.Nil(err)
	assertSimpleRequestResponse(assert, &message)
}

func TestDecoderMsgpackSimpleRequestResponse(t *testing.T) {
	assert := assert.New(t)

	var message Message
	assert.NotNil(message.Valid())

	output := bytes.NewBuffer(simpleRequestResponseMsgpack)
	decoder := NewDecoder(output, Msgpack)
	err := decoder.Decode(&message)
	assert.Nil(err)
	t.Logf("%s", &message)
	assertSimpleRequestResponse(assert, &message)
}

func TestEncoderMsgpack(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized bytes.Buffer
		encoder := NewEncoder(&serialized, Msgpack)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoder(&serialized, Msgpack)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestEncoderBytesMsgpack(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized []byte
		encoder := NewEncoderBytes(&serialized, Msgpack)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoderBytes(serialized, Msgpack)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestEncoderJSON(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized bytes.Buffer
		encoder := NewEncoder(&serialized, JSON)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoder(&serialized, JSON)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(MessageType(0), deserialized.Type)
		assert.Nil(deserialized.DeduceType())
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestEncoderBytesJSON(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized []byte
		encoder := NewEncoderBytes(&serialized, JSON)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoderBytes(serialized, JSON)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(MessageType(0), deserialized.Type)
		assert.Nil(deserialized.DeduceType())
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestFormatHandle(t *testing.T) {
	assert := assert.New(t)

	assert.NotNil(JSON.handle())
	assert.NotNil(Msgpack.handle())
	assert.Nil(Format(999).handle())
}

// BenchmarkMsgpackToJSONUsingEncoder benchmarks the typical use case for a WebPA server:
// reading a msgpack message then writing the JSON equivalent of that message.
func BenchmarkMsgpackToJSONUsingEncoder(b *testing.B) {
	b.StopTimer()

	var (
		decoder        = NewDecoder(nil, Msgpack)
		encoder        = NewEncoder(nil, JSON)
		bufferedReader = bufio.NewReader(nil)
		message        Message
		output         bytes.Buffer
	)

	b.StartTimer()
	for i := 0; i < b.N; i++ {
		bufferedReader.Reset(bytes.NewReader(simpleRequestResponseMsgpack))
		decoder.Reset(bufferedReader)

		if err := decoder.Decode(&message); err != nil {
			b.Fatal(err)
		}

		output.Reset()
		encoder.Reset(&output)
		if err := encoder.Encode(&message); err != nil {
			b.Fatal(err)
		}
	}
}
<commit_msg>Test for Format.String<commit_after>package wrp

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"fmt"
	"github.com/stretchr/testify/assert"
	"strconv"
	"testing"
)

func encodeBase64(input string) string {
	var output bytes.Buffer
	encoder := base64.NewEncoder(base64.StdEncoding, &output)
	_, err := encoder.Write([]byte(input))
	if err != nil {
		panic(err)
	}

	if err = encoder.Close(); err != nil {
		panic(err)
	}

	return output.String()
}

var (
	// simpleRequestResponseMsgpack is a hand-coded example of a valid request/response message
	// in msgpack format
	simpleRequestResponseMsgpack = []byte{
		0x85, 0xa8, 0x6d, 0x73, 0x67, 0x5f, 0x74, 0x79,
		0x70, 0x65, 0x03, 0xb0, 0x74, 0x72, 0x61, 0x6e,
		0x73, 0x61, 0x63, 0x74, 0x69, 0x6f, 0x6e, 0x5f,
		0x75, 0x75, 0x69, 0x64, 0xd9, 0x24, 0x39, 0x34,
		0x34, 0x37, 0x32, 0x34, 0x31, 0x63, 0x2d, 0x35,
		0x32, 0x33, 0x38, 0x2d, 0x34, 0x63, 0x62, 0x39,
		0x2d, 0x39, 0x62, 0x61, 0x61, 0x2d, 0x37, 0x30,
		0x37, 0x36, 0x65, 0x33, 0x32, 0x33, 0x32, 0x38,
		0x39, 0x39, 0xa6, 0x73, 0x6f, 0x75, 0x72, 0x63,
		0x65, 0xd9, 0x26, 0x64, 0x6e, 0x73, 0x3a, 0x77,
		0x65, 0x62, 0x70, 0x61, 0x2e, 0x63, 0x6f, 0x6d,
		0x63, 0x61, 0x73, 0x74, 0x2e, 0x63, 0x6f, 0x6d,
		0x2f, 0x76, 0x32, 0x2d, 0x64, 0x65, 0x76, 0x69,
		0x63, 0x65, 0x2d, 0x63, 0x6f, 0x6e, 0x66, 0x69,
		0x67, 0xa4, 0x64, 0x65, 0x73, 0x74, 0xb2, 0x73,
		0x65, 0x72, 0x69, 0x61, 0x6c, 0x3a, 0x31, 0x32,
		0x33, 0x34, 0x2f, 0x63, 0x6f, 0x6e, 0x66, 0x69,
		0x67, 0xa7, 0x70, 0x61, 0x79, 0x6c, 0x6f, 0x61,
		0x64, 0xc4, 0x45, 0x7b, 0x20, 0x22, 0x6e, 0x61,
		0x6d, 0x65, 0x73, 0x22, 0x3a, 0x20, 0x5b, 0x20,
		0x22, 0x44, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2e,
		0x58, 0x5f, 0x43, 0x49, 0x53, 0x43, 0x4f, 0x5f,
		0x43, 0x4f, 0x4d, 0x5f, 0x53, 0x65, 0x63, 0x75,
		0x72, 0x69, 0x74, 0x79, 0x2e, 0x46, 0x69, 0x72,
		0x65, 0x77, 0x61, 0x6c, 0x6c, 0x2e, 0x46, 0x69,
		0x72, 0x65, 0x77, 0x61, 0x6c, 0x6c, 0x4c, 0x65,
		0x76, 0x65, 0x6c, 0x22, 0x20, 0x5d, 0x20, 0x7d,
	}

	// expectedPayload is the payload within simpleRequestResponseMsgpack
	expectedPayload       = `{ "names": [ "Device.X_CISCO_COM_Security.Firewall.FirewallLevel" ] }`
	expectedPayloadBase64 = encodeBase64(expectedPayload)

	expectedStatus      = int64(123)
	expectedSource      = "mac:112233445566"
	expectedDestination = "dns:somewhere.com/webhook"

	encoderTestData = []struct {
		original     Message
		expectedJSON string
	}{
		{
			original: Message{
				Type:   AuthMessageType,
				Status: &expectedStatus,
			},
			expectedJSON: `
				{"status": 123}
			`,
		},
		{
			original: Message{
				Type:        SimpleEventMessageType,
				Destination: expectedDestination,
				Payload:     []byte(expectedPayload),
			},
			expectedJSON: fmt.Sprintf(`{
				"dest": "%s",
				"payload": "%s"
			}`, expectedDestination, expectedPayloadBase64),
		},
		{
			original: Message{
				Type:        SimpleRequestResponseMessageType,
				Source:      expectedSource,
				Destination: expectedDestination,
				Payload:     []byte(expectedPayload),
			},
			expectedJSON: fmt.Sprintf(`{
				"source": "%s",
				"dest": "%s",
				"payload": "%s"
			}`, expectedSource, expectedDestination, expectedPayloadBase64),
		},
	}
)

// assertSimpleRequestResponse is an assertion specific to the hand-coded request/response message
func assertSimpleRequestResponse(assert *assert.Assertions, actual *Message) {
	assert.Equal(SimpleRequestResponseMessageType, actual.Type)
	assert.Nil(actual.Status)
	assert.Equal("dns:webpa.comcast.com/v2-device-config", actual.Source)
	assert.Equal("serial:1234/config", actual.Destination)
	assert.Equal("9447241c-5238-4cb9-9baa-7076e3232899", actual.TransactionUUID)
	assert.Equal([]byte(expectedPayload), actual.Payload)
	assert.Nil(actual.Valid())
}

// assertStringValue runs some sanity checks on the String() representation of a message.
// This is important because we want certain items output in logs.
func assertStringValue(assert *assert.Assertions, actual *Message) {
	stringValue := actual.String()
	assert.Contains(stringValue, actual.Type.String())
	assert.Contains(stringValue, actual.Source)
	assert.Contains(stringValue, actual.Destination)
	assert.Contains(stringValue, fmt.Sprintf("%v", actual.Payload))

	if actual.Status != nil {
		assert.Contains(stringValue, strconv.FormatInt(*actual.Status, 10))
	} else {
		assert.Contains(stringValue, "nil")
	}
}

func TestFormatString(t *testing.T) {
	assert := assert.New(t)
	assert.NotEmpty(JSON.String())
	assert.NotEmpty(Msgpack.String())
	assert.NotEqual(JSON.String(), Msgpack.String())
}

func TestDecoderBytesMsgpackSimpleRequestResponse(t *testing.T) {
	assert := assert.New(t)

	var message Message
	assert.NotNil(message.Valid())

	decoder := NewDecoderBytes(simpleRequestResponseMsgpack, Msgpack)
	err := decoder.Decode(&message)
	assert.Nil(err)
	assertSimpleRequestResponse(assert, &message)
}

func TestDecoderMsgpackSimpleRequestResponse(t *testing.T) {
	assert := assert.New(t)

	var message Message
	assert.NotNil(message.Valid())

	output := bytes.NewBuffer(simpleRequestResponseMsgpack)
	decoder := NewDecoder(output, Msgpack)
	err := decoder.Decode(&message)
	assert.Nil(err)
	t.Logf("%s", &message)
	assertSimpleRequestResponse(assert, &message)
}

func TestEncoderMsgpack(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized bytes.Buffer
		encoder := NewEncoder(&serialized, Msgpack)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoder(&serialized, Msgpack)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestEncoderBytesMsgpack(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized []byte
		encoder := NewEncoderBytes(&serialized, Msgpack)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoderBytes(serialized, Msgpack)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestEncoderJSON(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized bytes.Buffer
		encoder := NewEncoder(&serialized, JSON)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoder(&serialized, JSON)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(MessageType(0), deserialized.Type)
		assert.Nil(deserialized.DeduceType())
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestEncoderBytesJSON(t *testing.T) {
	assert := assert.New(t)

	for _, record := range encoderTestData {
		t.Logf("%#v", record)

		var serialized []byte
		encoder := NewEncoderBytes(&serialized, JSON)
		assert.Nil(encoder.Encode(&record.original))
		assert.NotEmpty(serialized)

		var deserialized Message
		decoder := NewDecoderBytes(serialized, JSON)
		assert.Nil(decoder.Decode(&deserialized))
		assert.Equal(MessageType(0), deserialized.Type)
		assert.Nil(deserialized.DeduceType())
		assert.Equal(record.original, deserialized)

		assertStringValue(assert, &deserialized)
	}
}

func TestFormatHandle(t *testing.T) {
	assert := assert.New(t)

	assert.NotNil(JSON.handle())
	assert.NotNil(Msgpack.handle())
	assert.Nil(Format(999).handle())
}

// BenchmarkMsgpackToJSONUsingEncoder benchmarks the typical use case for a WebPA server:
// reading a msgpack message then writing the JSON equivalent of that message.
func BenchmarkMsgpackToJSONUsingEncoder(b *testing.B) {
	b.StopTimer()

	var (
		decoder        = NewDecoder(nil, Msgpack)
		encoder        = NewEncoder(nil, JSON)
		bufferedReader = bufio.NewReader(nil)
		message        Message
		output         bytes.Buffer
	)

	b.StartTimer()
	for i := 0; i < b.N; i++ {
		bufferedReader.Reset(bytes.NewReader(simpleRequestResponseMsgpack))
		decoder.Reset(bufferedReader)

		if err := decoder.Decode(&message); err != nil {
			b.Fatal(err)
		}

		output.Reset()
		encoder.Reset(&output)
		if err := encoder.Encode(&message); err != nil {
			b.Fatal(err)
		}
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"fmt"
	"io"
)

const (
	// MaxCFilterDataSize is the maximum byte size of a committed filter.
	MaxCFilterDataSize = 65536
)
type MsgCFilter struct {
	Data []byte
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
func (msg *MsgCFilter) BtcDecode(r io.Reader, pver uint32) error {
	var err error
	msg.Data, err = ReadVarBytes(r, pver, MaxCFilterDataSize,
	    "cfilter data")
	return err
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
func (msg *MsgCFilter) BtcEncode(w io.Writer, pver uint32) error {
	size := len(msg.Data)
	if size > MaxCFilterDataSize {
		str := fmt.Sprintf("cfilter size too large for message "+
			"[size %v, max %v]", size, MaxCFilterDataSize)
		return messageError("MsgCFilter.BtcEncode", str)
	}

	return WriteVarBytes(w, pver, msg.Data)
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgCFilter) Command() string {
	return CmdCFilter
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgCFilter) MaxPayloadLength(pver uint32) uint32 {
	return uint32(VarIntSerializeSize(MaxCFilterDataSize)) +
	    MaxCFilterDataSize
}

// NewMsgFilterAdd returns a new bitcoin filteradd message that conforms to the
// Message interface. See MsgCFilter for details.
func NewMsgCFilter(data []byte) *MsgCFilter {
	return &MsgCFilter{
		Data: data,
	}
}
<commit_msg>define MsgCFilter.Deserialize(), required by btcrpcclient<commit_after>// Copyright (c) 2017 The btcsuite developers
// Use of this source code is governed by an ISC
// license that can be found in the LICENSE file.

package wire

import (
	"fmt"
	"io"
)

const (
	// MaxCFilterDataSize is the maximum byte size of a committed filter.
	MaxCFilterDataSize = 65536
)
type MsgCFilter struct {
	Data []byte
}

// BtcDecode decodes r using the bitcoin protocol encoding into the receiver.
// This is part of the Message interface implementation.
func (msg *MsgCFilter) BtcDecode(r io.Reader, pver uint32) error {
	var err error
	msg.Data, err = ReadVarBytes(r, pver, MaxCFilterDataSize,
	    "cfilter data")
	return err
}

// BtcEncode encodes the receiver to w using the bitcoin protocol encoding.
// This is part of the Message interface implementation.
func (msg *MsgCFilter) BtcEncode(w io.Writer, pver uint32) error {
	size := len(msg.Data)
	if size > MaxCFilterDataSize {
		str := fmt.Sprintf("cfilter size too large for message "+
			"[size %v, max %v]", size, MaxCFilterDataSize)
		return messageError("MsgCFilter.BtcEncode", str)
	}

	return WriteVarBytes(w, pver, msg.Data)
}

// Deserialize decodes a filter from r into the receiver using a format that is
// suitable for long-term storage such as a database. This function differs
// from BtcDecode in that BtcDecode decodes from the bitcoin wire protocol as
// it was sent across the network.  The wire encoding can technically differ
// depending on the protocol version and doesn't even really need to match the
// format of a stored filter at all. As of the time this comment was written,
// the encoded filter is the same in both instances, but there is a distinct
// difference and separating the two allows the API to be flexible enough to
// deal with changes.
func (msg *MsgCFilter) Deserialize(r io.Reader) error {
	// At the current time, there is no difference between the wire encoding
	// and the stable long-term storage format.  As a result, make use of
	// BtcDecode.
	return msg.BtcDecode(r, 0)
}

// Command returns the protocol command string for the message.  This is part
// of the Message interface implementation.
func (msg *MsgCFilter) Command() string {
	return CmdCFilter
}

// MaxPayloadLength returns the maximum length the payload can be for the
// receiver.  This is part of the Message interface implementation.
func (msg *MsgCFilter) MaxPayloadLength(pver uint32) uint32 {
	return uint32(VarIntSerializeSize(MaxCFilterDataSize)) +
	    MaxCFilterDataSize
}

// NewMsgFilterAdd returns a new bitcoin filteradd message that conforms to the
// Message interface. See MsgCFilter for details.
func NewMsgCFilter(data []byte) *MsgCFilter {
	return &MsgCFilter{
		Data: data,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"flag"
	"log"

	"pixur.org/pixur/be/imaging"
	"pixur.org/pixur/be/schema"
	sdb "pixur.org/pixur/be/schema/db"
	tab "pixur.org/pixur/be/schema/tables"
	beconfig "pixur.org/pixur/be/server/config"
	"pixur.org/pixur/be/status"
)

func handlePic(db sdb.DB, picId int64) status.S {
	j, err := tab.NewJob(context.Background(), db)
	if err != nil {
		return status.InternalError(err, "cant make job")
	}
	defer j.Rollback()

	pics, err := j.FindPics(sdb.Opts{Lock: sdb.LockWrite, Prefix: tab.PicsPrimary{&picId}})
	if err != nil {
		return status.InternalError(err, "can't find pics")
	}
	if len(pics) != 1 {
		return status.NotFound(nil, "can't find pic")
	}
	p := pics[0]

	if p.FileSize == 0 {
		return nil
	}
	p.FileSize = 0
	p.Mime = schema.Pic_UNKNOWN
	p.Width = 0
	p.Height = 0
	p.AnimationInfo = nil

	if err := j.UpdatePic(p); err != nil {
		return status.InternalError(err, "can't update pic")
	}

	if err := j.Commit(); err != nil {
		return status.InternalError(err, "can't commit")
	}
	log.Println(p.GetVarPicID())

	return nil
}

func run() error {
	olddb, err := sdb.Open(beconfig.Conf.DbName, beconfig.Conf.DbConfig)
	if err != nil {
		return err
	}
	defer olddb.Close()

	oldj, err := tab.NewJob(context.Background(), olddb)
	if err != nil {
		return err
	}
	iiidd := int64(0)
	pics, err := oldj.FindPics(sdb.Opts{
		Start: tab.PicsPrimary{&iiidd},
		Lock:  sdb.LockNone,
	})
	oldj.Rollback()
	if err != nil {
		return err
	}

	for _, p := range pics {
		if err := handlePic(olddb, p.PicId); err != nil {
			return err
		}
	}

	return nil
}

func imageFormatToMime(f imaging.ImageFormat) (schema.Pic_File_Mime, status.S) {
	switch {
	case f.IsJpeg():
		return schema.Pic_File_JPEG, nil
	case f.IsGif():
		return schema.Pic_File_GIF, nil
	case f.IsPng():
		return schema.Pic_File_PNG, nil
	case f.IsWebm():
		return schema.Pic_File_WEBM, nil
	default:
		return schema.Pic_File_UNKNOWN, status.InvalidArgument(nil, "Unknown image type", f)
	}
}

func main() {
	flag.Parse()

	if err := run(); err != nil {
		log.Println(err)
	}
}
<commit_msg>tools: remove thumbnail scan<commit_after><|endoftext|>


<commit_before>package moulasaservice

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/moul/showcase"
)

func init() {
	http.HandleFunc("/", indexHandler)

	for name := range moulshowcase.Actions() {
		http.HandleFunc(fmt.Sprintf("/%s", name), actionHandler)
	}
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	var payload struct {
		Services []string `json:"services"`
	}
	payload.Services = make([]string, 0)
	for action := range moulshowcase.Actions() {
		payload.Services = append(payload.Services, fmt.Sprintf("/%s", action))
	}
	enc := json.NewEncoder(w)
	if err := enc.Encode(payload); err != nil {
		http.Error(w, fmt.Sprintf("json encode error: %v\n", err), 500)
	}
}

func actionHandler(w http.ResponseWriter, r *http.Request) {
	path := strings.TrimLeft(r.URL.Path, "/")
	if fn, found := moulshowcase.Actions()[path]; found {
		// parse CLI arguments
		u, err := url.Parse(r.URL.String())
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to parse url %q: %v", r.URL.String(), err), 500)
		}

		// call action
		ret, err := fn(u.RawQuery)

		// render result
		if err != nil {
			http.Error(w, fmt.Sprintf("service error: %v\n", err), 500)
		} else {
			switch ret.ContentType {
			case "application/json":
				w.Header().Set("Content-Type", "application/json; charset=utf-8")
				enc := json.NewEncoder(w)
				if err := enc.Encode(&(ret.Body)); err != nil {
					http.Error(w, fmt.Sprintf("json encode error: %v\n", err), 500)
				}
			case "text/plain":
				w.Header().Set("Content-Type", "text/plain; charset=utf-8")
				fmt.Fprintf(w, "%s", ret.Body)
			}
		}
	} else {
		http.NotFound(w, r)
	}
}
<commit_msg>Passing r.Body as stdin in appspot<commit_after>package moulasaservice

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"github.com/moul/showcase"
)

func init() {
	http.HandleFunc("/", indexHandler)

	for name := range moulshowcase.Actions() {
		http.HandleFunc(fmt.Sprintf("/%s", name), actionHandler)
	}
}

func indexHandler(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	var payload struct {
		Services []string `json:"services"`
	}
	payload.Services = make([]string, 0)
	for action := range moulshowcase.Actions() {
		payload.Services = append(payload.Services, fmt.Sprintf("/%s", action))
	}
	enc := json.NewEncoder(w)
	if err := enc.Encode(payload); err != nil {
		http.Error(w, fmt.Sprintf("json encode error: %v\n", err), 500)
	}
}

func actionHandler(w http.ResponseWriter, r *http.Request) {
	path := strings.TrimLeft(r.URL.Path, "/")
	if fn, found := moulshowcase.Actions()[path]; found {
		// parse CLI arguments
		u, err := url.Parse(r.URL.String())
		if err != nil {
			http.Error(w, fmt.Sprintf("failed to parse url %q: %v", r.URL.String(), err), 500)
		}

		// call action
		ret, err := fn(u.RawQuery, r.Body)

		// render result
		if err != nil {
			http.Error(w, fmt.Sprintf("service error: %v\n", err), 500)
		} else {
			switch ret.ContentType {
			case "application/json":
				w.Header().Set("Content-Type", "application/json; charset=utf-8")
				enc := json.NewEncoder(w)
				if err := enc.Encode(&(ret.Body)); err != nil {
					http.Error(w, fmt.Sprintf("json encode error: %v\n", err), 500)
				}
			case "text/plain":
				w.Header().Set("Content-Type", "text/plain; charset=utf-8")
				fmt.Fprintf(w, "%s", ret.Body)
			}
		}
	} else {
		http.NotFound(w, r)
	}
}
<|endoftext|>


<commit_before>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
*/

//
// NOTE: DO NOT EDIT. This file was generated by genwrap.go from the proton header files.
// Update the generator and re-run if you need to modify this code.
//

package amqp

// Version check for proton library.
// Done here because this is the lowest-level dependency for all the proton Go packages.

// #include <proton/version.h>
// #if PN_VERSION_MAJOR == 0 && PN_VERSION_MINOR < 27
// #error packages qpid.apache.org/... require Proton-C library version 0.10 or greater
// #endif
import "C"
<commit_msg>PROTON-2017: [go] fix proton-c version check<commit_after>/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing,
software distributed under the License is distributed on an
"AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
KIND, either express or implied.  See the License for the
specific language governing permissions and limitations
under the License.
*/

//
// NOTE: DO NOT EDIT. This file was generated by genwrap.go from the proton header files.
// Update the generator and re-run if you need to modify this code.
//

package amqp

// Version check for compatible proton-c library.
//
// NOTE: the required version should NOT be increased unless the Go
// library is modified to require some new proton-c API. That hasn't
// happened for a long time.

// #include <proton/version.h>
// #if PN_VERSION_MAJOR == 0 && PN_VERSION_MINOR < 10
// #error packages qpid.apache.org/... require Proton-C library version 0.10 or greater
// #endif
import "C"
<|endoftext|>


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

package discovery

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/prometheus/prometheus/config"
	yaml "gopkg.in/yaml.v2"
)

func TestTargetSetRecreatesTargetGroupsEveryRun(t *testing.T) {

	verifyPresence := func(tgroups map[string]*config.TargetGroup, name string, present bool) {
		if _, ok := tgroups[name]; ok != present {
			msg := ""
			if !present {
				msg = "not "
			}
			t.Fatalf("'%s' should %sbe present in TargetSet.tgroups: %s", name, msg, tgroups)
		}
	}

	cfg := &config.ServiceDiscoveryConfig{}

	sOne := `
static_configs:
- targets: ["foo:9090"]
- targets: ["bar:9090"]
`
	if err := yaml.Unmarshal([]byte(sOne), cfg); err != nil {
		t.Fatalf("Unable to load YAML config sOne: %s", err)
	}
	called := make(chan struct{})

	ts := NewTargetSet(&mockSyncer{
		sync: func([]*config.TargetGroup) { called <- struct{}{} },
	})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go ts.Run(ctx)

	ts.UpdateProviders(ProvidersFromConfig(*cfg, nil))
	<-called

	verifyPresence(ts.tgroups, "static/0/0", true)
	verifyPresence(ts.tgroups, "static/0/1", true)

	sTwo := `
static_configs:
- targets: ["foo:9090"]
`
	if err := yaml.Unmarshal([]byte(sTwo), cfg); err != nil {
		t.Fatalf("Unable to load YAML config sTwo: %s", err)
	}

	ts.UpdateProviders(ProvidersFromConfig(*cfg, nil))
	<-called

	verifyPresence(ts.tgroups, "static/0/0", true)
	verifyPresence(ts.tgroups, "static/0/1", false)
}

type mockSyncer struct {
	sync func(tgs []*config.TargetGroup)
}

func (s *mockSyncer) Sync(tgs []*config.TargetGroup) {
	if s.sync != nil {
		s.sync(tgs)
	}
}

type mockTargetProvider struct {
	callCount *uint32
}

func (tp mockTargetProvider) Run(ctx context.Context, up chan<- []*config.TargetGroup) {
	atomic.AddUint32(tp.callCount, 1)
	up <- []*config.TargetGroup{{Source: "dummySource"}}
}

func TestTargetSetRunsSameTargetProviderMultipleTimes(t *testing.T) {

	var wg sync.WaitGroup

	wg.Add(2)

	ts1 := NewTargetSet(&mockSyncer{
		sync: func([]*config.TargetGroup) { wg.Done() },
	})

	ts2 := NewTargetSet(&mockSyncer{
		sync: func([]*config.TargetGroup) { wg.Done() },
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	tp := mockTargetProvider{}
	var callCount uint32
	tp.callCount = &callCount

	targetProviders := map[string]TargetProvider{}
	targetProviders["testProvider"] = tp

	go ts1.Run(ctx)
	go ts2.Run(ctx)

	ts1.UpdateProviders(targetProviders)
	ts2.UpdateProviders(targetProviders)
	wg.Wait()

	if callCount != 2 {
		t.Errorf("Was expecting 2 calls received %v", callCount)
	}
}
<commit_msg>Add discovery test (#3417)<commit_after>// Copyright 2016 The Prometheus Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package discovery

import (
	"context"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/prometheus/prometheus/config"
	yaml "gopkg.in/yaml.v2"
)

func TestSingleTargetSetWithSingleProviderOnlySendsNewTargetGroups(t *testing.T) {

	testCases := [][]update{
		[]update{}, // No updates.

		[]update{{[]string{}, 0}}, // Empty initials.

		[]update{{[]string{}, 6000}}, // Empty initials with a delay.

		[]update{{[]string{"initial1", "initial2"}, 0}}, // Initials only.

		[]update{{[]string{"initial1", "initial2"}, 6000}}, // Initials only but after a delay.

		[]update{ // Initials and new groups.
			{[]string{"initial1", "initial2"}, 0},
			{[]string{"update1", "update2"}, 0},
		},

		[]update{ // Initials and new groups after a delay.
			{[]string{"initial1", "initial2"}, 6000},
			{[]string{"update1", "update2"}, 500},
		},

		[]update{
			{[]string{"initial1", "initial2"}, 100},
			{[]string{"update1", "update2", "update3", "update4", "update5", "update6", "update7", "update8", "update9", "update10", "update11"}, 100},
		},

		[]update{
			{[]string{"initial1"}, 10},
			{[]string{"update1"}, 45},
			{[]string{"update2", "update3", "update4"}, 0},
			{[]string{"update5"}, 10},
			{[]string{"update6", "update7", "update8", "update9"}, 70},
		},

		[]update{
			{[]string{"initial1", "initial2"}, 5},
			{[]string{}, 100},
			{[]string{"update1", "update2"}, 100},
			{[]string{"update3", "update4", "update5"}, 70},
		},
	}

	for i, updates := range testCases {

		expectedGroups := make(map[string]struct{})
		for _, update := range updates {
			for _, target := range update.targets {
				expectedGroups[target] = struct{}{}
			}
		}

		var wg sync.WaitGroup
		wg.Add(1)

		isFirstSyncCall := true
		var initialGroups []*config.TargetGroup
		var syncedGroups []*config.TargetGroup

		targetSet := NewTargetSet(&mockSyncer{
			sync: func(tgs []*config.TargetGroup) {
				syncedGroups = tgs

				if isFirstSyncCall {
					isFirstSyncCall = false
					initialGroups = tgs
				}

				if len(tgs) == len(expectedGroups) {
					// All the groups are sent, we can start asserting.
					wg.Done()
				}
			},
		})

		ctx, cancel := context.WithCancel(context.Background())
		defer cancel()

		tp := newMockTargetProvider(updates)
		targetProviders := map[string]TargetProvider{}
		targetProviders["testProvider"] = tp

		go targetSet.Run(ctx)
		targetSet.UpdateProviders(targetProviders)

		finalize := make(chan struct{})
		go func() {
			defer close(finalize)
			wg.Wait()
		}()

		select {
		case <-time.After(20000 * time.Millisecond):
			t.Errorf("In test case %v: Test timed out after 20000 millisecond. All targets should be sent within the timeout", i)

		case <-finalize:

			if *tp.callCount != 1 {
				t.Errorf("In test case %v: TargetProvider Run should be called once only, was called %v times", i, *tp.callCount)
			}

			if len(updates) > 0 && updates[0].interval > 5000 {
				// If the initial set of targets never arrive or arrive after 5 seconds.
				// The first sync call should receive empty set of targets.
				if len(initialGroups) != 0 {
					t.Errorf("In test case %v: Expecting 0 initial target groups, received %v", i, len(initialGroups))
				}
			}

			if len(syncedGroups) != len(expectedGroups) {
				t.Errorf("In test case %v: Expecting %v target groups in total, received %v", i, len(expectedGroups), len(syncedGroups))
			}

			for _, tg := range syncedGroups {
				if _, ok := expectedGroups[tg.Source]; ok == false {
					t.Errorf("In test case %v: '%s' does not exist in expected target groups: %s", i, tg.Source, expectedGroups)
				} else {
					delete(expectedGroups, tg.Source) // Remove used targets from the map.
				}
			}
		}
	}
}

func TestTargetSetRecreatesTargetGroupsEveryRun(t *testing.T) {

	verifyPresence := func(tgroups map[string]*config.TargetGroup, name string, present bool) {
		if _, ok := tgroups[name]; ok != present {
			msg := ""
			if !present {
				msg = "not "
			}
			t.Fatalf("'%s' should %sbe present in TargetSet.tgroups: %s", name, msg, tgroups)
		}
	}

	cfg := &config.ServiceDiscoveryConfig{}

	sOne := `
static_configs:
- targets: ["foo:9090"]
- targets: ["bar:9090"]
`
	if err := yaml.Unmarshal([]byte(sOne), cfg); err != nil {
		t.Fatalf("Unable to load YAML config sOne: %s", err)
	}
	called := make(chan struct{})

	ts := NewTargetSet(&mockSyncer{
		sync: func([]*config.TargetGroup) { called <- struct{}{} },
	})
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	go ts.Run(ctx)

	ts.UpdateProviders(ProvidersFromConfig(*cfg, nil))
	<-called

	verifyPresence(ts.tgroups, "static/0/0", true)
	verifyPresence(ts.tgroups, "static/0/1", true)

	sTwo := `
static_configs:
- targets: ["foo:9090"]
`
	if err := yaml.Unmarshal([]byte(sTwo), cfg); err != nil {
		t.Fatalf("Unable to load YAML config sTwo: %s", err)
	}

	ts.UpdateProviders(ProvidersFromConfig(*cfg, nil))
	<-called

	verifyPresence(ts.tgroups, "static/0/0", true)
	verifyPresence(ts.tgroups, "static/0/1", false)
}

func TestTargetSetRunsSameTargetProviderMultipleTimes(t *testing.T) {

	var wg sync.WaitGroup

	wg.Add(2)

	ts1 := NewTargetSet(&mockSyncer{
		sync: func([]*config.TargetGroup) { wg.Done() },
	})

	ts2 := NewTargetSet(&mockSyncer{
		sync: func([]*config.TargetGroup) { wg.Done() },
	})

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	tp := newMockTargetProvider([]update{{[]string{"initial1", "initial2"}, 10}})
	targetProviders := map[string]TargetProvider{}
	targetProviders["testProvider"] = tp

	go ts1.Run(ctx)
	go ts2.Run(ctx)

	ts1.UpdateProviders(targetProviders)
	ts2.UpdateProviders(targetProviders)
	wg.Wait()

	if *tp.callCount != 2 {
		t.Errorf("Was expecting 2 calls received %v", tp.callCount)
	}
}

type mockSyncer struct {
	sync func(tgs []*config.TargetGroup)
}

func (s *mockSyncer) Sync(tgs []*config.TargetGroup) {
	if s.sync != nil {
		s.sync(tgs)
	}
}

type mockTargetProvider struct {
	callCount *uint32
	updates   []update
	up        chan<- []*config.TargetGroup
}

type update struct {
	targets  []string
	interval time.Duration
}

func newMockTargetProvider(updates []update) mockTargetProvider {
	var callCount uint32

	tp := mockTargetProvider{
		callCount: &callCount,
		updates:   updates,
	}

	return tp
}

func (tp mockTargetProvider) Run(ctx context.Context, up chan<- []*config.TargetGroup) {
	atomic.AddUint32(tp.callCount, 1)
	tp.up = up
	tp.sendUpdates()
}

func (tp mockTargetProvider) sendUpdates() {
	for _, update := range tp.updates {

		time.Sleep(update.interval * time.Millisecond)

		tgs := make([]*config.TargetGroup, len(update.targets))
		for i, tg := range update.targets {
			tgs[i] = &config.TargetGroup{Source: tg}
		}

		tp.up <- tgs
	}
}
<|endoftext|>


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

package typegraph

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"

	"github.com/serulian/compiler/compilergraph"
)

// Value string example:
//
// [9242167c-2d57-4212-8aea-fedf32bd708e]T00010000G000039[04f97d44-d8fc-4a7d-9c46-955d1bd5add6]F
//  ^ Type ID of SomeType                ^ Nullable
//                                                ^ G000039 = generic with 39 chars in length for subreference
//                                        ^ 0001 = 1 generic
//                                            ^ 0000 = 0 parameters
// represents
//
// SomeType<SomeGeneric>?
//
// where SomeType is 9242167c-2d57-4212-8aea-fedf32bd708e
// and SomeGeneric is 04f97d44-d8fc-4a7d-9c46-955d1bd5add6
type typeReferenceHeaderSlot struct {
	index     int  // The slot index.
	length    int  // The length of this slot.
	writeable bool // Whether this slot is writeable.
}

var (
	trhSlotStartTypeId        = typeReferenceHeaderSlot{0, 1, false}
	trhSlotTypeId             = typeReferenceHeaderSlot{1, compilergraph.NodeIDLength, true}
	trhSlotEndTypeId          = typeReferenceHeaderSlot{2, 1, false}
	trhSlotFlagNullable       = typeReferenceHeaderSlot{3, 1, true}
	trhSlotGenericCount       = typeReferenceHeaderSlot{4, 4, true}
	trhSlotParameterCount     = typeReferenceHeaderSlot{5, 4, true}
	trhSlotSubReferenceMarker = typeReferenceHeaderSlot{6, 0, false}
)

var typeReferenceHeaderSlots = [...]typeReferenceHeaderSlot{
	trhSlotStartTypeId,
	trhSlotTypeId,
	trhSlotEndTypeId,
	trhSlotFlagNullable,
	trhSlotGenericCount,
	trhSlotParameterCount,
	trhSlotSubReferenceMarker,
}

// The size of the length prefix for subreferences.
const typeRefValueSubReferenceLength = 6

// subReferenceKind represents the kinds of supported subreferences
type subReferenceKind int

const (
	subReferenceGeneric subReferenceKind = iota
	subReferenceParameter
)

// withFlag returns a copy of this type reference with the flag at the given slot replaced with the
// specified rune.
func (tr *TypeReference) withFlag(flagSlot typeReferenceHeaderSlot, value rune) TypeReference {
	return tr.replaceSlot(flagSlot, string(value))
}

// replaceSlot replaces the given slot with the specified value.
func (tr *TypeReference) replaceSlot(slot typeReferenceHeaderSlot, value string) TypeReference {
	if !slot.writeable {
		panic(fmt.Sprintf("Cannot write to slot %v", slot))
	}

	location := getSlotLocation(slot)
	size := slot.length

	return TypeReference{
		layer: tr.layer,
		value: tr.value[:location] + value + tr.value[location+size:],
	}
}

// getSlotAsInt returns the int found at the given slot.
func (tr *TypeReference) getSlotAsInt(slot typeReferenceHeaderSlot) int {
	strValue := tr.getSlot(slot)

	// Special common case.
	if strValue == "0000" {
		return 0
	}

	i, err := strconv.Atoi(strValue)
	if err != nil {
		panic(fmt.Sprintf("Expected int value for slot %v, found: %v", slot, strValue))
	}
	return i
}

// getSlot returns the string found at the given slot.
func (tr *TypeReference) getSlot(slot typeReferenceHeaderSlot) string {
	location := getSlotLocation(slot)
	size := slot.length

	return tr.value[location : location+size]
}

// lengthPrefixedValue returns this type reference's value string, prefixed with its length value
// padded to the correct number of digits.
func (tr *TypeReference) lengthPrefixedValue() string {
	return padNumberToString(len(tr.value), typeRefValueSubReferenceLength) + tr.value
}

// withSubReference returns a copy of this type reference with the given subreference added.
func (tr *TypeReference) withSubReference(kind subReferenceKind, subref TypeReference) TypeReference {
	slot, kindRune := getSubReferenceSlotAndChar(kind)
	count := tr.getSlotAsInt(slot)

	// Add to the count and append the subreference.
	updated := tr.replaceSlot(slot, padNumberToString(count+1, slot.length))
	return TypeReference{
		layer: tr.layer,
		value: updated.value + string(kindRune) + subref.lengthPrefixedValue(),
	}
}

// getSubReferences returns all sub references of the given kind on this type reference.
func (tr *TypeReference) getSubReferences(kind subReferenceKind) []TypeReference {
	slot, kindChar := getSubReferenceSlotAndChar(kind)

	// Find the number of applicable subreferences.
	count := tr.getSlotAsInt(slot)
	subrefs := make([]TypeReference, count)

	if count == 0 {
		return subrefs
	}

	// For all subreferences loop and filter.
	var currentIndex int = getSlotLocation(trhSlotSubReferenceMarker)
	var collectedCount int = 0

	for {
		// Retrieve the rune for the subreference (G or P)
		subReferenceChar := tr.value[currentIndex]

		// Retrieve the length of the subreference's value string.
		subReferenceLengthStr := tr.value[currentIndex+1 : currentIndex+1+typeRefValueSubReferenceLength]
		subReferenceLength, err := strconv.Atoi(subReferenceLengthStr)
		if err != nil {
			panic(fmt.Sprintf("Expected int value for subreference length, found: %v", subReferenceLengthStr))
		}

		// Move the current index forward to the point at which the subreference value string beings.
		currentIndex = currentIndex + 1 + typeRefValueSubReferenceLength
		subReferenceValue := tr.value[currentIndex : currentIndex+subReferenceLength]

		// If we have found a subreference of the known kind, add it.
		if subReferenceChar == kindChar {
			subrefs[collectedCount] = TypeReference{
				layer: tr.layer,
				value: subReferenceValue,
			}

			collectedCount = collectedCount + 1
		}

		// Once we've found all the subreferences of the expected kind, we're done.
		if collectedCount == count {
			return subrefs
		}

		// Move the current index past the subreference value string.
		currentIndex = currentIndex + subReferenceLength
	}
}

// getSubReferenceSlotAndChar returns the count slot and character for the given kind of subreference.
func getSubReferenceSlotAndChar(kind subReferenceKind) (typeReferenceHeaderSlot, uint8) {
	switch kind {
	case subReferenceGeneric:
		return trhSlotGenericCount, 'G'

	case subReferenceParameter:
		return trhSlotParameterCount, 'P'

	default:
		panic("Unknown kind of subreference")
		return trhSlotGenericCount, '_'
	}
}

// getSlotLocation returns the slot location (0-indexed) in the value string.
func getSlotLocation(slot typeReferenceHeaderSlot) int {
	// TODO(jschorr): Add a map here to be a cache?
	var location int
	for _, currentSlot := range typeReferenceHeaderSlots {
		if currentSlot == slot {
			return location
		}

		location = location + currentSlot.length
	}

	panic("Could not retrieve location for slot")
	return -1
}

// buildTypeReferenceValue returns a string value for representing the given type reference data in a single
// string.
func buildTypeReferenceValue(typeNode compilergraph.GraphNode, nullable bool, generics ...TypeReference) string {
	var buffer bytes.Buffer

	// Referenced type ID.
	buffer.WriteByte('[')
	buffer.WriteString(string(typeNode.NodeId))
	buffer.WriteRune(']')

	// Nullable: 'T' or 'F'.
	if nullable {
		buffer.WriteRune('T')
	} else {
		buffer.WriteRune('F')
	}

	// Generic count and parameter count.
	buffer.WriteString(padNumberToString(len(generics), trhSlotGenericCount.length))
	buffer.WriteString(padNumberToString(0, trhSlotParameterCount.length))

	if len(generics) == 0 {
		return buffer.String()
	}

	// For each generic, add the length of its value string and then the value string itself.
	for _, generic := range generics {
		buffer.WriteRune('G')
		buffer.WriteString(padNumberToString(len(generic.value), typeRefValueSubReferenceLength))
		buffer.WriteString(generic.value)
	}

	return buffer.String()
}

// padNumberToString converts the given number into a string, padding it to ensure it has the
// specified number of digits.
func padNumberToString(number int, digits int) string {
	numberString := strconv.Itoa(number)
	return strings.Repeat("0", digits-len(numberString)) + numberString
}
<commit_msg>Type references are fully ascii<commit_after>// Copyright 2015 The Serulian Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package typegraph

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"

	"github.com/serulian/compiler/compilergraph"
)

// Value string example:
//
// [9242167c-2d57-4212-8aea-fedf32bd708e]T00010000G000039[04f97d44-d8fc-4a7d-9c46-955d1bd5add6]F
//  ^ Type ID of SomeType                ^ Nullable
//                                                ^ G000039 = generic with 39 chars in length for subreference
//                                        ^ 0001 = 1 generic
//                                            ^ 0000 = 0 parameters
// represents
//
// SomeType<SomeGeneric>?
//
// where SomeType is 9242167c-2d57-4212-8aea-fedf32bd708e
// and SomeGeneric is 04f97d44-d8fc-4a7d-9c46-955d1bd5add6
type typeReferenceHeaderSlot struct {
	index     int  // The slot index.
	length    int  // The length of this slot.
	writeable bool // Whether this slot is writeable.
}

var (
	trhSlotStartTypeId        = typeReferenceHeaderSlot{0, 1, false}
	trhSlotTypeId             = typeReferenceHeaderSlot{1, compilergraph.NodeIDLength, true}
	trhSlotEndTypeId          = typeReferenceHeaderSlot{2, 1, false}
	trhSlotFlagNullable       = typeReferenceHeaderSlot{3, 1, true}
	trhSlotGenericCount       = typeReferenceHeaderSlot{4, 4, true}
	trhSlotParameterCount     = typeReferenceHeaderSlot{5, 4, true}
	trhSlotSubReferenceMarker = typeReferenceHeaderSlot{6, 0, false}
)

var typeReferenceHeaderSlots = [...]typeReferenceHeaderSlot{
	trhSlotStartTypeId,
	trhSlotTypeId,
	trhSlotEndTypeId,
	trhSlotFlagNullable,
	trhSlotGenericCount,
	trhSlotParameterCount,
	trhSlotSubReferenceMarker,
}

// The size of the length prefix for subreferences.
const typeRefValueSubReferenceLength = 6

// subReferenceKind represents the kinds of supported subreferences
type subReferenceKind int

const (
	subReferenceGeneric subReferenceKind = iota
	subReferenceParameter
)

// withFlag returns a copy of this type reference with the flag at the given slot replaced with the
// specified rune.
func (tr *TypeReference) withFlag(flagSlot typeReferenceHeaderSlot, value rune) TypeReference {
	return tr.replaceSlot(flagSlot, string(value))
}

// replaceSlot replaces the given slot with the specified value.
func (tr *TypeReference) replaceSlot(slot typeReferenceHeaderSlot, value string) TypeReference {
	if !slot.writeable {
		panic(fmt.Sprintf("Cannot write to slot %v", slot))
	}

	location := getSlotLocation(slot)
	size := slot.length

	return TypeReference{
		layer: tr.layer,
		value: tr.value[:location] + value + tr.value[location+size:],
	}
}

// getSlotAsInt returns the int found at the given slot.
func (tr *TypeReference) getSlotAsInt(slot typeReferenceHeaderSlot) int {
	strValue := tr.getSlot(slot)

	// Special common case.
	if strValue == "0000" {
		return 0
	}

	i, err := strconv.Atoi(strValue)
	if err != nil {
		panic(fmt.Sprintf("Expected int value for slot %v, found: %v", slot, strValue))
	}
	return i
}

// getSlot returns the string found at the given slot.
func (tr *TypeReference) getSlot(slot typeReferenceHeaderSlot) string {
	location := getSlotLocation(slot)
	size := slot.length

	return tr.value[location : location+size]
}

// lengthPrefixedValue returns this type reference's value string, prefixed with its length value
// padded to the correct number of digits.
func (tr *TypeReference) lengthPrefixedValue() string {
	return padNumberToString(len(tr.value), typeRefValueSubReferenceLength) + tr.value
}

// withSubReference returns a copy of this type reference with the given subreference added.
func (tr *TypeReference) withSubReference(kind subReferenceKind, subref TypeReference) TypeReference {
	slot, kindRune := getSubReferenceSlotAndChar(kind)
	count := tr.getSlotAsInt(slot)

	// Add to the count and append the subreference.
	updated := tr.replaceSlot(slot, padNumberToString(count+1, slot.length))
	return TypeReference{
		layer: tr.layer,
		value: updated.value + string(kindRune) + subref.lengthPrefixedValue(),
	}
}

// getSubReferences returns all sub references of the given kind on this type reference.
func (tr *TypeReference) getSubReferences(kind subReferenceKind) []TypeReference {
	slot, kindChar := getSubReferenceSlotAndChar(kind)

	// Find the number of applicable subreferences.
	count := tr.getSlotAsInt(slot)
	subrefs := make([]TypeReference, count)

	if count == 0 {
		return subrefs
	}

	// For all subreferences loop and filter.
	var currentIndex int = getSlotLocation(trhSlotSubReferenceMarker)
	var collectedCount int = 0

	for {
		// Retrieve the character for the subreference (G or P)
		subReferenceChar := tr.value[currentIndex]

		// Retrieve the length of the subreference's value string.
		subReferenceLengthStr := tr.value[currentIndex+1 : currentIndex+1+typeRefValueSubReferenceLength]
		subReferenceLength, err := strconv.Atoi(subReferenceLengthStr)
		if err != nil {
			panic(fmt.Sprintf("Expected int value for subreference length, found: %v", subReferenceLengthStr))
		}

		// Move the current index forward to the point at which the subreference value string beings.
		currentIndex = currentIndex + 1 + typeRefValueSubReferenceLength
		subReferenceValue := tr.value[currentIndex : currentIndex+subReferenceLength]

		// If we have found a subreference of the known kind, add it.
		if subReferenceChar == kindChar {
			subrefs[collectedCount] = TypeReference{
				layer: tr.layer,
				value: subReferenceValue,
			}

			collectedCount = collectedCount + 1
		}

		// Once we've found all the subreferences of the expected kind, we're done.
		if collectedCount == count {
			return subrefs
		}

		// Move the current index past the subreference value string.
		currentIndex = currentIndex + subReferenceLength
	}
}

// getSubReferenceSlotAndChar returns the count slot and character for the given kind of subreference.
func getSubReferenceSlotAndChar(kind subReferenceKind) (typeReferenceHeaderSlot, uint8) {
	switch kind {
	case subReferenceGeneric:
		return trhSlotGenericCount, 'G'

	case subReferenceParameter:
		return trhSlotParameterCount, 'P'

	default:
		panic("Unknown kind of subreference")
		return trhSlotGenericCount, '_'
	}
}

// getSlotLocation returns the slot location (0-indexed) in the value string.
func getSlotLocation(slot typeReferenceHeaderSlot) int {
	// TODO(jschorr): Add a map here to be a cache?
	var location int
	for _, currentSlot := range typeReferenceHeaderSlots {
		if currentSlot == slot {
			return location
		}

		location = location + currentSlot.length
	}

	panic("Could not retrieve location for slot")
	return -1
}

// buildTypeReferenceValue returns a string value for representing the given type reference data in a single
// string.
func buildTypeReferenceValue(typeNode compilergraph.GraphNode, nullable bool, generics ...TypeReference) string {
	var buffer bytes.Buffer

	// Referenced type ID.
	buffer.WriteByte('[')
	buffer.WriteString(string(typeNode.NodeId))
	buffer.WriteByte(']')

	// Nullable: 'T' or 'F'.
	if nullable {
		buffer.WriteByte('T')
	} else {
		buffer.WriteByte('F')
	}

	// Generic count and parameter count.
	buffer.WriteString(padNumberToString(len(generics), trhSlotGenericCount.length))
	buffer.WriteString(padNumberToString(0, trhSlotParameterCount.length))

	if len(generics) == 0 {
		return buffer.String()
	}

	// For each generic, add the length of its value string and then the value string itself.
	for _, generic := range generics {
		buffer.WriteByte('G')
		buffer.WriteString(padNumberToString(len(generic.value), typeRefValueSubReferenceLength))
		buffer.WriteString(generic.value)
	}

	return buffer.String()
}

// padNumberToString converts the given number into a string, padding it to ensure it has the
// specified number of digits.
func padNumberToString(number int, digits int) string {
	numberString := strconv.Itoa(number)
	return strings.Repeat("0", digits-len(numberString)) + numberString
}
<|endoftext|>


<commit_before>package middleware

// Ported from Goji's middleware, source:
// https://github.com/zenazn/goji/tree/master/web/middleware

import (
	"bytes"
	"fmt"
	"os"
)

var (
	// Normal colors
	nBlack   = []byte{'\033', '[', '3', '0', 'm'}
	nRed     = []byte{'\033', '[', '3', '1', 'm'}
	nGreen   = []byte{'\033', '[', '3', '2', 'm'}
	nYellow  = []byte{'\033', '[', '3', '3', 'm'}
	nBlue    = []byte{'\033', '[', '3', '4', 'm'}
	nMagenta = []byte{'\033', '[', '3', '5', 'm'}
	nCyan    = []byte{'\033', '[', '3', '6', 'm'}
	nWhite   = []byte{'\033', '[', '3', '7', 'm'}
	// Bright colors
	bBlack   = []byte{'\033', '[', '3', '0', ';', '1', 'm'}
	bRed     = []byte{'\033', '[', '3', '1', ';', '1', 'm'}
	bGreen   = []byte{'\033', '[', '3', '2', ';', '1', 'm'}
	bYellow  = []byte{'\033', '[', '3', '3', ';', '1', 'm'}
	bBlue    = []byte{'\033', '[', '3', '4', ';', '1', 'm'}
	bMagenta = []byte{'\033', '[', '3', '5', ';', '1', 'm'}
	bCyan    = []byte{'\033', '[', '3', '6', ';', '1', 'm'}
	bWhite   = []byte{'\033', '[', '3', '7', ';', '1', 'm'}

	reset = []byte{'\033', '[', '0', 'm'}
)

var isTTY bool

func init() {
	// This is sort of cheating: if stdout is a character device, we assume
	// that means it's a TTY. Unfortunately, there are many non-TTY
	// character devices, but fortunately stdout is rarely set to any of
	// them.
	//
	// We could solve this properly by pulling in a dependency on
	// code.google.com/p/go.crypto/ssh/terminal, for instance, but as a
	// heuristic for whether to print in color or in black-and-white, I'd
	// really rather not.
	fi, err := os.Stdout.Stat()
	if err == nil {
		m := os.ModeDevice | os.ModeCharDevice
		isTTY = fi.Mode()&m == m
	}
}

// colorWrite
func cW(buf *bytes.Buffer, color []byte, s string, args ...interface{}) {
	if isTTY {
		buf.Write(color)
	}
	fmt.Fprintf(buf, s, args...)
	if isTTY {
		buf.Write(reset)
	}
}
<commit_msg>middleware: replace bytes.Buffer with io.Writer<commit_after>package middleware

// Ported from Goji's middleware, source:
// https://github.com/zenazn/goji/tree/master/web/middleware

import (
	"fmt"
	"io"
	"os"
)

var (
	// Normal colors
	nBlack   = []byte{'\033', '[', '3', '0', 'm'}
	nRed     = []byte{'\033', '[', '3', '1', 'm'}
	nGreen   = []byte{'\033', '[', '3', '2', 'm'}
	nYellow  = []byte{'\033', '[', '3', '3', 'm'}
	nBlue    = []byte{'\033', '[', '3', '4', 'm'}
	nMagenta = []byte{'\033', '[', '3', '5', 'm'}
	nCyan    = []byte{'\033', '[', '3', '6', 'm'}
	nWhite   = []byte{'\033', '[', '3', '7', 'm'}
	// Bright colors
	bBlack   = []byte{'\033', '[', '3', '0', ';', '1', 'm'}
	bRed     = []byte{'\033', '[', '3', '1', ';', '1', 'm'}
	bGreen   = []byte{'\033', '[', '3', '2', ';', '1', 'm'}
	bYellow  = []byte{'\033', '[', '3', '3', ';', '1', 'm'}
	bBlue    = []byte{'\033', '[', '3', '4', ';', '1', 'm'}
	bMagenta = []byte{'\033', '[', '3', '5', ';', '1', 'm'}
	bCyan    = []byte{'\033', '[', '3', '6', ';', '1', 'm'}
	bWhite   = []byte{'\033', '[', '3', '7', ';', '1', 'm'}

	reset = []byte{'\033', '[', '0', 'm'}
)

var isTTY bool

func init() {
	// This is sort of cheating: if stdout is a character device, we assume
	// that means it's a TTY. Unfortunately, there are many non-TTY
	// character devices, but fortunately stdout is rarely set to any of
	// them.
	//
	// We could solve this properly by pulling in a dependency on
	// code.google.com/p/go.crypto/ssh/terminal, for instance, but as a
	// heuristic for whether to print in color or in black-and-white, I'd
	// really rather not.
	fi, err := os.Stdout.Stat()
	if err == nil {
		m := os.ModeDevice | os.ModeCharDevice
		isTTY = fi.Mode()&m == m
	}
}

// colorWrite
func cW(w io.Writer, color []byte, s string, args ...interface{}) {
	if isTTY {
		w.Write(color)
	}
	fmt.Fprintf(w, s, args...)
	if isTTY {
		w.Write(reset)
	}
}
<|endoftext|>


<commit_before>package main

import "errors"

var DijkstraMapCache [DungeonNCells]int

func (g *game) Autoexplore(ev event) error {
	if mons := g.MonsterInLOS(); mons.Exists() {
		return errors.New("You cannot auto-explore while there are monsters in view.")
	}
	if g.ExclusionsMap[g.Player.Pos] {
		return errors.New("You cannot auto-explore while in an excluded area.")
	}
	if g.AllExplored() {
		return errors.New("Nothing left to explore.")
	}
	sources := g.AutoexploreSources()
	if len(sources) == 0 {
		return errors.New("Some excluded places remain unexplored.")
	}
	g.BuildAutoexploreMap(sources)
	n, finished := g.NextAuto()
	if finished || n == nil {
		return errors.New("You cannot reach safely some places.")
	}
	g.Autoexploring = true
	g.AutoHalt = false
	return g.MovePlayer(*n, ev)
}

func (g *game) AllExplored() bool {
	np := &normalPath{game: g}
	for i, c := range g.Dungeon.Cells {
		pos := idxtopos(i)
		if c.T == WallCell {
			if len(np.Neighbors(pos)) == 0 {
				continue
			}
		}
		_, okc := g.Collectables[pos]
		if !c.Explored || g.Simellas[pos] > 0 || okc {
			return false
		} else if _, ok := g.Rods[pos]; ok {
			return false
		}
	}
	return true
}

func (g *game) AutoexploreSources() []int {
	sources := []int{}
	np := &normalPath{game: g}
	for i, c := range g.Dungeon.Cells {
		pos := idxtopos(i)
		if c.T == WallCell {
			if len(np.Neighbors(pos)) == 0 {
				continue
			}
		}
		if g.ExclusionsMap[pos] {
			continue
		}
		_, okc := g.Collectables[pos]
		if !c.Explored || g.Simellas[pos] > 0 || okc {
			sources = append(sources, i)
		} else if _, ok := g.Rods[pos]; ok {
			sources = append(sources, i)
		}

	}
	return sources
}

func (g *game) BuildAutoexploreMap(sources []int) {
	ap := &autoexplorePath{game: g}
	g.AutoExploreDijkstra(ap, sources)
	g.DijkstraMapRebuild = false
}

func (g *game) NextAuto() (next *position, finished bool) {
	ap := &autoexplorePath{game: g}
	if DijkstraMapCache[g.Player.Pos.idx()] == unreachable {
		return nil, false
	}
	neighbors := ap.Neighbors(g.Player.Pos)
	if len(neighbors) == 0 {
		return nil, false
	}
	n := neighbors[0]
	ncost := DijkstraMapCache[n.idx()]
	for _, pos := range neighbors[1:] {
		cost := DijkstraMapCache[pos.idx()]
		if cost < ncost {
			n = pos
			ncost = cost
		}
	}
	if ncost >= DijkstraMapCache[g.Player.Pos.idx()] {
		finished = true
	}
	next = &n
	return next, finished
}
<commit_msg>Changed a silly? error message<commit_after>package main

import "errors"

var DijkstraMapCache [DungeonNCells]int

func (g *game) Autoexplore(ev event) error {
	if mons := g.MonsterInLOS(); mons.Exists() {
		return errors.New("You cannot auto-explore while there are monsters in view.")
	}
	if g.ExclusionsMap[g.Player.Pos] {
		return errors.New("You cannot auto-explore while in an excluded area.")
	}
	if g.AllExplored() {
		return errors.New("Nothing left to explore.")
	}
	sources := g.AutoexploreSources()
	if len(sources) == 0 {
		return errors.New("Some excluded places remain unexplored.")
	}
	g.BuildAutoexploreMap(sources)
	n, finished := g.NextAuto()
	if finished || n == nil {
		return errors.New("You cannot reach some places safely.")
	}
	g.Autoexploring = true
	g.AutoHalt = false
	return g.MovePlayer(*n, ev)
}

func (g *game) AllExplored() bool {
	np := &normalPath{game: g}
	for i, c := range g.Dungeon.Cells {
		pos := idxtopos(i)
		if c.T == WallCell {
			if len(np.Neighbors(pos)) == 0 {
				continue
			}
		}
		_, okc := g.Collectables[pos]
		if !c.Explored || g.Simellas[pos] > 0 || okc {
			return false
		} else if _, ok := g.Rods[pos]; ok {
			return false
		}
	}
	return true
}

func (g *game) AutoexploreSources() []int {
	sources := []int{}
	np := &normalPath{game: g}
	for i, c := range g.Dungeon.Cells {
		pos := idxtopos(i)
		if c.T == WallCell {
			if len(np.Neighbors(pos)) == 0 {
				continue
			}
		}
		if g.ExclusionsMap[pos] {
			continue
		}
		_, okc := g.Collectables[pos]
		if !c.Explored || g.Simellas[pos] > 0 || okc {
			sources = append(sources, i)
		} else if _, ok := g.Rods[pos]; ok {
			sources = append(sources, i)
		}

	}
	return sources
}

func (g *game) BuildAutoexploreMap(sources []int) {
	ap := &autoexplorePath{game: g}
	g.AutoExploreDijkstra(ap, sources)
	g.DijkstraMapRebuild = false
}

func (g *game) NextAuto() (next *position, finished bool) {
	ap := &autoexplorePath{game: g}
	if DijkstraMapCache[g.Player.Pos.idx()] == unreachable {
		return nil, false
	}
	neighbors := ap.Neighbors(g.Player.Pos)
	if len(neighbors) == 0 {
		return nil, false
	}
	n := neighbors[0]
	ncost := DijkstraMapCache[n.idx()]
	for _, pos := range neighbors[1:] {
		cost := DijkstraMapCache[pos.idx()]
		if cost < ncost {
			n = pos
			ncost = cost
		}
	}
	if ncost >= DijkstraMapCache[g.Player.Pos.idx()] {
		finished = true
	}
	next = &n
	return next, finished
}
<|endoftext|>


<commit_before>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"os"
	"path"
	"syscall"
	"time"
)

type state struct {
	subInodeToRequiredInode map[uint64]uint64
}

func (sub *Sub) buildUpdateRequest(request *subproto.UpdateRequest) {
	fmt.Println("buildUpdateRequest()") // TODO(rgooch): Delete debugging.
	subFS := sub.fileSystem
	requiredImage := sub.herd.getImage(sub.requiredImage)
	requiredFS := requiredImage.FileSystem
	filter := requiredImage.Filter
	request.Triggers = requiredImage.Triggers
	var state state
	state.subInodeToRequiredInode = make(map[uint64]uint64)
	var rusageStart, rusageStop syscall.Rusage
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStart)
	compareDirectories(request, &state,
		&subFS.DirectoryInode, &requiredFS.DirectoryInode,
		"/", filter)
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStop) // HACK
	cpuTime := time.Duration(rusageStop.Utime.Sec)*time.Second +
		time.Duration(rusageStop.Utime.Usec)*time.Microsecond -
		time.Duration(rusageStart.Utime.Sec)*time.Second -
		time.Duration(rusageStart.Utime.Usec)*time.Microsecond
	fmt.Printf("Build update request took: %s user CPU time\n", cpuTime)
}

func compareDirectories(request *subproto.UpdateRequest, state *state,
	subDirectory, requiredDirectory *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	// First look for entries that should be deleted.
	if subDirectory != nil {
		for name := range subDirectory.EntriesByName {
			pathname := path.Join(myPathName, name)
			if filter.Match(pathname) {
				continue
			}
			if _, ok := requiredDirectory.EntriesByName[name]; !ok {
				request.PathsToDelete = append(request.PathsToDelete, pathname)
				fmt.Printf("Delete: %s\n", pathname) // HACK
			}
		}
	}
	for name, requiredEntry := range requiredDirectory.EntriesByName {
		pathname := path.Join(myPathName, name)
		if filter.Match(pathname) {
			continue
		}
		var subEntry *filesystem.DirectoryEntry
		if subDirectory != nil {
			if se, ok := subDirectory.EntriesByName[name]; ok {
				subEntry = se
			}
		}
		if subEntry == nil {
			addEntry(request, state, requiredEntry, pathname)
		} else {
			compareEntries(request, state, subEntry, requiredEntry, pathname,
				filter)
		}
		// If a directory: descend (possibly with the directory for the sub).
		requiredInode := requiredEntry.Inode()
		if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
			var subInode *filesystem.DirectoryInode
			if subEntry != nil {
				if si, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
					subInode = si
				}
			}
			compareDirectories(request, state, subInode, requiredInode,
				pathname, filter)
		}
	}
}

func addEntry(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		makeDirectory(request, requiredInode, myPathName, true)
		fmt.Printf("Add directory: %s...\n", myPathName) // HACK
	} else {
		fmt.Printf("Add entry: %s...\n", myPathName) // HACK
		// TODO(rgooch): Add entry.
	}
}

func makeDirectory(request *subproto.UpdateRequest,
	requiredInode *filesystem.DirectoryInode, pathName string, create bool) {
	var newdir subproto.Directory
	newdir.Name = pathName
	newdir.Mode = requiredInode.Mode
	newdir.Uid = requiredInode.Uid
	newdir.Gid = requiredInode.Gid
	if create {
		request.DirectoriesToMake = append(request.DirectoriesToMake, newdir)
	} else {
		request.DirectoriesToChange = append(request.DirectoriesToMake, newdir)
	}
}

func compareEntries(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry,
	myPathName string, filter *filter.Filter) {
	switch requiredInode := requiredEntry.Inode().(type) {
	case *filesystem.RegularInode:
		compareRegularFile(request, state, subEntry,
			requiredInode, requiredEntry.InodeNumber, myPathName)
		return
	case *filesystem.SymlinkInode:
		compareSymlink(request, state, subEntry,
			requiredInode, requiredEntry.InodeNumber, myPathName)
		return
	case *filesystem.Inode:
		compareFile(request, state, subEntry,
			requiredInode, requiredEntry.InodeNumber, myPathName)
		return
	case *filesystem.DirectoryInode:
		compareDirectory(request, state, subEntry, requiredInode, myPathName,
			filter)
		return
	default:
		panic("Unsupported entry type")
	}
}

func compareRegularFile(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry,
	requiredInode *filesystem.RegularInode, requiredInodeNumber uint64,
	myPathName string) {
	if subInode, ok := subEntry.Inode().(*filesystem.RegularInode); ok {
		if requiredInum, ok :=
			state.subInodeToRequiredInode[subEntry.InodeNumber]; ok {
			if requiredInum != requiredInodeNumber {
				//
				fmt.Printf("Different links: %s...\n", myPathName) // HACK
			}
		} else {
			state.subInodeToRequiredInode[subEntry.InodeNumber] =
				requiredInodeNumber
		}
		sameMetadata := filesystem.CompareRegularInodesMetadata(
			subInode, requiredInode, os.Stdout)
		sameData := filesystem.CompareRegularInodesData(subInode,
			requiredInode, os.Stdout)
		if sameMetadata && sameData {
			return
		}
		fmt.Printf("Different rfile: %s...\n", myPathName) // HACK
	} else {
		fmt.Printf("Delete+add rfile: %s...\n", myPathName) // HACK
	}
	// TODO(rgooch): Delete entry and replace.
}

func compareSymlink(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry,
	requiredInode *filesystem.SymlinkInode, requiredInodeNumber uint64,
	myPathName string) {
	if subInode, ok := subEntry.Inode().(*filesystem.SymlinkInode); ok {
		if requiredInum, ok :=
			state.subInodeToRequiredInode[subEntry.InodeNumber]; ok {
			if requiredInum != requiredInodeNumber {
				fmt.Printf("Different links: %s...\n", myPathName) // HACK
			}
		} else {
			state.subInodeToRequiredInode[subEntry.InodeNumber] =
				requiredInodeNumber
		}
		if filesystem.CompareSymlinkInodes(subInode, requiredInode, os.Stdout) {
			return
		}
		fmt.Printf("Different symlink: %s...\n", myPathName) // HACK
	} else {
		fmt.Printf("Add symlink: %s...\n", myPathName) // HACK
	}
	// TODO(rgooch): Delete entry and replace.
}

func compareFile(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry,
	requiredInode *filesystem.Inode, requiredInodeNumber uint64,
	myPathName string) {
	if subInode, ok := subEntry.Inode().(*filesystem.Inode); ok {
		if requiredInum, ok :=
			state.subInodeToRequiredInode[subEntry.InodeNumber]; ok {
			if requiredInum != requiredInodeNumber {
				fmt.Printf("Different links: %s...\n", myPathName) // HACK
			}
		} else {
			state.subInodeToRequiredInode[subEntry.InodeNumber] =
				requiredInodeNumber
		}
		if filesystem.CompareInodes(subInode, requiredInode, os.Stdout) {
			return
		}
		fmt.Printf("Different file: %s...\n", myPathName) // HACK
	} else {
		fmt.Printf("Add file: %s...\n", myPathName) // HACK
	}
	// TODO(rgooch): Delete entry and replace.
}

func compareDirectory(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry,
	requiredInode *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	if subInode, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
		if filesystem.CompareDirectoriesMetadata(subInode, requiredInode,
			os.Stdout) {
			return
		}
		makeDirectory(request, requiredInode, myPathName, false)
		fmt.Printf("Different directory: %s...\n", myPathName) // HACK
	} else {
		request.PathsToDelete = append(request.PathsToDelete, myPathName)
		makeDirectory(request, requiredInode, myPathName, true)
		fmt.Printf("Replace non-directory: %s...\n", myPathName) // HACK
	}
}
<commit_msg>Incremental work on buildUpdateRequest(): add relinking and metadata updates.<commit_after>package herd

import (
	"fmt"
	"github.com/Symantec/Dominator/lib/filesystem"
	"github.com/Symantec/Dominator/lib/filter"
	subproto "github.com/Symantec/Dominator/proto/sub"
	"os"
	"path"
	"syscall"
	"time"
)

type state struct {
	requiredInodeToSubInode map[uint64]uint64
	inodesChanged           map[uint64]bool // Required inode number.
	subFS                   *filesystem.FileSystem
	requiredFS              *filesystem.FileSystem
}

func (sub *Sub) buildUpdateRequest(request *subproto.UpdateRequest) {
	fmt.Println("buildUpdateRequest()") // TODO(rgooch): Delete debugging.
	var state state
	state.subFS = &sub.fileSystem.FileSystem
	requiredImage := sub.herd.getImage(sub.requiredImage)
	state.requiredFS = requiredImage.FileSystem
	filter := requiredImage.Filter
	request.Triggers = requiredImage.Triggers
	state.requiredInodeToSubInode = make(map[uint64]uint64)
	state.inodesChanged = make(map[uint64]bool)
	var rusageStart, rusageStop syscall.Rusage
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStart)
	compareDirectories(request, &state,
		&state.subFS.DirectoryInode, &state.requiredFS.DirectoryInode,
		"/", filter)
	syscall.Getrusage(syscall.RUSAGE_SELF, &rusageStop) // HACK
	cpuTime := time.Duration(rusageStop.Utime.Sec)*time.Second +
		time.Duration(rusageStop.Utime.Usec)*time.Microsecond -
		time.Duration(rusageStart.Utime.Sec)*time.Second -
		time.Duration(rusageStart.Utime.Usec)*time.Microsecond
	fmt.Printf("Build update request took: %s user CPU time\n", cpuTime)
}

func compareDirectories(request *subproto.UpdateRequest, state *state,
	subDirectory, requiredDirectory *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	// First look for entries that should be deleted.
	if subDirectory != nil {
		for name := range subDirectory.EntriesByName {
			pathname := path.Join(myPathName, name)
			if filter.Match(pathname) {
				continue
			}
			if _, ok := requiredDirectory.EntriesByName[name]; !ok {
				request.PathsToDelete = append(request.PathsToDelete, pathname)
				fmt.Printf("Delete: %s\n", pathname) // HACK
			}
		}
	}
	for name, requiredEntry := range requiredDirectory.EntriesByName {
		pathname := path.Join(myPathName, name)
		if filter.Match(pathname) {
			continue
		}
		var subEntry *filesystem.DirectoryEntry
		if subDirectory != nil {
			if se, ok := subDirectory.EntriesByName[name]; ok {
				subEntry = se
			}
		}
		if subEntry == nil {
			addEntry(request, state, requiredEntry, pathname)
		} else {
			compareEntries(request, state, subEntry, requiredEntry, pathname,
				filter)
		}
		// If a directory: descend (possibly with the directory for the sub).
		requiredInode := requiredEntry.Inode()
		if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
			var subInode *filesystem.DirectoryInode
			if subEntry != nil {
				if si, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
					subInode = si
				}
			}
			compareDirectories(request, state, subInode, requiredInode,
				pathname, filter)
		}
	}
}

func addEntry(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	requiredInode := requiredEntry.Inode()
	if requiredInode, ok := requiredInode.(*filesystem.DirectoryInode); ok {
		makeDirectory(request, requiredInode, myPathName, true)
	} else {
		addInode(request, state, requiredEntry, myPathName)
	}
}

func compareEntries(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry,
	myPathName string, filter *filter.Filter) {
	var sameType, sameMetadata, sameData bool
	switch requiredInode := requiredEntry.Inode().(type) {
	case *filesystem.RegularInode:
		sameType, sameMetadata, sameData =
			compareRegularFile(request, state, subEntry, requiredInode,
				myPathName)
	case *filesystem.SymlinkInode:
		sameType, sameMetadata, sameData =
			compareSymlink(request, state, subEntry, requiredInode, myPathName)
	case *filesystem.Inode:
		sameType, sameMetadata, sameData =
			compareFile(request, state, subEntry, requiredInode, myPathName)
	case *filesystem.DirectoryInode:
		compareDirectory(request, state, subEntry, requiredInode, myPathName,
			filter)
		return
	default:
		panic("Unsupported entry type")
	}
	if sameType && sameData && sameMetadata {
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	if sameType && sameData {
		updateMetadata(request, state, subEntry, requiredEntry, myPathName)
		relink(request, state, subEntry, requiredEntry, myPathName)
		return
	}
	request.PathsToDelete = append(request.PathsToDelete, myPathName)
	addInode(request, state, requiredEntry, myPathName)
}

func compareRegularFile(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry, requiredInode *filesystem.RegularInode,
	myPathName string) (sameType, sameMetadata, sameData bool) {
	if subInode, ok := subEntry.Inode().(*filesystem.RegularInode); ok {
		sameType = true
		sameMetadata = filesystem.CompareRegularInodesMetadata(
			subInode, requiredInode, nil)
		sameData = filesystem.CompareRegularInodesData(subInode,
			requiredInode, os.Stdout)
	}
	return
}

func compareSymlink(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry, requiredInode *filesystem.SymlinkInode,
	myPathName string) (sameType, sameMetadata, sameData bool) {
	if subInode, ok := subEntry.Inode().(*filesystem.SymlinkInode); ok {
		sameType = true
		sameMetadata = filesystem.CompareSymlinkInodesMetadata(subInode,
			requiredInode, nil)
		sameData = filesystem.CompareSymlinkInodesData(subInode, requiredInode,
			os.Stdout)
	}
	return
}

func compareFile(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry, requiredInode *filesystem.Inode,
	myPathName string) (sameType, sameMetadata, sameData bool) {
	if subInode, ok := subEntry.Inode().(*filesystem.Inode); ok {
		sameType = true
		sameMetadata = filesystem.CompareInodesMetadata(subInode, requiredInode,
			nil)
		sameData = filesystem.CompareInodesData(subInode, requiredInode,
			os.Stdout)
	}
	return
}

func compareDirectory(request *subproto.UpdateRequest, state *state,
	subEntry *filesystem.DirectoryEntry,
	requiredInode *filesystem.DirectoryInode,
	myPathName string, filter *filter.Filter) {
	if subInode, ok := subEntry.Inode().(*filesystem.DirectoryInode); ok {
		if filesystem.CompareDirectoriesMetadata(subInode, requiredInode, nil) {
			return
		}
		makeDirectory(request, requiredInode, myPathName, false)
	} else {
		request.PathsToDelete = append(request.PathsToDelete, myPathName)
		makeDirectory(request, requiredInode, myPathName, true)
		fmt.Printf("Replace non-directory: %s...\n", myPathName) // HACK
	}
}

func relink(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	subInum, ok := state.requiredInodeToSubInode[requiredEntry.InodeNumber]
	if !ok {
		state.requiredInodeToSubInode[requiredEntry.InodeNumber] =
			subEntry.InodeNumber
		return
	}
	if subInum == subEntry.InodeNumber {
		return
	}
	var hardlink subproto.Hardlink
	hardlink.Source = myPathName
	hardlink.Target = state.subFS.InodeToFilenamesTable[subInum][0]
	request.HardlinksToMake = append(request.HardlinksToMake, hardlink)
	fmt.Printf("Make link: %s => %s\n", hardlink.Source,
		hardlink.Target) // HACK
}

func updateMetadata(request *subproto.UpdateRequest, state *state,
	subEntry, requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	if changed := state.inodesChanged[requiredEntry.InodeNumber]; changed {
		return
	}
	var inode subproto.Inode
	inode.Name = myPathName
	inode.GenericInode = requiredEntry.Inode()
	request.InodesToChange = append(request.InodesToChange, inode)
	state.inodesChanged[requiredEntry.InodeNumber] = true
	fmt.Printf("Update metadata: %s\n", myPathName) // HACK
}

func makeDirectory(request *subproto.UpdateRequest,
	requiredInode *filesystem.DirectoryInode, pathName string, create bool) {
	var newdir subproto.Directory
	newdir.Name = pathName
	newdir.Mode = requiredInode.Mode
	newdir.Uid = requiredInode.Uid
	newdir.Gid = requiredInode.Gid
	if create {
		request.DirectoriesToMake = append(request.DirectoriesToMake, newdir)
		fmt.Printf("Add directory: %s...\n", pathName) // HACK
	} else {
		request.DirectoriesToChange = append(request.DirectoriesToMake, newdir)
		fmt.Printf("Change directory: %s...\n", pathName) // HACK
	}
}

func addInode(request *subproto.UpdateRequest, state *state,
	requiredEntry *filesystem.DirectoryEntry, myPathName string) {
	fmt.Printf("Add entry: %s...\n", myPathName) // HACK
	// TODO(rgooch): Add entry.
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/olekukonko/tablewriter"

	"gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"
)

type config struct {
	BaseURL  string             `yaml:"baseURL"`
	Formats  map[string]format  `yaml:"formats,flow"`
	Elements map[string]element `yaml:"elements,flow"`
}

type element struct {
	ID     string   `yaml:"id"`
	File   string   `yaml:"file"`
	Meta   bool     `yaml:"meta"`
	Name   string   `yaml:"name"`
	Files  []string `yaml:"files"`
	Parent string   `yaml:"parent"`
}

type format struct {
	ID  string `yaml:"ext"`
	Loc string `yaml:"loc"`
}

var (
	app        = kingpin.New("download-geofabrik", "A command-line tool for downloading OSM files.")
	Fconfig    = app.Flag("config", "Set Config file.").Default("./geofabrik.yml").Short('c').String()
	nodownload = app.Flag("nodownload", "Do not download file (test only)").Short('n').Bool()
	verbose    = app.Flag("verbose", "Be verbose").Short('v').Bool()

	update = app.Command("update", "Update geofabrik.yml from github")
	url    = update.Flag("url", "Url for config source").Default("https://raw.githubusercontent.com/julien-noblet/download-geofabrik/stable/geofabrik.yml").String()

	list = app.Command("list", "Show elements available")
	lmd  = list.Flag("markdown", "generate list in Markdown format").Bool()

	download = app.Command("download", "Download element") //TODO : add d as command
	delement = download.Arg("element", "OSM element").Required().String()
	dosmBz2  = download.Flag("osm.bz2", "Download osm.bz2 if available").Short('B').Bool()
	dshpZip  = download.Flag("shp.zip", "Download shp.zip if available").Short('S').Bool()
	dosmPbf  = download.Flag("osm.pbf", "Download osm.pbf (default)").Short('P').Bool()
	doshPbf  = download.Flag("osh.pbf", "Download osh.pbf (default)").Short('H').Bool()
	dstate   = download.Flag("state", "Download state.txt file").Short('s').Bool()
	dpoly    = download.Flag("poly", "Download poly file").Short('p').Bool()
)

func (e *element) hasParent() bool {
	return len(e.Parent) != 0
}

func miniFormats(s []string) string {
	res := make([]string, 6)
	for _, item := range s {
		switch item {
		case "state":
			res[0] = "s"
		case "osm.pbf":
			res[1] = "P"
		case "osm.bz2":
			res[2] = "B"
		case "osh.pbf":
			res[3] = "H"
		case "poly":
			res[4] = "p"
		case "shp.zip":
			res[5] = "S"
		}
	}

	return strings.Join(res, "")
}

func downloadFromURL(url string, fileName string) {
	if *verbose == true {
		log.Println(" Downloading", url, "to", fileName)
	}

	if *nodownload == false {
		// TODO: check file existence first with io.IsExist
		output, err := os.Create(fileName)
		if err != nil {
			log.Fatalln(" Error while creating ", fileName, "-", err)
			return
		}
		defer output.Close()

		response, err := http.Get(url)
		if err != nil {
			log.Fatalln(" Error while downloading ", url, "-", err)
			return
		}
		defer response.Body.Close()

		n, err := io.Copy(output, response.Body)
		if err != nil {
			log.Fatalln(" Error while downloading ", url, "-", err)
			return
		}

		if *verbose == true {
			log.Println(" ", n, "bytes downloaded.")
		}
	}
}
func elem2preURL(c config, e element) string {
	var res string
	if e.hasParent() {
		res = elem2preURL(c, findElem(c, e.Parent)) + "/"
		if e.File != "" { //TODO use file in config???
			res = res + e.File
		} else {
			res = res + e.ID
		}
	} else {
		res = c.BaseURL + "/" + e.ID
	}
	return res
}

func elem2URL(c config, e element, ext string) string {
	res := elem2preURL(c, e)
	res += c.Formats[ext].Loc
	if !stringInSlice(ext, e.Files) {
		log.Fatalln(" Error!!! " + res + " not exist")
	}

	return res
}

func findElem(c config, e string) element {
	res := c.Elements[e]
	if res.ID == "" {
		log.Fatalln(" " + e + " is not in config! Please use \"list\" command!")
	}
	return res
}
func getFormats() []string {
	var formatFile []string
	if *dosmPbf {
		formatFile = append(formatFile, "osm.pbf")
	}
	if *doshPbf {
		formatFile = append(formatFile, "osh.pbf")
	}
	if *dosmBz2 {
		formatFile = append(formatFile, "osm.bz2")
	}
	if *dshpZip {
		formatFile = append(formatFile, "shp.zip")
	}
	if *dstate {
		formatFile = append(formatFile, "state")
	}
	if *dpoly {
		formatFile = append(formatFile, "poly")
	}
	if len(formatFile) == 0 {
		formatFile = append(formatFile, "osm.pbf")
	}
	return formatFile
}

func listAllRegions(c config, format string) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetHeader([]string{"ShortName", "Is in", "Long Name", "formats"})
	if format == "Markdown" {
		table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false})
		table.SetCenterSeparator("|")
	}
	keys := make(sort.StringSlice, len(c.Elements))
	i := 0
	for k := range c.Elements {
		keys[i] = k
		i++
	}
	keys.Sort()
	for _, item := range keys {
		table.Append([]string{item, c.Elements[c.Elements[item].Parent].Name, c.Elements[item].Name, miniFormats(c.Elements[item].Files)})
	}
	table.Render()
	fmt.Printf("Total elements: %#v\n", len(c.Elements))
}

func loadConfig(configFile string) config {
	filename, _ := filepath.Abs(configFile)
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalln(" File error: %v ", err)
		os.Exit(1)
	}
	var myConfig config
	err = yaml.Unmarshal(file, &myConfig)
	if err != nil {
		panic(err)
	}
	return myConfig

}
func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

func UpdateConfig(url string, myconfig string) {
	downloadFromURL(url, myconfig)
	fmt.Println("Congratulation, you have the latest geofabrik.yml")
}

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	case list.FullCommand():
		var format = ""
		if *lmd {
			format = "Markdown"
		}
		listAllRegions(loadConfig(*Fconfig), format)
	case update.FullCommand():
		UpdateConfig(*url, *Fconfig)
	case download.FullCommand():
		formatFile := getFormats()
		for _, format := range formatFile {
			downloadFromURL(elem2URL(loadConfig(*Fconfig), findElem(loadConfig(*Fconfig), *delement), format), *delement+"."+format)
		}
	}
}
<commit_msg>fix: update url<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"github.com/olekukonko/tablewriter"

	"gopkg.in/alecthomas/kingpin.v2"
	"gopkg.in/yaml.v2"
)

type config struct {
	BaseURL  string             `yaml:"baseURL"`
	Formats  map[string]format  `yaml:"formats,flow"`
	Elements map[string]element `yaml:"elements,flow"`
}

type element struct {
	ID     string   `yaml:"id"`
	File   string   `yaml:"file"`
	Meta   bool     `yaml:"meta"`
	Name   string   `yaml:"name"`
	Files  []string `yaml:"files"`
	Parent string   `yaml:"parent"`
}

type format struct {
	ID  string `yaml:"ext"`
	Loc string `yaml:"loc"`
}

var (
	app        = kingpin.New("download-geofabrik", "A command-line tool for downloading OSM files.")
	Fconfig    = app.Flag("config", "Set Config file.").Default("./geofabrik.yml").Short('c').String()
	nodownload = app.Flag("nodownload", "Do not download file (test only)").Short('n').Bool()
	verbose    = app.Flag("verbose", "Be verbose").Short('v').Bool()

	update = app.Command("update", "Update geofabrik.yml from github")
	url    = update.Flag("url", "Url for config source").Default("https://raw.githubusercontent.com/julien-noblet/download-geofabrik/master/geofabrik.yml").String()

	list = app.Command("list", "Show elements available")
	lmd  = list.Flag("markdown", "generate list in Markdown format").Bool()

	download = app.Command("download", "Download element") //TODO : add d as command
	delement = download.Arg("element", "OSM element").Required().String()
	dosmBz2  = download.Flag("osm.bz2", "Download osm.bz2 if available").Short('B').Bool()
	dshpZip  = download.Flag("shp.zip", "Download shp.zip if available").Short('S').Bool()
	dosmPbf  = download.Flag("osm.pbf", "Download osm.pbf (default)").Short('P').Bool()
	doshPbf  = download.Flag("osh.pbf", "Download osh.pbf (default)").Short('H').Bool()
	dstate   = download.Flag("state", "Download state.txt file").Short('s').Bool()
	dpoly    = download.Flag("poly", "Download poly file").Short('p').Bool()
)

func (e *element) hasParent() bool {
	return len(e.Parent) != 0
}

func miniFormats(s []string) string {
	res := make([]string, 6)
	for _, item := range s {
		switch item {
		case "state":
			res[0] = "s"
		case "osm.pbf":
			res[1] = "P"
		case "osm.bz2":
			res[2] = "B"
		case "osh.pbf":
			res[3] = "H"
		case "poly":
			res[4] = "p"
		case "shp.zip":
			res[5] = "S"
		}
	}

	return strings.Join(res, "")
}

func downloadFromURL(url string, fileName string) {
	if *verbose == true {
		log.Println(" Downloading", url, "to", fileName)
	}

	if *nodownload == false {
		// TODO: check file existence first with io.IsExist
		output, err := os.Create(fileName)
		if err != nil {
			log.Fatalln(" Error while creating ", fileName, "-", err)
			return
		}
		defer output.Close()

		response, err := http.Get(url)
		if err != nil {
			log.Fatalln(" Error while downloading ", url, "-", err)
			return
		}
		defer response.Body.Close()

		n, err := io.Copy(output, response.Body)
		if err != nil {
			log.Fatalln(" Error while downloading ", url, "-", err)
			return
		}

		if *verbose == true {
			log.Println(" ", n, "bytes downloaded.")
		}
	}
}
func elem2preURL(c config, e element) string {
	var res string
	if e.hasParent() {
		res = elem2preURL(c, findElem(c, e.Parent)) + "/"
		if e.File != "" { //TODO use file in config???
			res = res + e.File
		} else {
			res = res + e.ID
		}
	} else {
		res = c.BaseURL + "/" + e.ID
	}
	return res
}

func elem2URL(c config, e element, ext string) string {
	res := elem2preURL(c, e)
	res += c.Formats[ext].Loc
	if !stringInSlice(ext, e.Files) {
		log.Fatalln(" Error!!! " + res + " not exist")
	}

	return res
}

func findElem(c config, e string) element {
	res := c.Elements[e]
	if res.ID == "" {
		log.Fatalln(" " + e + " is not in config! Please use \"list\" command!")
	}
	return res
}
func getFormats() []string {
	var formatFile []string
	if *dosmPbf {
		formatFile = append(formatFile, "osm.pbf")
	}
	if *doshPbf {
		formatFile = append(formatFile, "osh.pbf")
	}
	if *dosmBz2 {
		formatFile = append(formatFile, "osm.bz2")
	}
	if *dshpZip {
		formatFile = append(formatFile, "shp.zip")
	}
	if *dstate {
		formatFile = append(formatFile, "state")
	}
	if *dpoly {
		formatFile = append(formatFile, "poly")
	}
	if len(formatFile) == 0 {
		formatFile = append(formatFile, "osm.pbf")
	}
	return formatFile
}

func listAllRegions(c config, format string) {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetHeader([]string{"ShortName", "Is in", "Long Name", "formats"})
	if format == "Markdown" {
		table.SetBorders(tablewriter.Border{Left: true, Top: false, Right: true, Bottom: false})
		table.SetCenterSeparator("|")
	}
	keys := make(sort.StringSlice, len(c.Elements))
	i := 0
	for k := range c.Elements {
		keys[i] = k
		i++
	}
	keys.Sort()
	for _, item := range keys {
		table.Append([]string{item, c.Elements[c.Elements[item].Parent].Name, c.Elements[item].Name, miniFormats(c.Elements[item].Files)})
	}
	table.Render()
	fmt.Printf("Total elements: %#v\n", len(c.Elements))
}

func loadConfig(configFile string) config {
	filename, _ := filepath.Abs(configFile)
	file, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatalln(" File error: %v ", err)
		os.Exit(1)
	}
	var myConfig config
	err = yaml.Unmarshal(file, &myConfig)
	if err != nil {
		panic(err)
	}
	return myConfig

}
func stringInSlice(a string, list []string) bool {
	for _, b := range list {
		if b == a {
			return true
		}
	}
	return false
}

func UpdateConfig(url string, myconfig string) {
	downloadFromURL(url, myconfig)
	fmt.Println("Congratulation, you have the latest geofabrik.yml")
}

func main() {
	switch kingpin.MustParse(app.Parse(os.Args[1:])) {

	case list.FullCommand():
		var format = ""
		if *lmd {
			format = "Markdown"
		}
		listAllRegions(loadConfig(*Fconfig), format)
	case update.FullCommand():
		UpdateConfig(*url, *Fconfig)
	case download.FullCommand():
		formatFile := getFormats()
		for _, format := range formatFile {
			downloadFromURL(elem2URL(loadConfig(*Fconfig), findElem(loadConfig(*Fconfig), *delement), format), *delement+"."+format)
		}
	}
}
<|endoftext|>