<commit_before>// Copyright 2015 go-swagger maintainers
//
// 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 runtime

import (
	"io"
	"io/ioutil"
	"net/url"
	"time"

	"github.com/go-openapi/strfmt"
)

// ClientRequestWriterFunc converts a function to a request writer interface
type ClientRequestWriterFunc func(ClientRequest, strfmt.Registry) error

// WriteToRequest adds data to the request
func (fn ClientRequestWriterFunc) WriteToRequest(req ClientRequest, reg strfmt.Registry) error {
	return fn(req, reg)
}

// ClientRequestWriter is an interface for things that know how to write to a request
type ClientRequestWriter interface {
	WriteToRequest(ClientRequest, strfmt.Registry) error
}

// ClientRequest is an interface for things that know how to
// add information to a swagger client request
type ClientRequest interface {
	SetHeaderParam(string, ...string) error

	SetQueryParam(string, ...string) error

	SetFormParam(string, ...string) error

	SetPathParam(string, string) error

	GetQueryParams() url.Values

	SetFileParam(string, ...NamedReadCloser) error

	SetBodyParam(interface{}) error

	SetTimeout(time.Duration) error

	GetMethod() string

	GetPath() string

	GetBody() []byte
}

// NamedReadCloser represents a named ReadCloser interface
type NamedReadCloser interface {
	io.ReadCloser
	Name() string
}

func NamedReader(name string, rdr io.Reader) NamedReadCloser {
	rc, ok := rdr.(io.ReadCloser)
	if !ok {
		rc = ioutil.NopCloser(rdr)
	}
	return &namedReadCloser{
		name: name,
		cr:   rc,
	}
}

type namedReadCloser struct {
	name string
	cr   io.ReadCloser
}

func (n *namedReadCloser) Close() error {
	return n.cr.Close()
}
func (n *namedReadCloser) Read(p []byte) (int, error) {
	return n.cr.Read(p)
}
func (n *namedReadCloser) Name() string {
	return n.name
}
<commit_msg>update doc comment<commit_after>// Copyright 2015 go-swagger maintainers
//
// 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 runtime

import (
	"io"
	"io/ioutil"
	"net/url"
	"time"

	"github.com/go-openapi/strfmt"
)

// ClientRequestWriterFunc converts a function to a request writer interface
type ClientRequestWriterFunc func(ClientRequest, strfmt.Registry) error

// WriteToRequest adds data to the request
func (fn ClientRequestWriterFunc) WriteToRequest(req ClientRequest, reg strfmt.Registry) error {
	return fn(req, reg)
}

// ClientRequestWriter is an interface for things that know how to write to a request
type ClientRequestWriter interface {
	WriteToRequest(ClientRequest, strfmt.Registry) error
}

// ClientRequest is an interface for things that know how to
// add information to a swagger client request
type ClientRequest interface {
	SetHeaderParam(string, ...string) error

	SetQueryParam(string, ...string) error

	SetFormParam(string, ...string) error

	SetPathParam(string, string) error

	GetQueryParams() url.Values

	SetFileParam(string, ...NamedReadCloser) error

	SetBodyParam(interface{}) error

	SetTimeout(time.Duration) error

	GetMethod() string

	GetPath() string

	GetBody() []byte
}

// NamedReadCloser represents a named ReadCloser interface
type NamedReadCloser interface {
	io.ReadCloser
	Name() string
}

// NamedReader creates a NamedReadCloser for use as file upload
func NamedReader(name string, rdr io.Reader) NamedReadCloser {
	rc, ok := rdr.(io.ReadCloser)
	if !ok {
		rc = ioutil.NopCloser(rdr)
	}
	return &namedReadCloser{
		name: name,
		cr:   rc,
	}
}

type namedReadCloser struct {
	name string
	cr   io.ReadCloser
}

func (n *namedReadCloser) Close() error {
	return n.cr.Close()
}
func (n *namedReadCloser) Read(p []byte) (int, error) {
	return n.cr.Read(p)
}
func (n *namedReadCloser) Name() string {
	return n.name
}
<|endoftext|>


<commit_before>// Copyright The OpenTelemetry 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.

// nolint:errcheck
package filelogreceiver

import (
	"context"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/observiq/nanojack"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/collector/component/componenttest"
	"go.opentelemetry.io/collector/config"
	"go.opentelemetry.io/collector/config/configtest"
	"go.opentelemetry.io/collector/consumer/consumertest"
	"go.opentelemetry.io/collector/pdata/plog"
	"go.opentelemetry.io/collector/service/servicetest"

	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/adapter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/entry"
)

func TestDefaultConfig(t *testing.T) {
	factory := NewFactory()
	cfg := factory.CreateDefaultConfig()
	require.NotNil(t, cfg, "failed to create default config")
	require.NoError(t, configtest.CheckConfigStruct(cfg))
}

func TestLoadConfig(t *testing.T) {
	factories, err := componenttest.NopFactories()
	assert.Nil(t, err)

	factory := NewFactory()
	factories.Receivers[typeStr] = factory
	cfg, err := servicetest.LoadConfigAndValidate(filepath.Join("testdata", "config.yaml"), factories)
	require.NoError(t, err)
	require.NotNil(t, cfg)

	assert.Equal(t, len(cfg.Receivers), 1)

	assert.Equal(t, testdataConfigYamlAsMap(), cfg.Receivers[config.NewComponentID("filelog")])
}

func TestCreateWithInvalidInputConfig(t *testing.T) {
	t.Parallel()

	cfg := testdataConfigYamlAsMap()
	cfg.Input["include"] = "not an array"

	_, err := NewFactory().CreateLogsReceiver(
		context.Background(),
		componenttest.NewNopReceiverCreateSettings(),
		cfg,
		new(consumertest.LogsSink),
	)
	require.Error(t, err, "receiver creation should fail if given invalid input config")
}

func TestReadStaticFile(t *testing.T) {
	t.Parallel()

	expectedTimestamp, _ := time.ParseInLocation("2006-01-02", "2020-08-25", time.Local)

	f := NewFactory()
	sink := new(consumertest.LogsSink)

	cfg := testdataConfigYamlAsMap()
	cfg.Converter.MaxFlushCount = 10
	cfg.Converter.FlushInterval = time.Millisecond

	converter := adapter.NewConverter()
	converter.Start()
	defer converter.Stop()

	var wg sync.WaitGroup
	wg.Add(1)
	go consumeNLogsFromConverter(converter.OutChannel(), 3, &wg)

	rcvr, err := f.CreateLogsReceiver(context.Background(), componenttest.NewNopReceiverCreateSettings(), cfg, sink)
	require.NoError(t, err, "failed to create receiver")
	require.NoError(t, rcvr.Start(context.Background(), componenttest.NewNopHost()))

	// Build the expected set by using adapter.Converter to translate entries
	// to pdata Logs.
	queueEntry := func(t *testing.T, c *adapter.Converter, msg string, severity entry.Severity) {
		e := entry.New()
		e.Timestamp = expectedTimestamp
		e.Set(entry.NewBodyField("msg"), msg)
		e.Severity = severity
		e.AddAttribute("file_name", "simple.log")
		require.NoError(t, c.Batch([]*entry.Entry{e}))
	}
	queueEntry(t, converter, "Something routine", entry.Info)
	queueEntry(t, converter, "Something bad happened!", entry.Error)
	queueEntry(t, converter, "Some details...", entry.Debug)

	dir, err := os.Getwd()
	require.NoError(t, err)
	t.Logf("Working Directory: %s", dir)

	wg.Wait()

	require.Eventually(t, expectNLogs(sink, 3), 2*time.Second, 5*time.Millisecond,
		"expected %d but got %d logs",
		3, sink.LogRecordCount(),
	)
	// TODO: Figure out a nice way to assert each logs entry content.
	// require.Equal(t, expectedLogs, sink.AllLogs())
	require.NoError(t, rcvr.Shutdown(context.Background()))
}

func TestReadRotatingFiles(t *testing.T) {

	tests := []rotationTest{
		{
			name:         "CopyTruncateTimestamped",
			copyTruncate: true,
			sequential:   false,
		},
		{
			name:         "MoveCreateTimestamped",
			copyTruncate: false,
			sequential:   false,
		},
		{
			name:         "CopyTruncateSequential",
			copyTruncate: true,
			sequential:   true,
		},
		{
			name:         "MoveCreateSequential",
			copyTruncate: false,
			sequential:   true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, tc.Run)
	}
}

type rotationTest struct {
	name         string
	copyTruncate bool
	sequential   bool
}

func (rt *rotationTest) Run(t *testing.T) {
	t.Parallel()

	tempDir := t.TempDir()

	f := NewFactory()
	sink := new(consumertest.LogsSink)

	cfg := testdataRotateTestYamlAsMap(tempDir)
	cfg.Converter.MaxFlushCount = 1
	cfg.Converter.FlushInterval = time.Millisecond

	// With a max of 100 logs per file and 1 backup file, rotation will occur
	// when more than 100 logs are written, and deletion when more than 200 are written.
	// Write 300 and validate that we got the all despite rotation and deletion.
	logger := newRotatingLogger(t, tempDir, 100, 1, rt.copyTruncate, rt.sequential)
	numLogs := 300

	// Build expected outputs
	expectedTimestamp, _ := time.ParseInLocation("2006-01-02", "2020-08-25", time.Local)
	converter := adapter.NewConverter()
	converter.Start()

	var wg sync.WaitGroup
	wg.Add(1)
	go consumeNLogsFromConverter(converter.OutChannel(), numLogs, &wg)

	rcvr, err := f.CreateLogsReceiver(context.Background(), componenttest.NewNopReceiverCreateSettings(), cfg, sink)
	require.NoError(t, err, "failed to create receiver")
	require.NoError(t, rcvr.Start(context.Background(), componenttest.NewNopHost()))

	for i := 0; i < numLogs; i++ {
		msg := fmt.Sprintf("This is a simple log line with the number %3d", i)

		// Build the expected set by converting entries to pdata Logs...
		e := entry.New()
		e.Timestamp = expectedTimestamp
		e.Set(entry.NewBodyField("msg"), msg)
		require.NoError(t, converter.Batch([]*entry.Entry{e}))

		// ... and write the logs lines to the actual file consumed by receiver.
		logger.Printf("2020-08-25 %s", msg)
		time.Sleep(time.Millisecond)
	}

	wg.Wait()
	require.Eventually(t, expectNLogs(sink, numLogs), 2*time.Second, 10*time.Millisecond,
		"expected %d but got %d logs",
		numLogs, sink.LogRecordCount(),
	)
	// TODO: Figure out a nice way to assert each logs entry content.
	// require.Equal(t, expectedLogs, sink.AllLogs())
	require.NoError(t, rcvr.Shutdown(context.Background()))
	converter.Stop()
}

func consumeNLogsFromConverter(ch <-chan plog.Logs, count int, wg *sync.WaitGroup) {
	defer wg.Done()

	n := 0
	for pLog := range ch {
		n += pLog.ResourceLogs().At(0).ScopeLogs().At(0).LogRecords().Len()

		if n == count {
			return
		}
	}
}

func newRotatingLogger(t *testing.T, tempDir string, maxLines, maxBackups int, copyTruncate, sequential bool) *log.Logger {
	path := filepath.Join(tempDir, "test.log")
	rotator := &nanojack.Logger{
		Filename:     path,
		MaxLines:     maxLines,
		MaxBackups:   maxBackups,
		CopyTruncate: copyTruncate,
		Sequential:   sequential,
	}

	t.Cleanup(func() { _ = rotator.Close() })

	return log.New(rotator, "", 0)
}

func expectNLogs(sink *consumertest.LogsSink, expected int) func() bool {
	return func() bool { return sink.LogRecordCount() == expected }
}

func testdataConfigYamlAsMap() *FileLogConfig {
	return &FileLogConfig{
		BaseConfig: adapter.BaseConfig{
			ReceiverSettings: config.NewReceiverSettings(config.NewComponentID(typeStr)),
			Operators: adapter.OperatorConfigs{
				map[string]interface{}{
					"type":  "regex_parser",
					"regex": "^(?P<time>\\d{4}-\\d{2}-\\d{2}) (?P<sev>[A-Z]*) (?P<msg>.*)$",
					"severity": map[string]interface{}{
						"parse_from": "attributes.sev",
					},
					"timestamp": map[string]interface{}{
						"layout":     "%Y-%m-%d",
						"parse_from": "attributes.time",
					},
				},
			},
			Converter: adapter.ConverterConfig{
				MaxFlushCount: 100,
				FlushInterval: 100 * time.Millisecond,
			},
		},
		Input: adapter.InputConfig{
			"include": []interface{}{
				"testdata/simple.log",
			},
			"start_at": "beginning",
		},
	}
}

func testdataRotateTestYamlAsMap(tempDir string) *FileLogConfig {
	return &FileLogConfig{
		BaseConfig: adapter.BaseConfig{
			ReceiverSettings: config.NewReceiverSettings(config.NewComponentID(typeStr)),
			Operators: adapter.OperatorConfigs{
				map[string]interface{}{
					"type":  "regex_parser",
					"regex": "^(?P<ts>\\d{4}-\\d{2}-\\d{2}) (?P<msg>[^\n]+)",
					"timestamp": map[interface{}]interface{}{
						"layout":     "%Y-%m-%d",
						"parse_from": "body.ts",
					},
				},
			},
			Converter: adapter.ConverterConfig{},
		},
		Input: adapter.InputConfig{
			"type": "file_input",
			"include": []interface{}{
				fmt.Sprintf("%s/*", tempDir),
			},
			"include_file_name": false,
			"poll_interval":     "10ms",
			"start_at":          "beginning",
		},
	}
}
<commit_msg>[receiver/filelog] Enable lint and fix issues (#11796)<commit_after>// Copyright The OpenTelemetry 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 filelogreceiver

import (
	"context"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"testing"
	"time"

	"github.com/observiq/nanojack"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/collector/component/componenttest"
	"go.opentelemetry.io/collector/config"
	"go.opentelemetry.io/collector/config/configtest"
	"go.opentelemetry.io/collector/consumer/consumertest"
	"go.opentelemetry.io/collector/pdata/plog"
	"go.opentelemetry.io/collector/service/servicetest"

	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/adapter"
	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/stanza/entry"
)

func TestDefaultConfig(t *testing.T) {
	factory := NewFactory()
	cfg := factory.CreateDefaultConfig()
	require.NotNil(t, cfg, "failed to create default config")
	require.NoError(t, configtest.CheckConfigStruct(cfg))
}

func TestLoadConfig(t *testing.T) {
	factories, err := componenttest.NopFactories()
	assert.Nil(t, err)

	factory := NewFactory()
	factories.Receivers[typeStr] = factory
	cfg, err := servicetest.LoadConfigAndValidate(filepath.Join("testdata", "config.yaml"), factories)
	require.NoError(t, err)
	require.NotNil(t, cfg)

	assert.Equal(t, len(cfg.Receivers), 1)

	assert.Equal(t, testdataConfigYamlAsMap(), cfg.Receivers[config.NewComponentID("filelog")])
}

func TestCreateWithInvalidInputConfig(t *testing.T) {
	t.Parallel()

	cfg := testdataConfigYamlAsMap()
	cfg.Input["include"] = "not an array"

	_, err := NewFactory().CreateLogsReceiver(
		context.Background(),
		componenttest.NewNopReceiverCreateSettings(),
		cfg,
		new(consumertest.LogsSink),
	)
	require.Error(t, err, "receiver creation should fail if given invalid input config")
}

func TestReadStaticFile(t *testing.T) {
	t.Parallel()

	expectedTimestamp, _ := time.ParseInLocation("2006-01-02", "2020-08-25", time.Local)

	f := NewFactory()
	sink := new(consumertest.LogsSink)

	cfg := testdataConfigYamlAsMap()
	cfg.Converter.MaxFlushCount = 10
	cfg.Converter.FlushInterval = time.Millisecond

	converter := adapter.NewConverter()
	converter.Start()
	defer converter.Stop()

	var wg sync.WaitGroup
	wg.Add(1)
	go consumeNLogsFromConverter(converter.OutChannel(), 3, &wg)

	rcvr, err := f.CreateLogsReceiver(context.Background(), componenttest.NewNopReceiverCreateSettings(), cfg, sink)
	require.NoError(t, err, "failed to create receiver")
	require.NoError(t, rcvr.Start(context.Background(), componenttest.NewNopHost()))

	// Build the expected set by using adapter.Converter to translate entries
	// to pdata Logs.
	queueEntry := func(t *testing.T, c *adapter.Converter, msg string, severity entry.Severity) {
		e := entry.New()
		e.Timestamp = expectedTimestamp
		require.NoError(t, e.Set(entry.NewBodyField("msg"), msg))
		e.Severity = severity
		e.AddAttribute("file_name", "simple.log")
		require.NoError(t, c.Batch([]*entry.Entry{e}))
	}
	queueEntry(t, converter, "Something routine", entry.Info)
	queueEntry(t, converter, "Something bad happened!", entry.Error)
	queueEntry(t, converter, "Some details...", entry.Debug)

	dir, err := os.Getwd()
	require.NoError(t, err)
	t.Logf("Working Directory: %s", dir)

	wg.Wait()

	require.Eventually(t, expectNLogs(sink, 3), 2*time.Second, 5*time.Millisecond,
		"expected %d but got %d logs",
		3, sink.LogRecordCount(),
	)
	// TODO: Figure out a nice way to assert each logs entry content.
	// require.Equal(t, expectedLogs, sink.AllLogs())
	require.NoError(t, rcvr.Shutdown(context.Background()))
}

func TestReadRotatingFiles(t *testing.T) {

	tests := []rotationTest{
		{
			name:         "CopyTruncateTimestamped",
			copyTruncate: true,
			sequential:   false,
		},
		{
			name:         "MoveCreateTimestamped",
			copyTruncate: false,
			sequential:   false,
		},
		{
			name:         "CopyTruncateSequential",
			copyTruncate: true,
			sequential:   true,
		},
		{
			name:         "MoveCreateSequential",
			copyTruncate: false,
			sequential:   true,
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, tc.Run)
	}
}

type rotationTest struct {
	name         string
	copyTruncate bool
	sequential   bool
}

func (rt *rotationTest) Run(t *testing.T) {
	t.Parallel()

	tempDir := t.TempDir()

	f := NewFactory()
	sink := new(consumertest.LogsSink)

	cfg := testdataRotateTestYamlAsMap(tempDir)
	cfg.Converter.MaxFlushCount = 1
	cfg.Converter.FlushInterval = time.Millisecond

	// With a max of 100 logs per file and 1 backup file, rotation will occur
	// when more than 100 logs are written, and deletion when more than 200 are written.
	// Write 300 and validate that we got the all despite rotation and deletion.
	logger := newRotatingLogger(t, tempDir, 100, 1, rt.copyTruncate, rt.sequential)
	numLogs := 300

	// Build expected outputs
	expectedTimestamp, _ := time.ParseInLocation("2006-01-02", "2020-08-25", time.Local)
	converter := adapter.NewConverter()
	converter.Start()

	var wg sync.WaitGroup
	wg.Add(1)
	go consumeNLogsFromConverter(converter.OutChannel(), numLogs, &wg)

	rcvr, err := f.CreateLogsReceiver(context.Background(), componenttest.NewNopReceiverCreateSettings(), cfg, sink)
	require.NoError(t, err, "failed to create receiver")
	require.NoError(t, rcvr.Start(context.Background(), componenttest.NewNopHost()))

	for i := 0; i < numLogs; i++ {
		msg := fmt.Sprintf("This is a simple log line with the number %3d", i)

		// Build the expected set by converting entries to pdata Logs...
		e := entry.New()
		e.Timestamp = expectedTimestamp
		require.NoError(t, e.Set(entry.NewBodyField("msg"), msg))
		require.NoError(t, converter.Batch([]*entry.Entry{e}))

		// ... and write the logs lines to the actual file consumed by receiver.
		logger.Printf("2020-08-25 %s", msg)
		time.Sleep(time.Millisecond)
	}

	wg.Wait()
	require.Eventually(t, expectNLogs(sink, numLogs), 2*time.Second, 10*time.Millisecond,
		"expected %d but got %d logs",
		numLogs, sink.LogRecordCount(),
	)
	// TODO: Figure out a nice way to assert each logs entry content.
	// require.Equal(t, expectedLogs, sink.AllLogs())
	require.NoError(t, rcvr.Shutdown(context.Background()))
	converter.Stop()
}

func consumeNLogsFromConverter(ch <-chan plog.Logs, count int, wg *sync.WaitGroup) {
	defer wg.Done()

	n := 0
	for pLog := range ch {
		n += pLog.ResourceLogs().At(0).ScopeLogs().At(0).LogRecords().Len()

		if n == count {
			return
		}
	}
}

func newRotatingLogger(t *testing.T, tempDir string, maxLines, maxBackups int, copyTruncate, sequential bool) *log.Logger {
	path := filepath.Join(tempDir, "test.log")
	rotator := &nanojack.Logger{
		Filename:     path,
		MaxLines:     maxLines,
		MaxBackups:   maxBackups,
		CopyTruncate: copyTruncate,
		Sequential:   sequential,
	}

	t.Cleanup(func() { _ = rotator.Close() })

	return log.New(rotator, "", 0)
}

func expectNLogs(sink *consumertest.LogsSink, expected int) func() bool {
	return func() bool { return sink.LogRecordCount() == expected }
}

func testdataConfigYamlAsMap() *FileLogConfig {
	return &FileLogConfig{
		BaseConfig: adapter.BaseConfig{
			ReceiverSettings: config.NewReceiverSettings(config.NewComponentID(typeStr)),
			Operators: adapter.OperatorConfigs{
				map[string]interface{}{
					"type":  "regex_parser",
					"regex": "^(?P<time>\\d{4}-\\d{2}-\\d{2}) (?P<sev>[A-Z]*) (?P<msg>.*)$",
					"severity": map[string]interface{}{
						"parse_from": "attributes.sev",
					},
					"timestamp": map[string]interface{}{
						"layout":     "%Y-%m-%d",
						"parse_from": "attributes.time",
					},
				},
			},
			Converter: adapter.ConverterConfig{
				MaxFlushCount: 100,
				FlushInterval: 100 * time.Millisecond,
			},
		},
		Input: adapter.InputConfig{
			"include": []interface{}{
				"testdata/simple.log",
			},
			"start_at": "beginning",
		},
	}
}

func testdataRotateTestYamlAsMap(tempDir string) *FileLogConfig {
	return &FileLogConfig{
		BaseConfig: adapter.BaseConfig{
			ReceiverSettings: config.NewReceiverSettings(config.NewComponentID(typeStr)),
			Operators: adapter.OperatorConfigs{
				map[string]interface{}{
					"type":  "regex_parser",
					"regex": "^(?P<ts>\\d{4}-\\d{2}-\\d{2}) (?P<msg>[^\n]+)",
					"timestamp": map[interface{}]interface{}{
						"layout":     "%Y-%m-%d",
						"parse_from": "body.ts",
					},
				},
			},
			Converter: adapter.ConverterConfig{},
		},
		Input: adapter.InputConfig{
			"type": "file_input",
			"include": []interface{}{
				fmt.Sprintf("%s/*", tempDir),
			},
			"include_file_name": false,
			"poll_interval":     "10ms",
			"start_at":          "beginning",
		},
	}
}
<|endoftext|>


<commit_before>package consumergroup

import (
	"errors"
	"fmt"
	"reflect"
	"runtime/debug"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/meitu/go-zookeeper/zk"
	"github.com/sirupsen/logrus"
)

const (
	cgInit = iota
	cgStart
	cgStopped
)

const (
	restartEvent = iota
	quitEvent
)

// ConsumerGroup consume message from Kafka with rebalancing supports
type ConsumerGroup struct {
	name           string
	storage        groupStorage
	topicConsumers map[string]*topicConsumer
	saramaClient   sarama.Client
	saramaConsumer sarama.Consumer

	id          string
	state       int
	wg          sync.WaitGroup
	stopCh      chan struct{}
	triggerCh   chan int
	triggerOnce *sync.Once
	owners      map[string]map[int32]string

	config *Config
	logger *logrus.Logger

	onLoad, onClose []func()
}

// NewConsumerGroup create the ConsumerGroup instance with config
func NewConsumerGroup(config *Config) (*ConsumerGroup, error) {
	if config == nil {
		return nil, errors.New("config can't be empty")
	}
	err := config.validate()
	if err != nil {
		return nil, fmt.Errorf("vaildate config failed, as %s", err)
	}

	cg := new(ConsumerGroup)
	cg.state = cgInit
	cg.config = config
	cg.id = config.ConsumerID
	if cg.id == "" {
		cg.id = genConsumerID()
	}
	cg.name = config.GroupID
	cg.triggerCh = make(chan int)
	cg.topicConsumers = make(map[string]*topicConsumer)
	cg.onLoad = make([]func(), 0)
	cg.onClose = make([]func(), 0)
	cg.storage = newZKGroupStorage(config.ZkList, config.ZkSessionTimeout)
	cg.logger = logrus.New()
	if _, ok := cg.storage.(*zkGroupStorage); ok {
		cg.storage.(*zkGroupStorage).Chroot(config.Chroot)
	}

	err = cg.initSaramaConsumer()
	if err != nil {
		return nil, fmt.Errorf("init sarama consumer, as %s", err)
	}
	cg.owners = make(map[string]map[int32]string)
	for _, topic := range config.TopicList {
		cg.topicConsumers[topic] = newTopicConsumer(cg, topic)
		cg.owners[topic] = make(map[int32]string)
	}
	return cg, nil
}

func (cg *ConsumerGroup) initSaramaConsumer() error {
	brokerList, err := cg.storage.getBrokerList()
	if err != nil {
		return err
	}
	if len(brokerList) == 0 {
		return errors.New("no broker alive")
	}
	cg.saramaClient, err = sarama.NewClient(brokerList, cg.config.SaramaConfig)
	if err != nil {
		return err
	}
	cg.saramaConsumer, err = sarama.NewConsumerFromClient(cg.saramaClient)
	return err
}

// Start would register ConsumerGroup, and rebalance would be triggered.
// ConsumerGroup computes the partitions which should be consumed by consumer's num, and start fetching message.
func (cg *ConsumerGroup) Start() error {
	// exit when failed to register the consumer
	err := cg.storage.registerConsumer(cg.name, cg.id, nil)
	if err != nil && err != zk.ErrNodeExists {
		return err
	}
	cg.wg.Add(1)
	go cg.start()
	return nil
}

// Stop would unregister ConsumerGroup, and rebalance would be triggered.
// The partitions which consumed by this ConsumerGroup would be assigned to others.
func (cg *ConsumerGroup) Stop() {
	cg.stop()
	cg.wg.Wait()
}

// SetLogger use to set the user's logger the consumer group
func (cg *ConsumerGroup) SetLogger(l *logrus.Logger) {
	if l != nil {
		cg.logger = l
	}
}

// IsStopped return whether the ConsumerGroup was stopped or not.
func (cg *ConsumerGroup) IsStopped() bool {
	return cg.state == cgStopped
}

func (cg *ConsumerGroup) callRecover() {
	if err := recover(); err != nil {
		cg.logger.WithFields(logrus.Fields{
			"group": cg.name,
			"err":   err,
			"stack": string(debug.Stack()),
		}).Error("Recover panic")
		cg.stop()
	}
}

func (cg *ConsumerGroup) start() {
	var wg sync.WaitGroup

	defer cg.callRecover()
	defer func() {
		cg.state = cgStopped
		err := cg.storage.deleteConsumer(cg.name, cg.id)
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to delete consumer from zk")
		}
		for _, tc := range cg.topicConsumers {
			close(tc.messages)
			close(tc.errors)
		}
		cg.wg.Done()
	}()

CONSUME_TOPIC_LOOP:
	for {
		cg.logger.WithField("group", cg.name).Info("Consumer group started")
		cg.triggerOnce = new(sync.Once)
		cg.stopCh = make(chan struct{})

		err := cg.watchRebalance()
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to watch rebalance")
			cg.stop()
			return
		}
		wg.Add(1)
		go func() {
			defer cg.callRecover()
			defer wg.Done()
			cg.autoReconnect(cg.storage.(*zkGroupStorage).sessionTimeout / 3)
		}()
		for _, consumer := range cg.topicConsumers {
			wg.Add(1)
			consumer.start()
			go func(tc *topicConsumer) {
				defer cg.callRecover()
				defer wg.Done()
				tc.wg.Wait()
				cg.logger.WithFields(logrus.Fields{
					"group": tc.group,
					"topic": tc.name,
				}).Info("Stop the topic consumer")
			}(consumer)
		}
		cg.state = cgStart
		for _, onLoadFunc := range cg.onLoad {
			onLoadFunc()
		}
		msg := <-cg.triggerCh
		for _, onCloseFunc := range cg.onClose {
			onCloseFunc()
		}
		switch msg {
		case restartEvent:
			close(cg.stopCh)
			// The stop channel was used to notify partition's consumer to stop consuming when rebalance is triggered.
			// So we should reinit when rebalance was triggered, as it would be closed.
			wg.Wait()
			continue CONSUME_TOPIC_LOOP
		case quitEvent:
			close(cg.stopCh)
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup is stopping")
			wg.Wait()
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup was stopped")
			return
		}
	}
}

func (cg *ConsumerGroup) stop() {
	cg.triggerOnce.Do(func() { cg.triggerCh <- quitEvent })
}

func (cg *ConsumerGroup) triggerRebalance() {
	cg.triggerOnce.Do(func() { cg.triggerCh <- restartEvent })
}

func (cg *ConsumerGroup) getPartitionConsumer(topic string, partition int32, nextOffset int64) (sarama.PartitionConsumer, error) {
	consumer, err := cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	if err == sarama.ErrOffsetOutOfRange {
		nextOffset = cg.config.OffsetAutoReset
		consumer, err = cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	}
	if err != nil {
		return nil, err
	}
	return consumer, nil
}

// GetMessages was used to get a unbuffered message's channel from specified topic
func (cg *ConsumerGroup) GetMessages(topic string) (<-chan *sarama.ConsumerMessage, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.messages, true
	}
	return nil, false
}

// GetErrors was used to get a unbuffered error's channel from specified topic
func (cg *ConsumerGroup) GetErrors(topic string) (<-chan *sarama.ConsumerError, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.errors, true
	}
	return nil, false
}

// OnLoad load callback function that runs after startup
func (cg *ConsumerGroup) OnLoad(cb func()) {
	cg.onLoad = append(cg.onLoad, cb)
}

// OnClose load callback function that runs before the end
func (cg *ConsumerGroup) OnClose(cb func()) {
	cg.onClose = append(cg.onClose, cb)
}

func (cg *ConsumerGroup) autoReconnect(interval time.Duration) {
	timer := time.NewTimer(interval)
	cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was started")
	defer cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was stopped")
	for {
		select {
		case <-cg.stopCh:
			return
		case <-timer.C:
			timer.Reset(interval)
			exist, err := cg.storage.existsConsumer(cg.name, cg.id)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to check consumer existence")
				break
			}
			if exist {
				break
			}
			err = cg.storage.registerConsumer(cg.name, cg.id, nil)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to re-register consumer")
			}
		}
	}
}

func (cg *ConsumerGroup) watchRebalance() error {
	consumersWatcher, err := cg.storage.watchConsumerList(cg.name)
	if err != nil {
		return err
	}
	topicsChange, topicWatchers, err := cg.watchTopics(cg.config.TopicList)
	if err != nil {
		return err
	}
	go func() {
		defer cg.callRecover()
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was started")
		select {
		case <-consumersWatcher.EvCh:
			cg.triggerRebalance()
			cg.logger.WithField("group", cg.name).Info("Trigger rebalance while consumers was changed")
			for _, tw := range topicWatchers {
				cg.storage.removeWatcher(tw)
			}
		case topic := <-topicsChange:
			cg.triggerRebalance()
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"topic": topic,
			}).Info("Trigger rebalance while partitions was changed")
			cg.storage.removeWatcher(consumersWatcher)
		case <-cg.stopCh:
		}
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was exited")
	}()
	return nil
}

// CommitOffset is used to commit offset when auto commit was disabled.
func (cg *ConsumerGroup) CommitOffset(topic string, partition int32, offset int64) error {
	if cg.config.OffsetAutoCommitEnable {
		return errors.New("commit offset take effect when offset auto commit was disabled")
	}
	return cg.storage.commitOffset(cg.name, topic, partition, offset)
}

// GetOffsets return the offset in memory for debug
func (cg *ConsumerGroup) GetOffsets() map[string]interface{} {
	topics := make(map[string]interface{})
	for topic, tc := range cg.topicConsumers {
		topics[topic] = tc.getOffsets()
	}
	return topics
}

// Owners return owners of all partitions
func (cg *ConsumerGroup) Owners() map[string]map[int32]string {
	return cg.owners
}

func (cg *ConsumerGroup) watchTopics(topics []string) (<-chan string, []*zk.Watcher, error) {
	ch := make(chan string)
	cases := make([]reflect.SelectCase, len(topics))
	watchers := make([]*zk.Watcher, len(topics))
	for i, topic := range topics {
		w, err := cg.storage.watchTopic(topic)
		if err != nil {
			return nil, nil, err
		}
		watchers[i] = w
		cases[i] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(w.EvCh)}
	}
	go func(cases []reflect.SelectCase, ch chan string, topics []string) {
		chosen, _, ok := reflect.Select(cases)
		if !ok {
			//the chosen channel has been closed.
			return
		}
		topic := topics[chosen]
		num, err := cg.storage.getPartitionsNum(topic)
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"topic": topic,
				"err":   err,
			}).Error("Failed to get partitions in zookeeper after topic metadata change")
			return
		}
		for {
			cg.saramaClient.RefreshMetadata(topic)
			partitions, err := cg.saramaClient.Partitions(topic)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"topic": topic,
					"err":   err,
				}).Error("Failed to get partitions in broker after topic metadata change")
				return
			}
			if len(partitions) == num {
				break
			}
			time.Sleep(100 * time.Millisecond)
		}

		ch <- topics[chosen]
	}(cases, ch, topics)
	return ch, watchers, nil
}
<commit_msg>MOD: simpify the return value<commit_after>package consumergroup

import (
	"errors"
	"fmt"
	"reflect"
	"runtime/debug"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/meitu/go-zookeeper/zk"
	"github.com/sirupsen/logrus"
)

const (
	cgInit = iota
	cgStart
	cgStopped
)

const (
	restartEvent = iota
	quitEvent
)

// ConsumerGroup consume message from Kafka with rebalancing supports
type ConsumerGroup struct {
	name           string
	storage        groupStorage
	topicConsumers map[string]*topicConsumer
	saramaClient   sarama.Client
	saramaConsumer sarama.Consumer

	id          string
	state       int
	wg          sync.WaitGroup
	stopCh      chan struct{}
	triggerCh   chan int
	triggerOnce *sync.Once
	owners      map[string]map[int32]string

	config *Config
	logger *logrus.Logger

	onLoad, onClose []func()
}

// NewConsumerGroup create the ConsumerGroup instance with config
func NewConsumerGroup(config *Config) (*ConsumerGroup, error) {
	if config == nil {
		return nil, errors.New("config can't be empty")
	}
	err := config.validate()
	if err != nil {
		return nil, fmt.Errorf("vaildate config failed, as %s", err)
	}

	cg := new(ConsumerGroup)
	cg.state = cgInit
	cg.config = config
	cg.id = config.ConsumerID
	if cg.id == "" {
		cg.id = genConsumerID()
	}
	cg.name = config.GroupID
	cg.triggerCh = make(chan int)
	cg.topicConsumers = make(map[string]*topicConsumer)
	cg.onLoad = make([]func(), 0)
	cg.onClose = make([]func(), 0)
	cg.storage = newZKGroupStorage(config.ZkList, config.ZkSessionTimeout)
	cg.logger = logrus.New()
	if _, ok := cg.storage.(*zkGroupStorage); ok {
		cg.storage.(*zkGroupStorage).Chroot(config.Chroot)
	}

	err = cg.initSaramaConsumer()
	if err != nil {
		return nil, fmt.Errorf("init sarama consumer, as %s", err)
	}
	cg.owners = make(map[string]map[int32]string)
	for _, topic := range config.TopicList {
		cg.topicConsumers[topic] = newTopicConsumer(cg, topic)
		cg.owners[topic] = make(map[int32]string)
	}
	return cg, nil
}

func (cg *ConsumerGroup) initSaramaConsumer() error {
	brokerList, err := cg.storage.getBrokerList()
	if err != nil {
		return err
	}
	if len(brokerList) == 0 {
		return errors.New("no broker alive")
	}
	cg.saramaClient, err = sarama.NewClient(brokerList, cg.config.SaramaConfig)
	if err != nil {
		return err
	}
	cg.saramaConsumer, err = sarama.NewConsumerFromClient(cg.saramaClient)
	return err
}

// Start would register ConsumerGroup, and rebalance would be triggered.
// ConsumerGroup computes the partitions which should be consumed by consumer's num, and start fetching message.
func (cg *ConsumerGroup) Start() error {
	// exit when failed to register the consumer
	err := cg.storage.registerConsumer(cg.name, cg.id, nil)
	if err != nil && err != zk.ErrNodeExists {
		return err
	}
	cg.wg.Add(1)
	go cg.start()
	return nil
}

// Stop would unregister ConsumerGroup, and rebalance would be triggered.
// The partitions which consumed by this ConsumerGroup would be assigned to others.
func (cg *ConsumerGroup) Stop() {
	cg.stop()
	cg.wg.Wait()
}

// SetLogger use to set the user's logger the consumer group
func (cg *ConsumerGroup) SetLogger(l *logrus.Logger) {
	if l != nil {
		cg.logger = l
	}
}

// IsStopped return whether the ConsumerGroup was stopped or not.
func (cg *ConsumerGroup) IsStopped() bool {
	return cg.state == cgStopped
}

func (cg *ConsumerGroup) callRecover() {
	if err := recover(); err != nil {
		cg.logger.WithFields(logrus.Fields{
			"group": cg.name,
			"err":   err,
			"stack": string(debug.Stack()),
		}).Error("Recover panic")
		cg.stop()
	}
}

func (cg *ConsumerGroup) start() {
	var wg sync.WaitGroup

	defer cg.callRecover()
	defer func() {
		cg.state = cgStopped
		err := cg.storage.deleteConsumer(cg.name, cg.id)
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to delete consumer from zk")
		}
		for _, tc := range cg.topicConsumers {
			close(tc.messages)
			close(tc.errors)
		}
		cg.wg.Done()
	}()

CONSUME_TOPIC_LOOP:
	for {
		cg.logger.WithField("group", cg.name).Info("Consumer group started")
		cg.triggerOnce = new(sync.Once)
		cg.stopCh = make(chan struct{})

		err := cg.watchRebalance()
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"err":   err,
			}).Error("Failed to watch rebalance")
			cg.stop()
			return
		}
		wg.Add(1)
		go func() {
			defer cg.callRecover()
			defer wg.Done()
			cg.autoReconnect(cg.storage.(*zkGroupStorage).sessionTimeout / 3)
		}()
		for _, consumer := range cg.topicConsumers {
			wg.Add(1)
			consumer.start()
			go func(tc *topicConsumer) {
				defer cg.callRecover()
				defer wg.Done()
				tc.wg.Wait()
				cg.logger.WithFields(logrus.Fields{
					"group": tc.group,
					"topic": tc.name,
				}).Info("Stop the topic consumer")
			}(consumer)
		}
		cg.state = cgStart
		for _, onLoadFunc := range cg.onLoad {
			onLoadFunc()
		}
		msg := <-cg.triggerCh
		for _, onCloseFunc := range cg.onClose {
			onCloseFunc()
		}
		switch msg {
		case restartEvent:
			close(cg.stopCh)
			// The stop channel was used to notify partition's consumer to stop consuming when rebalance is triggered.
			// So we should reinit when rebalance was triggered, as it would be closed.
			wg.Wait()
			continue CONSUME_TOPIC_LOOP
		case quitEvent:
			close(cg.stopCh)
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup is stopping")
			wg.Wait()
			cg.logger.WithField("group", cg.name).Info("ConsumerGroup was stopped")
			return
		}
	}
}

func (cg *ConsumerGroup) stop() {
	cg.triggerOnce.Do(func() { cg.triggerCh <- quitEvent })
}

func (cg *ConsumerGroup) triggerRebalance() {
	cg.triggerOnce.Do(func() { cg.triggerCh <- restartEvent })
}

func (cg *ConsumerGroup) getPartitionConsumer(topic string, partition int32, nextOffset int64) (sarama.PartitionConsumer, error) {
	consumer, err := cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	if err == sarama.ErrOffsetOutOfRange {
		nextOffset = cg.config.OffsetAutoReset
		consumer, err = cg.saramaConsumer.ConsumePartition(topic, partition, nextOffset)
	}
	return consumer, err
}

// GetMessages was used to get a unbuffered message's channel from specified topic
func (cg *ConsumerGroup) GetMessages(topic string) (<-chan *sarama.ConsumerMessage, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.messages, true
	}
	return nil, false
}

// GetErrors was used to get a unbuffered error's channel from specified topic
func (cg *ConsumerGroup) GetErrors(topic string) (<-chan *sarama.ConsumerError, bool) {
	if topicConsumer, ok := cg.topicConsumers[topic]; ok {
		return topicConsumer.errors, true
	}
	return nil, false
}

// OnLoad load callback function that runs after startup
func (cg *ConsumerGroup) OnLoad(cb func()) {
	cg.onLoad = append(cg.onLoad, cb)
}

// OnClose load callback function that runs before the end
func (cg *ConsumerGroup) OnClose(cb func()) {
	cg.onClose = append(cg.onClose, cb)
}

func (cg *ConsumerGroup) autoReconnect(interval time.Duration) {
	timer := time.NewTimer(interval)
	cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was started")
	defer cg.logger.WithField("group", cg.name).Info("The auto-reconnect consumer thread was stopped")
	for {
		select {
		case <-cg.stopCh:
			return
		case <-timer.C:
			timer.Reset(interval)
			exist, err := cg.storage.existsConsumer(cg.name, cg.id)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to check consumer existence")
				break
			}
			if exist {
				break
			}
			err = cg.storage.registerConsumer(cg.name, cg.id, nil)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"group": cg.name,
					"err":   err,
				}).Error("Failed to re-register consumer")
			}
		}
	}
}

func (cg *ConsumerGroup) watchRebalance() error {
	consumersWatcher, err := cg.storage.watchConsumerList(cg.name)
	if err != nil {
		return err
	}
	topicsChange, topicWatchers, err := cg.watchTopics(cg.config.TopicList)
	if err != nil {
		return err
	}
	go func() {
		defer cg.callRecover()
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was started")
		select {
		case <-consumersWatcher.EvCh:
			cg.triggerRebalance()
			cg.logger.WithField("group", cg.name).Info("Trigger rebalance while consumers was changed")
			for _, tw := range topicWatchers {
				cg.storage.removeWatcher(tw)
			}
		case topic := <-topicsChange:
			cg.triggerRebalance()
			cg.logger.WithFields(logrus.Fields{
				"group": cg.name,
				"topic": topic,
			}).Info("Trigger rebalance while partitions was changed")
			cg.storage.removeWatcher(consumersWatcher)
		case <-cg.stopCh:
		}
		cg.logger.WithField("group", cg.name).Info("Rebalance watcher thread was exited")
	}()
	return nil
}

// CommitOffset is used to commit offset when auto commit was disabled.
func (cg *ConsumerGroup) CommitOffset(topic string, partition int32, offset int64) error {
	if cg.config.OffsetAutoCommitEnable {
		return errors.New("commit offset take effect when offset auto commit was disabled")
	}
	return cg.storage.commitOffset(cg.name, topic, partition, offset)
}

// GetOffsets return the offset in memory for debug
func (cg *ConsumerGroup) GetOffsets() map[string]interface{} {
	topics := make(map[string]interface{})
	for topic, tc := range cg.topicConsumers {
		topics[topic] = tc.getOffsets()
	}
	return topics
}

// Owners return owners of all partitions
func (cg *ConsumerGroup) Owners() map[string]map[int32]string {
	return cg.owners
}

func (cg *ConsumerGroup) watchTopics(topics []string) (<-chan string, []*zk.Watcher, error) {
	ch := make(chan string)
	cases := make([]reflect.SelectCase, len(topics))
	watchers := make([]*zk.Watcher, len(topics))
	for i, topic := range topics {
		w, err := cg.storage.watchTopic(topic)
		if err != nil {
			return nil, nil, err
		}
		watchers[i] = w
		cases[i] = reflect.SelectCase{Dir: reflect.SelectRecv, Chan: reflect.ValueOf(w.EvCh)}
	}
	go func(cases []reflect.SelectCase, ch chan string, topics []string) {
		chosen, _, ok := reflect.Select(cases)
		if !ok {
			//the chosen channel has been closed.
			return
		}
		topic := topics[chosen]
		num, err := cg.storage.getPartitionsNum(topic)
		if err != nil {
			cg.logger.WithFields(logrus.Fields{
				"topic": topic,
				"err":   err,
			}).Error("Failed to get partitions in zookeeper after topic metadata change")
			return
		}
		for {
			cg.saramaClient.RefreshMetadata(topic)
			partitions, err := cg.saramaClient.Partitions(topic)
			if err != nil {
				cg.logger.WithFields(logrus.Fields{
					"topic": topic,
					"err":   err,
				}).Error("Failed to get partitions in broker after topic metadata change")
				return
			}
			if len(partitions) == num {
				break
			}
			time.Sleep(100 * time.Millisecond)
		}

		ch <- topics[chosen]
	}(cases, ch, topics)
	return ch, watchers, nil
}
<|endoftext|>


<commit_before>// Copyright © 2016 Thomas Rabaix <thomas.rabaix@gmail.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package commands

import (
	"github.com/mitchellh/cli"
	helper "github.com/rande/gitlab-ci-helper"
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
)

func Test_Project_Builds_Artifacts(t *testing.T) {
	fpBuilds, err := os.Open("../fixtures/builds.json")
	assert.NoError(t, err)

	fpProjects, err := os.Open("../fixtures/projects.json")
	assert.NoError(t, err)

	fpProject, err := os.Open("../fixtures/project.json")
	assert.NoError(t, err)

	fpCommits, err := os.Open("../fixtures/commits.json")
	assert.NoError(t, err)

	fpArchive, err := os.Open("../fixtures/artifacts.zip")
	assert.NoError(t, err)

	headers := http.Header{
		"Content-Type": []string{"application/json"},
	}

	reqs := []*helper.FakeRequest{
		{
			Path:   "/api/v3/projects/3",
			Method: "GET",
			Response: &http.Response{
				Body:   fpProject,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects",
			Method: "GET",
			Response: &http.Response{
				Body:   fpProjects,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects/3/builds",
			Method: "GET",
			Response: &http.Response{
				Body:   fpBuilds,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects/3/repository/commits/889935cf4d3e7558ae6c0d4dd62e20ea600f5a57/builds",
			Method: "GET",
			Response: &http.Response{
				Body:   fpCommits,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects/3/builds/69/artifacts",
			Method: "GET",
			Response: &http.Response{
				Body: fpArchive,
				Header: http.Header{
					"Content-Type": []string{"application/zip"},
				},
			},
		},
	}

	envs := map[string]string{}

	helper.WrapperTestCommand(reqs, envs, t, func(ts *httptest.Server) {
		ui := &cli.MockUi{}
		c := &ProjectBuildArtifactCommand{
			Ui: ui,
		}

		code := c.Run([]string{"-project", "3", "-ref", "889935cf4d3e7558ae6c0d4dd62e20ea600f5a57", "-job", "rubocop"})

		assert.Equal(t, 0, code)

		expected := "Found project: Diaspora/Diaspora Project Site (id: 3)\nFound build - stage:test status:canceled id:69\nDownloading artifacts... (artifacts.zip)\nDone!\n"
		assert.Equal(t, expected, ui.OutputWriter.String())
		assert.Equal(t, "", ui.ErrorWriter.String())

		os.Remove(c.ArtifactsFile)
	})

}<commit_msg>fix(test): adjust cs<commit_after>// Copyright © 2016 Thomas Rabaix <thomas.rabaix@gmail.com>.
//
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.

package commands

import (
	"github.com/mitchellh/cli"
	helper "github.com/rande/gitlab-ci-helper"
	"github.com/stretchr/testify/assert"
	"net/http"
	"net/http/httptest"
	"os"
	"testing"
)

func Test_Project_Builds_Artifacts(t *testing.T) {
	fpBuilds, err := os.Open("../fixtures/builds.json")
	assert.NoError(t, err)

	fpProjects, err := os.Open("../fixtures/projects.json")
	assert.NoError(t, err)

	fpProject, err := os.Open("../fixtures/project.json")
	assert.NoError(t, err)

	fpCommits, err := os.Open("../fixtures/commits.json")
	assert.NoError(t, err)

	fpArchive, err := os.Open("../fixtures/artifacts.zip")
	assert.NoError(t, err)

	headers := http.Header{
		"Content-Type": []string{"application/json"},
	}

	reqs := []*helper.FakeRequest{
		{
			Path:   "/api/v3/projects/3",
			Method: "GET",
			Response: &http.Response{
				Body:   fpProject,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects",
			Method: "GET",
			Response: &http.Response{
				Body:   fpProjects,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects/3/builds",
			Method: "GET",
			Response: &http.Response{
				Body:   fpBuilds,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects/3/repository/commits/889935cf4d3e7558ae6c0d4dd62e20ea600f5a57/builds",
			Method: "GET",
			Response: &http.Response{
				Body:   fpCommits,
				Header: headers,
			},
		},
		{
			Path:   "/api/v3/projects/3/builds/69/artifacts",
			Method: "GET",
			Response: &http.Response{
				Body: fpArchive,
				Header: http.Header{
					"Content-Type": []string{"application/zip"},
				},
			},
		},
	}

	envs := map[string]string{}

	helper.WrapperTestCommand(reqs, envs, t, func(ts *httptest.Server) {
		ui := &cli.MockUi{}
		c := &ProjectBuildArtifactCommand{
			Ui: ui,
		}

		code := c.Run([]string{"-project", "3", "-ref", "889935cf4d3e7558ae6c0d4dd62e20ea600f5a57", "-job", "rubocop"})

		assert.Equal(t, 0, code)

		expected := "Found project: Diaspora/Diaspora Project Site (id: 3)\nFound build - stage:test status:canceled id:69\nDownloading artifacts... (artifacts.zip)\nDone!\n"
		assert.Equal(t, expected, ui.OutputWriter.String())
		assert.Equal(t, "", ui.ErrorWriter.String())

		os.Remove(c.ArtifactsFile)
	})

}
<|endoftext|>


<commit_before>package redis2kvstore

import (
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"runtime/debug"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gallir/bytebufferpool"
	"github.com/gallir/smart-relayer/lib"
	"github.com/gallir/smart-relayer/redis"
	"github.com/gallir/smart-relayer/redis/radix.improved/redis"
	"github.com/golang/snappy"
)

// Server is the thread that listen for clients' connections
type Server struct {
	sync.Mutex
	config   lib.RelayerConfig
	done     chan bool
	exiting  bool
	listener net.Listener

	client *http.Client

	lastConnection time.Time
	lastError      time.Time
	running        int64
}

const (
	strContentType      = "application/octet-stream"
	maxConnections      = 2
	requestBufferSize   = 1024 * 10
	maxConnectionsTries = 3
	connectionRetry     = 5 * time.Second
	errorsFrame         = 10 * time.Second
	maxErrors           = 10 // Limit of errors to restart the connection
	connectTimeout      = 15 * time.Second
	defaultExpire       = 2 * 60 * 60 // 2h
	retryTime           = 100 * time.Millisecond
	waitingForExit      = 2 * time.Second
)

var (
	errBadCmd      = errors.New("ERR bad command")
	errKO          = errors.New("fatal error")
	errOverloaded  = errors.New("http server overloaded")
	errNotFound    = errors.New("Not found")
	respOK         = redis.NewRespSimple("OK")
	respBadCommand = redis.NewResp(errBadCmd)
	respKO         = redis.NewResp(errKO)
	commands       map[string]*redis.Resp

	pool = bytebufferpool.Pool{}
)

func init() {
	commands = map[string]*redis.Resp{
		"PING":    respOK,
		"HMSET":   respOK,
		"EXPIRE":  respOK,
		"HGET":    respOK,
		"HGETALL": respOK,
	}
}

// New creates a new Redis local server
func New(c lib.RelayerConfig, done chan bool) (*Server, error) {
	srv := &Server{
		done: done,
		client: &http.Client{
			Timeout: 10 * time.Second,
			Transport: &http.Transport{
				DisableKeepAlives:  false,
				MaxIdleConns:       1024,
				IdleConnTimeout:    30 * time.Second,
				DisableCompression: true,
				DialContext: (&net.Dialer{
					Timeout:   5 * time.Second,
					KeepAlive: 30 * time.Second,
					DualStack: true,
				}).DialContext,
			},
		},
	}
	srv.Reload(&c)

	return srv, nil
}

// Reload the configuration
func (srv *Server) Reload(c *lib.RelayerConfig) (err error) {
	srv.Lock()
	defer srv.Unlock()

	srv.config = *c

	if srv.config.MaxConnections <= 0 {
		srv.config.MaxConnections = maxConnections
	}

	return nil
}

// Start accepts incoming connections on the Listener
func (srv *Server) Start() (e error) {
	srv.Lock()
	defer srv.Unlock()

	srv.listener, e = lib.NewListener(srv.config)
	if e != nil {
		return e
	}

	// Serve clients
	go func(l net.Listener) {
		defer srv.listener.Close()
		for {
			netConn, e := l.Accept()
			if e != nil {
				if netErr, ok := e.(net.Error); ok && netErr.Timeout() {
					// Paranoid, ignore timeout errors
					log.Println("redis2kvstore ERROR: timeout at local listener", srv.config.ListenHost(), e)
					continue
				}
				if srv.exiting {
					log.Println("redis2kvstore: exiting local listener", srv.config.ListenHost())
					return
				}
				log.Fatalln("redis2kvstore ERROR: emergency error in local listener", srv.config.ListenHost(), e)
				return
			}
			go srv.handleConnection(netConn)
		}
	}(srv.listener)

	return
}

// Exit closes the listener and send done to main
func (srv *Server) Exit() {
	srv.exiting = true

	if srv.listener != nil {
		srv.listener.Close()
	}

	retry := 0
	for retry < 10 {
		n := atomic.LoadInt64(&srv.running)
		if n == 0 {
			break
		}
		log.Printf("redis2kvstore Waiting that %d process are still running", n)
		time.Sleep(waitingForExit)
		retry++
	}

	if n := atomic.LoadInt64(&srv.running); n > 0 {
		log.Printf("redis2kvstore ERROR: %d messages lost", n)
	}

	// finishing the server
	srv.done <- true
}

func (srv *Server) handleConnection(netCon net.Conn) {
	defer netCon.Close()

	reader := redis.NewRespReader(netCon)

	pending := getPending()
	defer func() {
		if len(pending) > 0 {
			for key, p := range pending {
				if p != nil && !p.Sent && len(p.Fields) > 0 {
					srv.send(key, defaultExpire, p)
				}
				delete(pending, key)
				putPoolHMSet(p) // Return Hmset to the pool
			}
		}
		putPending(pending)
	}()

	for {
		r := reader.Read()

		if r.IsType(redis.IOErr) {
			if redis.IsTimeout(r) {
				// Paranoid, don't close it just log it
				log.Println("redis2kvstore: Local client listen timeout at", srv.config.Listen)
				continue
			}
			// Connection was closed
			return
		}

		req := lib.NewRequest(r, &srv.config)
		if req == nil {
			respBadCommand.WriteTo(netCon)
			continue
		}

		validCommand, ok := commands[req.Command]
		if !ok {
			respBadCommand.WriteTo(netCon)
			continue
		}

		switch req.Command {
		case "HMSET":
			if len(req.Items) < 4 || len(req.Items)%2 != 0 {
				respKO.WriteTo(netCon)
				continue
			}
			validCommand.WriteTo(netCon)

			key, _ := req.Items[1].Str()
			if _, ok := pending[key]; !ok {
				pending[key] = getPoolHMSet()
			}
			pending[key].processItems(req.Items[2:])
		case "EXPIRE":
			if len(req.Items) != 3 {
				respKO.WriteTo(netCon)
				continue
			}

			key, _ := req.Items[1].Str()
			p, ok := pending[key]
			if !ok || key == "" {
				log.Printf("redis2kvstore ERROR: Invalid key %s", key)
				respBadCommand.WriteTo(netCon)
				continue
			}

			expire, _ := req.Items[2].Int()
			if expire == 0 {
				expire = defaultExpire
			}
			p.Sent = true
			go srv.send(key, expire, p.clone())
			validCommand.WriteTo(netCon)
		case "HGETALL":
			if len(req.Items) != 2 {
				respKO.WriteTo(netCon)
				continue
			}
			key, _ := req.Items[1].Str()

			// Return information that is in memory
			if m, ok := pending[key]; ok {
				if r, err := m.getAllAsRedis(); err == nil {
					r.WriteTo(netCon)
					continue
				}
			}

			// If is not in memory we go to the cluster
			items, err := srv.getHGetAll(key)
			if err != nil {
				redis.NewResp(err).WriteTo(netCon)
				continue
			}
			items.WriteTo(netCon)
			items.ReleaseBuffers()
		case "HGET":
			if len(req.Items) != 3 {
				respKO.WriteTo(netCon)
				continue
			}
			key, _ := req.Items[1].Str()
			item, _ := req.Items[2].Str()

			// Return information that is in memory
			if m, ok := pending[key]; ok {
				if r, err := m.getOneAsRedis(item); err == nil {
					r.WriteTo(netCon)
					continue
				}
			}

			// If is not in memory we go to the cluster
			g, err := srv.getHGet(key, item)
			if err != nil {
				if err == errNotFound {
					redis.NewResp(nil).WriteTo(netCon)
					continue
				}
				redis.NewResp(err).WriteTo(netCon)
				continue
			}
			g.WriteTo(netCon)
			g.ReleaseBuffers()
		}

		for _, i := range req.Items {
			i.ReleaseBuffers()
		}

	}
}

func (srv *Server) send(key string, expire int, p *Hmset) {
	defer func(lenFields int) {
		if r := recover(); r != nil {
			log.Printf("redis2kvstore: Recovered in send [%s, %d] %s: %s\n", key, lenFields, r, debug.Stack())
		}
	}(len(p.Fields))

	// Send back to the pool the Hmset
	defer putPoolHMSet(p)

	// Get bytes pool for the compression
	w := pool.Get()
	defer pool.Put(w)

	// Increase the number of running process before create a new hmset
	atomic.AddInt64(&srv.running, 1)
	defer atomic.AddInt64(&srv.running, -1)

	url := fmt.Sprintf("%s/%s/%ds", srv.config.URL, key, expire)
	b, _ := p.Marshal()

	if srv.config.Gzip > 0 {
		gzWriter := lib.GetGzipWriterLevel(w, srv.config.Gzip)
		gzWriter.Write(b)
		gzWriter.Close()
		lib.PutGzipWriter(gzWriter)
	} else if srv.config.Compress {
		w.Write(compress.Bytes(b))
	} else {
		w.Write(b)
	}
	if w.Len() <= 0 {
		log.Printf("redis2kvstore ERROR empty body: %s", url)
		return
	}

	for i := 0; i < maxConnectionsTries; i++ {
		resp, err := srv.client.Post(url, strContentType, bytes.NewReader(w.B))
		if err == nil {
			defer resp.Body.Close()
			if resp.StatusCode == 200 {
				// Success
				return
			}
			log.Printf("redis2kvstore ERROR post: [%d] %s %s", resp.StatusCode, url, err)
		} else {
			log.Printf("redis2kvstore ERROR connect: %s %s", url, err)
		}
		time.Sleep(retryTime)
	}
}

// get will get via http the content of the key, this content will be
// an slice of bytes. The function will uncompress it based on the configuration
// IMPORTANT: this functions use a sync.Pool for the &Hmet{}. You should send
// the struct back to the pull after use it.
func (srv *Server) get(key string) (*Hmset, error) {
	buf := pool.Get()
	defer pool.Put(buf)

	url := fmt.Sprintf("%s/get/%s", srv.config.URL, key)
	resp, err := http.Get(url)
	if err != nil {
		log.Printf("redis2kvstore ERROR connect: %s %s", url, err)
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Error: %s %s", url, resp.Status)
	}

	lib.Debugf("redis2kvstore: get %s", url)

	if srv.config.Gzip > 0 || srv.config.Gunzip {
		gzReader, err := lib.GetGzipReader(resp.Body)
		if err != nil {
			return nil, err
		}
		buf.ReadFrom(gzReader)
		gzReader.Close()
		lib.PutGzipReader(gzReader)
	} else if srv.config.Compress {
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		buf.B, err = snappy.Decode(buf.B, b[len(redis.MarkerSnappy):])
		if err != nil {
			return nil, err
		}
	} else {
		var err error
		buf.B, err = ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
	}
	if buf.Len() <= 0 {
		return nil, fmt.Errorf("Empty response: %s", url)
	}

	// Notice that you should return m to the pool
	m := getPoolHMSet()
	if err := m.Unmarshal(buf.B); err != nil {
		return nil, err
	}
	return m, nil
}

func (srv *Server) getHGetAll(key string) (*redis.Resp, error) {
	var m *Hmset
	var err error

	for i := 0; i < maxConnectionsTries; i++ {
		m, err = srv.get(key)
		if err == nil {
			defer putPoolHMSet(m)
			return m.getAllAsRedis()
		}
		time.Sleep(retryTime * 2)
	}

	return nil, err
}

func (srv *Server) getHGet(key, field string) (*redis.Resp, error) {
	var m *Hmset
	var err error

	for i := 0; i < maxConnectionsTries; i++ {
		m, err = srv.get(key)
		if err == nil {
			defer putPoolHMSet(m)
			return m.getOneAsRedis(field)
		}
		time.Sleep(retryTime * 2)
	}

	return nil, err
}
<commit_msg>Be sure that we are reading all the body response<commit_after>package redis2kvstore

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"runtime/debug"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gallir/bytebufferpool"
	"github.com/gallir/smart-relayer/lib"
	"github.com/gallir/smart-relayer/redis"
	"github.com/gallir/smart-relayer/redis/radix.improved/redis"
	"github.com/golang/snappy"
)

// Server is the thread that listen for clients' connections
type Server struct {
	sync.Mutex
	config   lib.RelayerConfig
	done     chan bool
	exiting  bool
	listener net.Listener

	client *http.Client

	lastConnection time.Time
	lastError      time.Time
	running        int64
}

const (
	strContentType      = "application/octet-stream"
	maxConnections      = 2
	requestBufferSize   = 1024 * 10
	maxConnectionsTries = 3
	connectionRetry     = 5 * time.Second
	errorsFrame         = 10 * time.Second
	maxErrors           = 10 // Limit of errors to restart the connection
	connectTimeout      = 15 * time.Second
	defaultExpire       = 2 * 60 * 60 // 2h
	retryTime           = 100 * time.Millisecond
	waitingForExit      = 2 * time.Second
)

var (
	errBadCmd      = errors.New("ERR bad command")
	errKO          = errors.New("fatal error")
	errOverloaded  = errors.New("http server overloaded")
	errNotFound    = errors.New("Not found")
	respOK         = redis.NewRespSimple("OK")
	respBadCommand = redis.NewResp(errBadCmd)
	respKO         = redis.NewResp(errKO)
	commands       map[string]*redis.Resp

	pool = bytebufferpool.Pool{}
)

func init() {
	commands = map[string]*redis.Resp{
		"PING":    respOK,
		"HMSET":   respOK,
		"EXPIRE":  respOK,
		"HGET":    respOK,
		"HGETALL": respOK,
	}
}

// New creates a new Redis local server
func New(c lib.RelayerConfig, done chan bool) (*Server, error) {
	srv := &Server{
		done: done,
		client: &http.Client{
			Timeout: 10 * time.Second,
			Transport: &http.Transport{
				DisableKeepAlives:  false,
				MaxIdleConns:       1024,
				IdleConnTimeout:    30 * time.Second,
				DisableCompression: true,
				DialContext: (&net.Dialer{
					Timeout:   5 * time.Second,
					KeepAlive: 30 * time.Second,
					DualStack: true,
				}).DialContext,
			},
		},
	}
	srv.Reload(&c)

	return srv, nil
}

// Reload the configuration
func (srv *Server) Reload(c *lib.RelayerConfig) (err error) {
	srv.Lock()
	defer srv.Unlock()

	srv.config = *c

	if srv.config.MaxConnections <= 0 {
		srv.config.MaxConnections = maxConnections
	}

	return nil
}

// Start accepts incoming connections on the Listener
func (srv *Server) Start() (e error) {
	srv.Lock()
	defer srv.Unlock()

	srv.listener, e = lib.NewListener(srv.config)
	if e != nil {
		return e
	}

	// Serve clients
	go func(l net.Listener) {
		defer srv.listener.Close()
		for {
			netConn, e := l.Accept()
			if e != nil {
				if netErr, ok := e.(net.Error); ok && netErr.Timeout() {
					// Paranoid, ignore timeout errors
					log.Println("redis2kvstore ERROR: timeout at local listener", srv.config.ListenHost(), e)
					continue
				}
				if srv.exiting {
					log.Println("redis2kvstore: exiting local listener", srv.config.ListenHost())
					return
				}
				log.Fatalln("redis2kvstore ERROR: emergency error in local listener", srv.config.ListenHost(), e)
				return
			}
			go srv.handleConnection(netConn)
		}
	}(srv.listener)

	return
}

// Exit closes the listener and send done to main
func (srv *Server) Exit() {
	srv.exiting = true

	if srv.listener != nil {
		srv.listener.Close()
	}

	retry := 0
	for retry < 10 {
		n := atomic.LoadInt64(&srv.running)
		if n == 0 {
			break
		}
		log.Printf("redis2kvstore Waiting that %d process are still running", n)
		time.Sleep(waitingForExit)
		retry++
	}

	if n := atomic.LoadInt64(&srv.running); n > 0 {
		log.Printf("redis2kvstore ERROR: %d messages lost", n)
	}

	// finishing the server
	srv.done <- true
}

func (srv *Server) handleConnection(netCon net.Conn) {
	defer netCon.Close()

	reader := redis.NewRespReader(netCon)

	pending := getPending()
	defer func() {
		if len(pending) > 0 {
			for key, p := range pending {
				if p != nil && !p.Sent && len(p.Fields) > 0 {
					srv.send(key, defaultExpire, p)
				}
				delete(pending, key)
				putPoolHMSet(p) // Return Hmset to the pool
			}
		}
		putPending(pending)
	}()

	for {
		r := reader.Read()

		if r.IsType(redis.IOErr) {
			if redis.IsTimeout(r) {
				// Paranoid, don't close it just log it
				log.Println("redis2kvstore: Local client listen timeout at", srv.config.Listen)
				continue
			}
			// Connection was closed
			return
		}

		req := lib.NewRequest(r, &srv.config)
		if req == nil {
			respBadCommand.WriteTo(netCon)
			continue
		}

		validCommand, ok := commands[req.Command]
		if !ok {
			respBadCommand.WriteTo(netCon)
			continue
		}

		switch req.Command {
		case "HMSET":
			if len(req.Items) < 4 || len(req.Items)%2 != 0 {
				respKO.WriteTo(netCon)
				continue
			}
			validCommand.WriteTo(netCon)

			key, _ := req.Items[1].Str()
			if _, ok := pending[key]; !ok {
				pending[key] = getPoolHMSet()
			}
			pending[key].processItems(req.Items[2:])
		case "EXPIRE":
			if len(req.Items) != 3 {
				respKO.WriteTo(netCon)
				continue
			}

			key, _ := req.Items[1].Str()
			p, ok := pending[key]
			if !ok || key == "" {
				log.Printf("redis2kvstore ERROR: Invalid key %s", key)
				respBadCommand.WriteTo(netCon)
				continue
			}

			expire, _ := req.Items[2].Int()
			if expire == 0 {
				expire = defaultExpire
			}
			p.Sent = true
			go srv.send(key, expire, p.clone())
			validCommand.WriteTo(netCon)
		case "HGETALL":
			if len(req.Items) != 2 {
				respKO.WriteTo(netCon)
				continue
			}
			key, _ := req.Items[1].Str()

			// Return information that is in memory
			if m, ok := pending[key]; ok {
				if r, err := m.getAllAsRedis(); err == nil {
					r.WriteTo(netCon)
					continue
				}
			}

			// If is not in memory we go to the cluster
			items, err := srv.getHGetAll(key)
			if err != nil {
				redis.NewResp(err).WriteTo(netCon)
				continue
			}
			items.WriteTo(netCon)
			items.ReleaseBuffers()
		case "HGET":
			if len(req.Items) != 3 {
				respKO.WriteTo(netCon)
				continue
			}
			key, _ := req.Items[1].Str()
			item, _ := req.Items[2].Str()

			// Return information that is in memory
			if m, ok := pending[key]; ok {
				if r, err := m.getOneAsRedis(item); err == nil {
					r.WriteTo(netCon)
					continue
				}
			}

			// If is not in memory we go to the cluster
			g, err := srv.getHGet(key, item)
			if err != nil {
				if err == errNotFound {
					redis.NewResp(nil).WriteTo(netCon)
					continue
				}
				redis.NewResp(err).WriteTo(netCon)
				continue
			}
			g.WriteTo(netCon)
			g.ReleaseBuffers()
		}

		for _, i := range req.Items {
			i.ReleaseBuffers()
		}

	}
}

func (srv *Server) send(key string, expire int, p *Hmset) {
	defer func(lenFields int) {
		if r := recover(); r != nil {
			log.Printf("redis2kvstore: Recovered in send [%s, %d] %s: %s\n", key, lenFields, r, debug.Stack())
		}
	}(len(p.Fields))

	// Send back to the pool the Hmset
	defer putPoolHMSet(p)

	// Get bytes pool for the compression
	w := pool.Get()
	defer pool.Put(w)

	// Increase the number of running process before create a new hmset
	atomic.AddInt64(&srv.running, 1)
	defer atomic.AddInt64(&srv.running, -1)

	url := fmt.Sprintf("%s/%s/%ds", srv.config.URL, key, expire)
	b, _ := p.Marshal()

	if srv.config.Gzip > 0 {
		gzWriter := lib.GetGzipWriterLevel(w, srv.config.Gzip)
		gzWriter.Write(b)
		gzWriter.Close()
		lib.PutGzipWriter(gzWriter)
	} else if srv.config.Compress {
		w.Write(compress.Bytes(b))
	} else {
		w.Write(b)
	}
	if w.Len() <= 0 {
		log.Printf("redis2kvstore ERROR empty body: %s", url)
		return
	}

	for i := 0; i < maxConnectionsTries; i++ {
		resp, err := srv.client.Post(url, strContentType, bytes.NewReader(w.B))
		if err == nil {
			// https://golang.org/pkg/net/http/#Response
			// ... The default HTTP client's Transport may not
			// reuse HTTP/1.x "keep-alive" TCP connections if the Body is
			// not read to completion and closed.
			defer resp.Body.Close()
			if _, err = io.Copy(ioutil.Discard, resp.Body); err != nil {
				log.Printf("redis2kvstore WARNING: %s %s", url, err)
			}
			if resp.StatusCode == 200 {
				// Success
				return
			}
			log.Printf("redis2kvstore ERROR post: [%d] %s %s", resp.StatusCode, url, err)
		} else {
			log.Printf("redis2kvstore ERROR connect: %s %s", url, err)
		}
		time.Sleep(retryTime)
	}
}

// get will get via http the content of the key, this content will be
// an slice of bytes. The function will uncompress it based on the configuration
// IMPORTANT: this functions use a sync.Pool for the &Hmet{}. You should send
// the struct back to the pull after use it.
func (srv *Server) get(key string) (*Hmset, error) {
	buf := pool.Get()
	defer pool.Put(buf)

	url := fmt.Sprintf("%s/get/%s", srv.config.URL, key)
	resp, err := http.Get(url)
	if err != nil {
		log.Printf("redis2kvstore ERROR connect: %s %s", url, err)
		return nil, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("Error: %s %s", url, resp.Status)
	}

	lib.Debugf("redis2kvstore: get %s", url)

	if srv.config.Gzip > 0 || srv.config.Gunzip {
		gzReader, err := lib.GetGzipReader(resp.Body)
		if err != nil {
			return nil, err
		}
		buf.ReadFrom(gzReader)
		gzReader.Close()
		lib.PutGzipReader(gzReader)
	} else if srv.config.Compress {
		b, err := ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
		buf.B, err = snappy.Decode(buf.B, b[len(redis.MarkerSnappy):])
		if err != nil {
			return nil, err
		}
	} else {
		var err error
		buf.B, err = ioutil.ReadAll(resp.Body)
		if err != nil {
			return nil, err
		}
	}
	if buf.Len() <= 0 {
		return nil, fmt.Errorf("Empty response: %s", url)
	}

	// Notice that you should return m to the pool
	m := getPoolHMSet()
	if err := m.Unmarshal(buf.B); err != nil {
		return nil, err
	}
	return m, nil
}

func (srv *Server) getHGetAll(key string) (*redis.Resp, error) {
	var m *Hmset
	var err error

	for i := 0; i < maxConnectionsTries; i++ {
		m, err = srv.get(key)
		if err == nil {
			defer putPoolHMSet(m)
			return m.getAllAsRedis()
		}
		time.Sleep(retryTime * 2)
	}

	return nil, err
}

func (srv *Server) getHGet(key, field string) (*redis.Resp, error) {
	var m *Hmset
	var err error

	for i := 0; i < maxConnectionsTries; i++ {
		m, err = srv.get(key)
		if err == nil {
			defer putPoolHMSet(m)
			return m.getOneAsRedis(field)
		}
		time.Sleep(retryTime * 2)
	}

	return nil, err
}
<|endoftext|>


<commit_before>// Copyright 2017 Google Inc.
//
// 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 envoy

// MeshConfig defines proxy mesh variables
type MeshConfig struct {
	// DiscoveryAddress is the DNS address for Envoy discovery service
	DiscoveryAddress string
	// MixerAddress is the DNS address for Istio Mixer service
	MixerAddress string
	// ProxyPort is the Envoy proxy port
	ProxyPort int
	// AdminPort is the administrative interface port
	AdminPort int
	// Envoy binary path
	BinaryPath string
	// Envoy config root path
	ConfigPath string
	// Envoy runtime config path
	RuntimePath string
	// Envoy access log path
	AccessLogPath string
}

// Config defines the schema for Envoy JSON configuration format
// See: https://lyft.github.io/envoy/docs/configuration/overview/overview.html
type Config struct {
	RootRuntime    RootRuntime    `json:"runtime"`
	Listeners      []Listener     `json:"listeners"`
	Admin          Admin          `json:"admin"`
	ClusterManager ClusterManager `json:"cluster_manager"`
}

// RootRuntime definition.
// See: https://lyft.github.io/envoy/docs/configuration/overview/overview.html
type RootRuntime struct {
	SymlinkRoot          string `json:"symlink_root"`
	Subdirectory         string `json:"subdirectory"`
	OverrideSubdirectory string `json:"override_subdirectory,omitempty"`
}

// Listener definition
type Listener struct {
	Port           int             `json:"port"`
	Filters        []NetworkFilter `json:"filters"`
	BindToPort     bool            `json:"bind_to_port"`
	UseOriginalDst bool            `json:"use_original_dst,omitempty"`
}

// Admin definition
type Admin struct {
	AccessLogPath string `json:"access_log_path"`
	Port          int    `json:"port"`
}

// ClusterManager definition
type ClusterManager struct {
	Clusters []Cluster `json:"clusters"`
	SDS      SDS       `json:"sds"`
}

// SDS is a service discovery service definition
type SDS struct {
	Cluster        Cluster `json:"cluster"`
	RefreshDelayMs int     `json:"refresh_delay_ms"`
}

// Cluster definition
type Cluster struct {
	Name                     string            `json:"name"`
	ServiceName              string            `json:"service_name,omitempty"`
	ConnectTimeoutMs         int               `json:"connect_timeout_ms"`
	Type                     string            `json:"type"`
	LbType                   string            `json:"lb_type"`
	MaxRequestsPerConnection int               `json:"max_requests_per_connection,omitempty"`
	Hosts                    []Host            `json:"hosts,omitempty"`
	Features                 string            `json:"features,omitempty"`
	CircuitBreaker           *CircuitBreaker   `json:"circuit_breaker,omitempty"`
	OutlierDetection         *OutlierDetection `json:"outlier_detection,omitempty"`
}

// CircuitBreaker definition
// See: https://lyft.github.io/envoy/docs/configuration/cluster_manager/cluster_circuit_breakers.html#circuit-breakers
type CircuitBreaker struct {
	MaxConnections    int `json:"max_connections,omitempty"`
	MaxPendingRequest int `json:"max_pending_requests,omitempty"`
	MaxRequests       int `json:"max_requests,omitempty"`
	MaxRetries        int `json:"max_retries,omitempty"`
}

// OutlierDetection definition
// See: https://lyft.github.io/envoy/docs/configuration/cluster_manager/cluster_runtime.html#outlier-detection
type OutlierDetection struct {
	ConsecutiveError   int `json:"consecutive_5xx,omitempty"`
	IntervalMS         int `json:"interval_ms,omitempty"`
	BaseEjectionTimeMS int `json:"base_ejection_time_ms,omitempty"`
	MaxEjectionPercent int `json:"max_ejection_percent,omitempty"`
}

// Filter definition
type Filter struct {
	Type   string      `json:"type"`
	Name   string      `json:"name"`
	Config interface{} `json:"config"`
}

// FilterRouterConfig definition
type FilterRouterConfig struct {
	// DynamicStats defaults to true
	DynamicStats bool `json:"dynamic_stats,omitempty"`
}

// FilterEndpointsConfig definition
type FilterEndpointsConfig struct {
	ServiceConfig string `json:"service_config,omitempty"`
	ServerConfig  string `json:"server_config,omitempty"`
}

// FilterFaultConfig definition
type FilterFaultConfig struct {
	Abort           *AbortFilter `json:"abort,omitempty"`
	Delay           *DelayFilter `json:"delay,omitempty"`
	Headers         []Header     `json:"headers,omitempty"`
	UpstreamCluster string       `json:"upstream_cluster,omitempty"`
}

// AbortFilter definition
type AbortFilter struct {
	Percent    int `json:"abort_percent,omitempty"`
	HTTPStatus int `json:"http_status,omitempty"`
}

// DelayFilter definition
type DelayFilter struct {
	Type     string `json:"type,omitempty"`
	Percent  int    `json:"fixed_delay_percent,omitempty"`
	Duration int    `json:"fixed_duration_ms,omitempty"`
}

// NetworkFilter definition
type NetworkFilter struct {
	Type   string              `json:"type"`
	Name   string              `json:"name"`
	Config NetworkFilterConfig `json:"config"`
}

// NetworkFilterConfig definition
type NetworkFilterConfig struct {
	CodecType         string      `json:"codec_type"`
	StatPrefix        string      `json:"stat_prefix"`
	GenerateRequestID bool        `json:"generate_request_id,omitempty"`
	RouteConfig       RouteConfig `json:"route_config"`
	Filters           []Filter    `json:"filters"`
	AccessLog         []AccessLog `json:"access_log"`
	Cluster           string      `json:"cluster,omitempty"`
}

// AccessLog definition.
type AccessLog struct {
	Path   string `json:"path"`
	Format string `json:"format,omitempty"`
	Filter string `json:"filter,omitempty"`
}

// RouteConfig definition
type RouteConfig struct {
	VirtualHosts []VirtualHost `json:"virtual_hosts"`
}

// VirtualHost definition
type VirtualHost struct {
	Name    string   `json:"name"`
	Domains []string `json:"domains"`
	Routes  []Route  `json:"routes"`
}

// Route definition
type Route struct {
	Runtime          *Runtime         `json:"runtime,omitempty"`
	Prefix           string           `json:"prefix"`
	PrefixRewrite    string           `json:"prefix_rewrite,omitempty"`
	Cluster          string           `json:"cluster"`
	WeightedClusters *WeightedCluster `json:"weighted_clusters,omitempty"`
	Headers          []Header         `json:"headers,omitempty"`
	TimeoutMS        int              `json:"timeout_ms,omitempty"`
	RetryPolicy      RetryPolicy      `json:"retry_policy,omitempty"`
}

// RetryPolicy definition
// See: https://lyft.github.io/envoy/docs/configuration/http_conn_man/route_config/route.html#retry-policy
type RetryPolicy struct {
	Policy     string `json:"retry_on"` //5xx,connect-failure,refused-stream
	NumRetries int    `json:"num_retries"`
}

// Runtime definition
type Runtime struct {
	Key     string `json:"key"`
	Default int    `json:"default"`
}

// WeightedCluster definition
// See https://lyft.github.io/envoy/docs/configuration/http_conn_man/route_config/route.html
type WeightedCluster struct {
	Clusters         []WeightedClusterEntry `json:"clusters"`
	RuntimeKeyPrefix string                 `json:"runtime_key_prefix,omitempty"`
}

// WeightedClusterEntry definition. Describes the format of each entry in the WeightedCluster
type WeightedClusterEntry struct {
	Name   string `json:"name"`
	Weight int    `json:"weight"`
}

// Header definition
type Header struct {
	Name  string `json:"name"`
	Value string `json:"value"`
	Regex bool   `json:"regex,omitempty"`
}

// Host definition
type Host struct {
	URL string `json:"url"`
}

// Constant values
const (
	LbTypeRoundRobin = "round_robin"
)

// ListenersByPort sorts listeners by port
type ListenersByPort []Listener

func (l ListenersByPort) Len() int {
	return len(l)
}

func (l ListenersByPort) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l ListenersByPort) Less(i, j int) bool {
	return l[i].Port < l[j].Port
}

// ClustersByName sorts clusters by name
type ClustersByName []Cluster

func (s ClustersByName) Len() int {
	return len(s)
}

func (s ClustersByName) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s ClustersByName) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

// HostsByName sorts clusters by name
type HostsByName []VirtualHost

func (s HostsByName) Len() int {
	return len(s)
}

func (s HostsByName) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s HostsByName) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

// RoutesByCluster sorts routes by name
type RoutesByCluster []Route

func (r RoutesByCluster) Len() int {
	return len(r)
}

func (r RoutesByCluster) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

func (r RoutesByCluster) Less(i, j int) bool {
	return r[i].Cluster < r[j].Cluster
}

// ByName implements sort
type ByName []Cluster

// Len length
func (a ByName) Len() int {
	return len(a)
}

// Swap elements
func (a ByName) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

// Less compare
func (a ByName) Less(i, j int) bool {
	return a[i].Name < a[j].Name
}

//ByHost implement sort
type ByHost []Host

// Len length
func (a ByHost) Len() int {
	return len(a)
}

// Swap elements
func (a ByHost) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

// Less compare
func (a ByHost) Less(i, j int) bool {
	return a[i].URL < a[j].URL
}
<commit_msg>unshuffling structs in resources<commit_after>// Copyright 2017 Google Inc.
//
// 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 envoy

// MeshConfig defines proxy mesh variables
type MeshConfig struct {
	// DiscoveryAddress is the DNS address for Envoy discovery service
	DiscoveryAddress string
	// MixerAddress is the DNS address for Istio Mixer service
	MixerAddress string
	// ProxyPort is the Envoy proxy port
	ProxyPort int
	// AdminPort is the administrative interface port
	AdminPort int
	// Envoy binary path
	BinaryPath string
	// Envoy config root path
	ConfigPath string
	// Envoy runtime config path
	RuntimePath string
	// Envoy access log path
	AccessLogPath string
}

// Config defines the schema for Envoy JSON configuration format
type Config struct {
	RootRuntime    RootRuntime    `json:"runtime"`
	Listeners      []Listener     `json:"listeners"`
	Admin          Admin          `json:"admin"`
	ClusterManager ClusterManager `json:"cluster_manager"`
}

// RootRuntime definition.
// See: https://lyft.github.io/envoy/docs/configuration/overview/overview.html
type RootRuntime struct {
	SymlinkRoot          string `json:"symlink_root"`
	Subdirectory         string `json:"subdirectory"`
	OverrideSubdirectory string `json:"override_subdirectory,omitempty"`
}

// AbortFilter definition
type AbortFilter struct {
	Percent    int `json:"abort_percent,omitempty"`
	HTTPStatus int `json:"http_status,omitempty"`
}

// DelayFilter definition
type DelayFilter struct {
	Type     string `json:"type,omitempty"`
	Percent  int    `json:"fixed_delay_percent,omitempty"`
	Duration int    `json:"fixed_duration_ms,omitempty"`
}

// Header definition
type Header struct {
	Name  string `json:"name"`
	Value string `json:"value"`
	Regex bool   `json:"regex,omitempty"`
}

// FilterEndpointsConfig definition
type FilterEndpointsConfig struct {
	ServiceConfig string `json:"service_config,omitempty"`
	ServerConfig  string `json:"server_config,omitempty"`
}

// FilterFaultConfig definition
type FilterFaultConfig struct {
	Abort           *AbortFilter `json:"abort,omitempty"`
	Delay           *DelayFilter `json:"delay,omitempty"`
	Headers         []Header     `json:"headers,omitempty"`
	UpstreamCluster string       `json:"upstream_cluster,omitempty"`
}

// FilterRouterConfig definition
type FilterRouterConfig struct {
	// DynamicStats defaults to true
	DynamicStats bool `json:"dynamic_stats,omitempty"`
}

// Filter definition
type Filter struct {
	Type   string      `json:"type"`
	Name   string      `json:"name"`
	Config interface{} `json:"config"`
}

// Runtime definition
type Runtime struct {
	Key     string `json:"key"`
	Default int    `json:"default"`
}

// Route definition
type Route struct {
	Runtime          *Runtime         `json:"runtime,omitempty"`
	Prefix           string           `json:"prefix"`
	PrefixRewrite    string           `json:"prefix_rewrite,omitempty"`
	Cluster          string           `json:"cluster"`
	WeightedClusters *WeightedCluster `json:"weighted_clusters,omitempty"`
	Headers          []Header         `json:"headers,omitempty"`
	TimeoutMS        int              `json:"timeout_ms,omitempty"`
	RetryPolicy      RetryPolicy      `json:"retry_policy,omitempty"`
}

// RetryPolicy definition
// See: https://lyft.github.io/envoy/docs/configuration/http_conn_man/route_config/route.html#retry-policy
type RetryPolicy struct {
	Policy     string `json:"retry_on"` //5xx,connect-failure,refused-stream
	NumRetries int    `json:"num_retries"`
}

// WeightedCluster definition
// See https://lyft.github.io/envoy/docs/configuration/http_conn_man/route_config/route.html
type WeightedCluster struct {
	Clusters         []WeightedClusterEntry `json:"clusters"`
	RuntimeKeyPrefix string                 `json:"runtime_key_prefix,omitempty"`
}

// WeightedClusterEntry definition. Describes the format of each entry in the WeightedCluster
type WeightedClusterEntry struct {
	Name   string `json:"name"`
	Weight int    `json:"weight"`
}

// VirtualHost definition
type VirtualHost struct {
	Name    string   `json:"name"`
	Domains []string `json:"domains"`
	Routes  []Route  `json:"routes"`
}

// RouteConfig definition
type RouteConfig struct {
	VirtualHosts []VirtualHost `json:"virtual_hosts"`
}

// AccessLog definition.
type AccessLog struct {
	Path   string `json:"path"`
	Format string `json:"format,omitempty"`
	Filter string `json:"filter,omitempty"`
}

// NetworkFilterConfig definition
type NetworkFilterConfig struct {
	CodecType         string      `json:"codec_type"`
	StatPrefix        string      `json:"stat_prefix"`
	GenerateRequestID bool        `json:"generate_request_id,omitempty"`
	RouteConfig       RouteConfig `json:"route_config"`
	Filters           []Filter    `json:"filters"`
	AccessLog         []AccessLog `json:"access_log"`
	Cluster           string      `json:"cluster,omitempty"`
}

// NetworkFilter definition
type NetworkFilter struct {
	Type   string              `json:"type"`
	Name   string              `json:"name"`
	Config NetworkFilterConfig `json:"config"`
}

// Listener definition
type Listener struct {
	Port           int             `json:"port"`
	Filters        []NetworkFilter `json:"filters"`
	BindToPort     bool            `json:"bind_to_port"`
	UseOriginalDst bool            `json:"use_original_dst,omitempty"`
}

// Admin definition
type Admin struct {
	AccessLogPath string `json:"access_log_path"`
	Port          int    `json:"port"`
}

// Host definition
type Host struct {
	URL string `json:"url"`
}

// Constant values
const (
	LbTypeRoundRobin = "round_robin"
)

// Cluster definition
type Cluster struct {
	Name                     string            `json:"name"`
	ServiceName              string            `json:"service_name,omitempty"`
	ConnectTimeoutMs         int               `json:"connect_timeout_ms"`
	Type                     string            `json:"type"`
	LbType                   string            `json:"lb_type"`
	MaxRequestsPerConnection int               `json:"max_requests_per_connection,omitempty"`
	Hosts                    []Host            `json:"hosts,omitempty"`
	Features                 string            `json:"features,omitempty"`
	CircuitBreaker           *CircuitBreaker   `json:"circuit_breaker,omitempty"`
	OutlierDetection         *OutlierDetection `json:"outlier_detection,omitempty"`
}

// CircuitBreaker definition
// See: https://lyft.github.io/envoy/docs/configuration/cluster_manager/cluster_circuit_breakers.html#circuit-breakers
type CircuitBreaker struct {
	MaxConnections    int `json:"max_connections,omitempty"`
	MaxPendingRequest int `json:"max_pending_requests,omitempty"`
	MaxRequests       int `json:"max_requests,omitempty"`
	MaxRetries        int `json:"max_retries,omitempty"`
}

// OutlierDetection definition
// See: https://lyft.github.io/envoy/docs/configuration/cluster_manager/cluster_runtime.html#outlier-detection
type OutlierDetection struct {
	ConsecutiveError   int `json:"consecutive_5xx,omitempty"`
	IntervalMS         int `json:"interval_ms,omitempty"`
	BaseEjectionTimeMS int `json:"base_ejection_time_ms,omitempty"`
	MaxEjectionPercent int `json:"max_ejection_percent,omitempty"`
}

// ListenersByPort sorts listeners by port
type ListenersByPort []Listener

func (l ListenersByPort) Len() int {
	return len(l)
}

func (l ListenersByPort) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l ListenersByPort) Less(i, j int) bool {
	return l[i].Port < l[j].Port
}

// ClustersByName sorts clusters by name
type ClustersByName []Cluster

func (s ClustersByName) Len() int {
	return len(s)
}

func (s ClustersByName) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s ClustersByName) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

// HostsByName sorts clusters by name
type HostsByName []VirtualHost

func (s HostsByName) Len() int {
	return len(s)
}

func (s HostsByName) Swap(i, j int) {
	s[i], s[j] = s[j], s[i]
}

func (s HostsByName) Less(i, j int) bool {
	return s[i].Name < s[j].Name
}

// RoutesByCluster sorts routes by name
type RoutesByCluster []Route

func (r RoutesByCluster) Len() int {
	return len(r)
}

func (r RoutesByCluster) Swap(i, j int) {
	r[i], r[j] = r[j], r[i]
}

func (r RoutesByCluster) Less(i, j int) bool {
	return r[i].Cluster < r[j].Cluster
}

// SDS is a service discovery service definition
type SDS struct {
	Cluster        Cluster `json:"cluster"`
	RefreshDelayMs int     `json:"refresh_delay_ms"`
}

// ClusterManager definition
type ClusterManager struct {
	Clusters []Cluster `json:"clusters"`
	SDS      SDS       `json:"sds"`
}

// ByName implements sort
type ByName []Cluster

// Len length
func (a ByName) Len() int {
	return len(a)
}

// Swap elements
func (a ByName) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

// Less compare
func (a ByName) Less(i, j int) bool {
	return a[i].Name < a[j].Name
}

//ByHost implement sort
type ByHost []Host

// Len length
func (a ByHost) Len() int {
	return len(a)
}

// Swap elements
func (a ByHost) Swap(i, j int) {
	a[i], a[j] = a[j], a[i]
}

// Less compare
func (a ByHost) Less(i, j int) bool {
	return a[i].URL < a[j].URL
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/andres-erbsen/rrtcp/clockprinter"
	"github.com/andres-erbsen/rrtcp/clockstation"
	"github.com/andres-erbsen/rrtcp/fnet"
)

var addr = flag.String("address", "", "address to connect to or listen at")
var listen = flag.Bool("l", false, "bind to the specified address and listen (default: connect)")
var frameSize = flag.Int("s", 1024, "frame size")
var numStreams = flag.Int("n", 5, "number of streams")
var duration = flag.Int("d", -1, "number of seconds to run program for, -1 means run forever")

func main() {
	flag.Parse()
	if len(flag.Args()) != 0 || *listen == false && *addr == "" {
		flag.Usage()
		os.Exit(1)
	}
	var stop chan os.Signal
	stop = make(chan os.Signal, 2)
	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)

	go func() {
		if *duration != -1 {
			timer := time.NewTimer(time.Second * time.Duration(*duration))
			// Wait for the timer to end, then give the stop signal
			<-timer.C
			stop <- syscall.SIGINT
		}
	}()

	if *listen {
		err := listener(*frameSize, *numStreams, *addr, stop)
		if err != nil {
			os.Exit(2)
		}
	} else {
		err := dialer(*frameSize, *numStreams, *addr, stop)
		if err != nil {
			os.Exit(2)
		}
	}
}

func listener(frameSize int, numStreams int, addr string, stop chan os.Signal) error {
	var cs *clockstation.ClockStation
	cs_running := false

	// Handle stop signals
	done := make(chan bool, 1)
	go func() {
		<-stop
		fmt.Println(cs)
		// TODO: Not sure if this boolean is the best solution
		if cs_running {
			cs.Stop()
		}
		done <- true
		fmt.Println("Stopped listener.")
	}()

	ln, err := net.Listen("tcp", addr)
	defer ln.Close()
	if err != nil {
		fmt.Fprintf(os.Stderr, "net.Listen(%q): %s\n", addr, err.Error())
		return err
	}
	rr := fnet.NewRoundRobin(frameSize)

	for i := 0; i < numStreams; i++ {
		c, err := ln.Accept()
		if err != nil {
			fmt.Fprintf(os.Stderr, "ln.Accept(): %s\n", err.Error())
			return err
		}
		fs := fnet.FromOrderedStream(c, frameSize)
		rr.AddConn(&fs)
	}
	fc := fnet.FrameConn(rr)
	defer fc.Stop()

	cs = clockstation.NewStation(fc, time.Tick(50*time.Millisecond))
	cs_running = true
	err = cs.Run(frameSize)
	if err != nil {
		fmt.Fprintf(os.Stderr, "clockstation.Run: %s\n", err.Error())
		return err
	}
	// Wait for listener to be stopped before returning
	<-done
	return nil
}

func dialer(frameSize int, numStreams int, addr string, stop chan os.Signal) error {
	var fc fnet.FrameConn
	fc_started := false

	// Handle stop signals
	go func() {
		<-stop
		// TODO: This should be a defer instead, when .Stop() can be called more than once freely
		if fc_started {
			fc.Stop()
		}
		fmt.Println("Stopped dialer.")
	}()

	rr := fnet.NewRoundRobin(frameSize)
	for i := 0; i < numStreams; i++ {
		c, err := net.Dial("tcp", addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "net.Dial(%q): %s\n", addr, err.Error())
			return err
		}
		fs := fnet.FromOrderedStream(c, frameSize)
		rr.AddConn(&fs)
	}
	fc = fnet.FrameConn(rr)
	fc_started = true

	err := clockprinter.Run(fc)
	if err != nil {
		fmt.Fprintf(os.Stderr, "clockprinter.Run: %s\n", err.Error())
		return err
	}
	return nil
}
<commit_msg>Fixed rrtcp (everything is context-style now)<commit_after>package main

import (
	"flag"
	"fmt"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"golang.org/x/net/context"

	"github.com/andres-erbsen/rrtcp/clockprinter"
	"github.com/andres-erbsen/rrtcp/clockstation"
	"github.com/andres-erbsen/rrtcp/fnet"
)

var addr = flag.String("address", "", "address to connect to or listen at")
var listen = flag.Bool("l", false, "bind to the specified address and listen (default: connect)")
var frameSize = flag.Int("s", 1024, "frame size")
var numStreams = flag.Int("n", 5, "number of streams")
var duration = flag.Duration("d", 0, "duration to run program for")

func cancelOnSignal(ctx context.Context, sig ...os.Signal) context.Context {
	signalCh := make(chan os.Signal)
	signal.Notify(signalCh, sig...)

	ctx2, cancel := context.WithCancel(ctx)

	go func() {
		select {
		case <-ctx2.Done():
			return
		case <-signalCh:
			cancel()
		}
	}()

	return ctx2
}

func main() {
	flag.Parse()
	if len(flag.Args()) != 0 || *listen == false && *addr == "" {
		flag.Usage()
		os.Exit(1)
	}

	ctx := cancelOnSignal(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	if duration != nil && *duration != time.Duration(0) {
		ctx, _ = context.WithTimeout(ctx, *duration)
	}

	if *listen {
		err := listener(ctx, *frameSize, *numStreams, *addr)
		if err != nil {
			os.Exit(2)
		}
	} else {
		err := dialer(ctx, *frameSize, *numStreams, *addr)
		if err != nil {
			os.Exit(2)
		}
	}
}

func listener(ctx context.Context, frameSize int, numStreams int, addr string) error {
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		fmt.Fprintf(os.Stderr, "net.Listen(%q): %s\n", addr, err.Error())
		return err
	}
	go func() {
		<-ctx.Done()
		ln.Close()
	}()

	rr := fnet.NewRoundRobin(frameSize)
	go func() {
		<-ctx.Done()
		rr.Stop()
	}()

	for i := 0; i < numStreams; i++ {
		c, err := ln.Accept()
		if err != nil {
			fmt.Fprintf(os.Stderr, "ln.Accept(): %s\n", err.Error())
			return err
		}
		fs := fnet.FromOrderedStream(c, frameSize)
		rr.AddConn(&fs)
	}
	fc := fnet.FrameConn(rr)

	if err = clockstation.Run(ctx, fc, time.Tick(50*time.Millisecond)); err != nil {
		return fmt.Errorf("clockstation.Run: %s\n", err.Error())
	}
	return nil
}

func dialer(ctx context.Context, frameSize int, numStreams int, addr string) error {
	rr := fnet.NewRoundRobin(frameSize)
	go func() {
		<-ctx.Done()
		rr.Stop()
	}()

	for i := 0; i < numStreams; i++ {
		c, err := net.Dial("tcp", addr)
		if err != nil {
			fmt.Fprintf(os.Stderr, "net.Dial(%q): %s\n", addr, err.Error())
			return err
		}
		fs := fnet.FromOrderedStream(c, frameSize)
		rr.AddConn(&fs)
	}
	fc := fnet.FrameConn(rr)

	if err := clockprinter.Run(ctx, fc); err != nil {
		fmt.Errorf("clockprinter.Run: %s\n", err.Error())
	}
	return nil
}
<|endoftext|>