<commit_before>// Copyright 2015-2016 trivago GmbH
//
// 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 consumer

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"io/ioutil"
	"strconv"
	"bytes"
	"strings"
	"sync"
	"time"
)

const (
	kinesisCredentialEnv    = "environment"
	kinesisCredentialStatic = "static"
	kinesisCredentialShared = "shared"
	kinesisCredentialNone   = "none"
	kinesisOffsetNewest     = "newest"
	kinesisOffsetOldest     = "oldest"
)

// Kinesis consumer plugin
// This consumer reads message from an AWS Kinesis stream.
// When attached to a fuse, this consumer will stop processing messages in case
// that fuse is burned.
// Configuration example
//
//  - "consumer.Kinesis":
//    KinesisStream: "default"
//    Region: "eu-west-1"
//    Endpoint: "kinesis.eu-west-1.amazonaws.com"
//    DefaultOffset: "Newest"
//    OffsetFile: ""
//    RecordsPerQuery: 100
//    RecordMessageDelimiter: ""
//    QuerySleepTimeMs: 1000
//    RetrySleepTimeSec: 4
//    CredentialType: "none"
//    CredentialId: ""
//    CredentialToken: ""
//    CredentialSecret: ""
//    CredentialFile: ""
//    CredentialProfile: ""
//
// KinesisStream defines the stream to read from.
// By default this is set to "default"
//
// Region defines the amazon region of your kinesis stream.
// By default this is set to "eu-west-1".
//
// Endpoint defines the amazon endpoint for your kinesis stream.
// By default this is et to "kinesis.eu-west-1.amazonaws.com"
//
// CredentialType defines the credentials that are to be used when
// connectiong to kensis. This can be one of the following: environment,
// static, shared, none.
// Static enables the parameters CredentialId, CredentialToken and
// CredentialSecretm shared enables the parameters CredentialFile and
// CredentialProfile. None will not use any credentials and environment
// will pull the credentials from environmental settings.
// By default this is set to none.
//
// DefaultOffset defines the message index to start reading from.
// Valid values are either "Newset", "Oldest", or a number.
// The default value is "Newest".
//
// OffsetFile defines a file to store the current offset per shard.
// By default this is set to "", i.e. it is disabled.
// If a file is set and found consuming will start after the stored
// offset.
//
// RecordsPerQuery defines the number of records to pull per query.
// By default this is set to 100.
//
// RecordMessageDelimiter defines the string to delimit messages within a
// record. By default this is set to "", i.e. it is disabled.
//
// QuerySleepTimeMs defines the number of milliseconds to sleep before
// trying to pull new records from a shard that did not return any records.
// By default this is set to 1000.
//
// RetrySleepTimeSec defines the number of seconds to wait after trying to
// reconnect to a shard. By default this is set to 4.
type Kinesis struct {
	core.ConsumerBase
	client          *kinesis.Kinesis
	config          *aws.Config
	offsets         map[string]string
	stream          string
	offsetType      string
	offsetFile      string
	defaultOffset   string
	recordsPerQuery int64
	delimiter       []byte
	sleepTime       time.Duration
	retryTime       time.Duration
	shardTime       time.Duration
	running         bool
}

func init() {
	shared.TypeRegistry.Register(Kinesis{})
}

// Configure initializes this consumer with values from a plugin config.
func (cons *Kinesis) Configure(conf core.PluginConfig) error {
	err := cons.ConsumerBase.Configure(conf)
	if err != nil {
		return err
	}

	cons.offsets = make(map[string]string)
	cons.stream = conf.GetString("KinesisStream", "default")
	cons.offsetFile = conf.GetString("OffsetFile", "")
	cons.recordsPerQuery = int64(conf.GetInt("RecordsPerQuery", 1000))
	cons.delimiter = []byte(conf.GetString("RecordMessageDelimiter", ""))
	cons.sleepTime = time.Duration(conf.GetInt("QuerySleepTimeMs", 1000)) * time.Millisecond
	cons.retryTime = time.Duration(conf.GetInt("RetrySleepTimeSec", 4)) * time.Second
	// 0 means don't
	cons.shardTime = time.Duration(conf.GetInt("CheckNewShardsSec", 0)) * time.Second

	// Config
	cons.config = aws.NewConfig()
	if endpoint := conf.GetString("Endpoint", "kinesis.eu-west-1.amazonaws.com"); endpoint != "" {
		cons.config.WithEndpoint(endpoint)
	}

	if region := conf.GetString("Region", "eu-west-1"); region != "" {
		cons.config.WithRegion(region)
	}

	// Credentials
	credentialType := strings.ToLower(conf.GetString("CredentialType", kinesisCredentialNone))
	switch credentialType {
	case kinesisCredentialEnv:
		cons.config.WithCredentials(credentials.NewEnvCredentials())

	case kinesisCredentialStatic:
		id := conf.GetString("CredentialId", "")
		token := conf.GetString("CredentialToken", "")
		secret := conf.GetString("CredentialSecret", "")
		cons.config.WithCredentials(credentials.NewStaticCredentials(id, secret, token))

	case kinesisCredentialShared:
		filename := conf.GetString("CredentialFile", "")
		profile := conf.GetString("CredentialProfile", "")
		cons.config.WithCredentials(credentials.NewSharedCredentials(filename, profile))

	case kinesisCredentialNone:
		// Nothing

	default:
		return fmt.Errorf("Unknwon CredentialType: %s", credentialType)
	}

	// Offset
	offsetValue := strings.ToLower(conf.GetString("DefaultOffset", kinesisOffsetNewest))
	switch offsetValue {
	case kinesisOffsetNewest:
		cons.offsetType = kinesis.ShardIteratorTypeLatest
		cons.defaultOffset = ""

	case kinesisOffsetOldest:
		cons.offsetType = kinesis.ShardIteratorTypeTrimHorizon
		cons.defaultOffset = ""

	default:
		cons.offsetType = kinesis.ShardIteratorTypeAtSequenceNumber
		_, err := strconv.ParseUint(offsetValue, 10, 64)
		if err != nil {
			Log.Error.Printf("Default offset must be \"%s\", \"%s\" or a number. %s given", kinesisOffsetNewest, kinesisOffsetOldest, offsetValue)
			offsetValue = "0"
		}
		cons.defaultOffset = offsetValue
	}

	if cons.offsetFile != "" {
		fileContents, err := ioutil.ReadFile(cons.offsetFile)
		if err != nil {
			return err
		}
		if err := json.Unmarshal(fileContents, &cons.offsets); err != nil {
			return err
		}
	}

	return nil
}

func (cons *Kinesis) storeOffsets() {
	if cons.offsetFile != "" {
		fileContents, err := json.Marshal(cons.offsets)
		if err == nil {
			ioutil.WriteFile(cons.offsetFile, fileContents, 0644)
		}
	}
}

func (cons *Kinesis) processShard(shardID string) {
	iteratorConfig := kinesis.GetShardIteratorInput{
		ShardId:                aws.String(shardID),
		ShardIteratorType:      aws.String(cons.offsetType),
		StreamName:             aws.String(cons.stream),
		StartingSequenceNumber: aws.String(cons.offsets[shardID]),
	}
	if *iteratorConfig.StartingSequenceNumber == "" {
		iteratorConfig.StartingSequenceNumber = nil
	} else {
		// starting sequence number requires ShardIteratorTypeAfterSequenceNumber
		iteratorConfig.ShardIteratorType = aws.String(kinesis.ShardIteratorTypeAfterSequenceNumber)
	}

	iterator, err := cons.client.GetShardIterator(&iteratorConfig)
	if err != nil || iterator.ShardIterator == nil {
		Log.Error.Printf("Failed to iterate shard %s:%s - %s", *iteratorConfig.StreamName, *iteratorConfig.ShardId, err.Error())
		if cons.running {
			time.AfterFunc(cons.retryTime, func() { cons.processShard(shardID) })
		}
		return // ### return, retry ###
	}

	recordConfig := kinesis.GetRecordsInput{
		ShardIterator: iterator.ShardIterator,
		Limit:         aws.Int64(cons.recordsPerQuery),
	}

	cons.AddWorker()
	defer cons.WorkerDone()

	for cons.running {
		cons.WaitOnFuse()
		result, err := cons.client.GetRecords(&recordConfig)
		if err != nil {
			Log.Error.Printf("Failed to get records from shard %s:%s - %s", *iteratorConfig.StreamName, *iteratorConfig.ShardId, err.Error())
			// Check if we reached throughput limit
			if AWSerr, isAWSerr := err.(awserr.Error); isAWSerr {
				if AWSerr.Code() == "ProvisionedThroughputExceededException" {
					time.Sleep(5 * time.Second)
				}
			}
		} else {
			if result.NextShardIterator == nil {
				Log.Warning.Printf("Shard %s:%s has been closed", *iteratorConfig.StreamName, *iteratorConfig.ShardId)
				return // ### return, closed ###
			}

			for _, record := range result.Records {
				if record == nil {
					continue // ### continue ###
				}

				seq, _ := strconv.ParseInt(*record.SequenceNumber, 10, 64)
				if len(cons.delimiter) > 0 {
					messages := bytes.Split(record.Data, cons.delimiter)
					for idx, msg := range messages {
						cons.Enqueue([]byte(msg), uint64(seq) + uint64(idx))
					}
				} else {
					cons.Enqueue(record.Data, uint64(seq))
				}
				cons.offsets[*iteratorConfig.ShardId] = *record.SequenceNumber
				cons.storeOffsets()
			}

			recordConfig.ShardIterator = result.NextShardIterator
		}
		time.Sleep(cons.sleepTime)
	}
}

func (cons *Kinesis) connect() error {
	cons.client = kinesis.New(session.New(cons.config))

	// Get shard ids for stream
	streamQuery := &kinesis.DescribeStreamInput{
		StreamName: aws.String(cons.stream),
	}

	streamInfo, err := cons.client.DescribeStream(streamQuery)
	if err != nil {
		return err
	}

	if streamInfo.StreamDescription == nil {
		return fmt.Errorf("StreamDescription could not be retrieved.")
	}

	cons.running = true
	for _, shard := range streamInfo.StreamDescription.Shards {
		if shard.ShardId == nil {
			return fmt.Errorf("ShardId could not be retrieved.")
		}

		if _, offsetStored := cons.offsets[*shard.ShardId]; !offsetStored {
			cons.offsets[*shard.ShardId] = cons.defaultOffset
		}
	}

	for shardID := range cons.offsets {
		Log.Debug.Printf("Starting kinesis consumer for %s:%s", cons.stream, shardID)
		go cons.processShard(shardID)
	}

	if cons.shardTime {
		return nil
	}

	for cons.running {
		streamInfo, err = cons.client.DescribeStream(streamQuery)
		if err != nil {
			return err
		}

		if streamInfo.StreamDescription == nil {
			Log.Warning.Printf("StreamDescription could not be retrieved.")
			continue
		}

		cons.running = true
		for _, shard := range streamInfo.StreamDescription.Shards {
			if shard.ShardId == nil {
				Log.Warning.Printf("ShardId could not be retrieved.")
				continue
			}

			if _, offsetStored := cons.offsets[*shard.ShardId]; !offsetStored {
				cons.offsets[*shard.ShardId] = cons.defaultOffset
				Log.Debug.Printf("Starting kinesis consumer for %s:%s", cons.stream, *shard.ShardId)
				go cons.processShard(*shard.ShardId)
			}
		}
		time.Sleep(cons.shardTime)
	}

	return nil
}

func (cons *Kinesis) close() {
	cons.running = false
	cons.WorkerDone()
}

// Consume listens to stdin.
func (cons *Kinesis) Consume(workers *sync.WaitGroup) {
	cons.AddMainWorker(workers)
	defer cons.close()

	if err := cons.connect(); err != nil {
		Log.Error.Print("Kinesis connection error: ", err)
	} else {
		cons.ControlLoop()
	}
}
<commit_msg>use time.AfterFunc to update kinesis shards<commit_after>// Copyright 2015-2016 trivago GmbH
//
// 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 consumer

import (
	"encoding/json"
	"fmt"
	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/kinesis"
	"github.com/trivago/gollum/core"
	"github.com/trivago/gollum/core/log"
	"github.com/trivago/gollum/shared"
	"io/ioutil"
	"strconv"
	"bytes"
	"strings"
	"sync"
	"time"
)

const (
	kinesisCredentialEnv    = "environment"
	kinesisCredentialStatic = "static"
	kinesisCredentialShared = "shared"
	kinesisCredentialNone   = "none"
	kinesisOffsetNewest     = "newest"
	kinesisOffsetOldest     = "oldest"
)

// Kinesis consumer plugin
// This consumer reads message from an AWS Kinesis stream.
// When attached to a fuse, this consumer will stop processing messages in case
// that fuse is burned.
// Configuration example
//
//  - "consumer.Kinesis":
//    KinesisStream: "default"
//    Region: "eu-west-1"
//    Endpoint: "kinesis.eu-west-1.amazonaws.com"
//    DefaultOffset: "Newest"
//    OffsetFile: ""
//    RecordsPerQuery: 100
//    RecordMessageDelimiter: ""
//    QuerySleepTimeMs: 1000
//    RetrySleepTimeSec: 4
//    CredentialType: "none"
//    CredentialId: ""
//    CredentialToken: ""
//    CredentialSecret: ""
//    CredentialFile: ""
//    CredentialProfile: ""
//
// KinesisStream defines the stream to read from.
// By default this is set to "default"
//
// Region defines the amazon region of your kinesis stream.
// By default this is set to "eu-west-1".
//
// Endpoint defines the amazon endpoint for your kinesis stream.
// By default this is et to "kinesis.eu-west-1.amazonaws.com"
//
// CredentialType defines the credentials that are to be used when
// connectiong to kensis. This can be one of the following: environment,
// static, shared, none.
// Static enables the parameters CredentialId, CredentialToken and
// CredentialSecretm shared enables the parameters CredentialFile and
// CredentialProfile. None will not use any credentials and environment
// will pull the credentials from environmental settings.
// By default this is set to none.
//
// DefaultOffset defines the message index to start reading from.
// Valid values are either "Newset", "Oldest", or a number.
// The default value is "Newest".
//
// OffsetFile defines a file to store the current offset per shard.
// By default this is set to "", i.e. it is disabled.
// If a file is set and found consuming will start after the stored
// offset.
//
// RecordsPerQuery defines the number of records to pull per query.
// By default this is set to 100.
//
// RecordMessageDelimiter defines the string to delimit messages within a
// record. By default this is set to "", i.e. it is disabled.
//
// QuerySleepTimeMs defines the number of milliseconds to sleep before
// trying to pull new records from a shard that did not return any records.
// By default this is set to 1000.
//
// RetrySleepTimeSec defines the number of seconds to wait after trying to
// reconnect to a shard. By default this is set to 4.
type Kinesis struct {
	core.ConsumerBase
	client          *kinesis.Kinesis
	config          *aws.Config
	offsets         map[string]string
	stream          string
	offsetType      string
	offsetFile      string
	defaultOffset   string
	recordsPerQuery int64
	delimiter       []byte
	sleepTime       time.Duration
	retryTime       time.Duration
	shardTime       time.Duration
	running         bool
}

func init() {
	shared.TypeRegistry.Register(Kinesis{})
}

// Configure initializes this consumer with values from a plugin config.
func (cons *Kinesis) Configure(conf core.PluginConfig) error {
	err := cons.ConsumerBase.Configure(conf)
	if err != nil {
		return err
	}

	cons.offsets = make(map[string]string)
	cons.stream = conf.GetString("KinesisStream", "default")
	cons.offsetFile = conf.GetString("OffsetFile", "")
	cons.recordsPerQuery = int64(conf.GetInt("RecordsPerQuery", 1000))
	cons.delimiter = []byte(conf.GetString("RecordMessageDelimiter", ""))
	cons.sleepTime = time.Duration(conf.GetInt("QuerySleepTimeMs", 1000)) * time.Millisecond
	cons.retryTime = time.Duration(conf.GetInt("RetrySleepTimeSec", 4)) * time.Second
	// 0 means don't
	cons.shardTime = time.Duration(conf.GetInt("CheckNewShardsSec", 0)) * time.Second

	// Config
	cons.config = aws.NewConfig()
	if endpoint := conf.GetString("Endpoint", "kinesis.eu-west-1.amazonaws.com"); endpoint != "" {
		cons.config.WithEndpoint(endpoint)
	}

	if region := conf.GetString("Region", "eu-west-1"); region != "" {
		cons.config.WithRegion(region)
	}

	// Credentials
	credentialType := strings.ToLower(conf.GetString("CredentialType", kinesisCredentialNone))
	switch credentialType {
	case kinesisCredentialEnv:
		cons.config.WithCredentials(credentials.NewEnvCredentials())

	case kinesisCredentialStatic:
		id := conf.GetString("CredentialId", "")
		token := conf.GetString("CredentialToken", "")
		secret := conf.GetString("CredentialSecret", "")
		cons.config.WithCredentials(credentials.NewStaticCredentials(id, secret, token))

	case kinesisCredentialShared:
		filename := conf.GetString("CredentialFile", "")
		profile := conf.GetString("CredentialProfile", "")
		cons.config.WithCredentials(credentials.NewSharedCredentials(filename, profile))

	case kinesisCredentialNone:
		// Nothing

	default:
		return fmt.Errorf("Unknwon CredentialType: %s", credentialType)
	}

	// Offset
	offsetValue := strings.ToLower(conf.GetString("DefaultOffset", kinesisOffsetNewest))
	switch offsetValue {
	case kinesisOffsetNewest:
		cons.offsetType = kinesis.ShardIteratorTypeLatest
		cons.defaultOffset = ""

	case kinesisOffsetOldest:
		cons.offsetType = kinesis.ShardIteratorTypeTrimHorizon
		cons.defaultOffset = ""

	default:
		cons.offsetType = kinesis.ShardIteratorTypeAtSequenceNumber
		_, err := strconv.ParseUint(offsetValue, 10, 64)
		if err != nil {
			Log.Error.Printf("Default offset must be \"%s\", \"%s\" or a number. %s given", kinesisOffsetNewest, kinesisOffsetOldest, offsetValue)
			offsetValue = "0"
		}
		cons.defaultOffset = offsetValue
	}

	if cons.offsetFile != "" {
		fileContents, err := ioutil.ReadFile(cons.offsetFile)
		if err != nil {
			return err
		}
		if err := json.Unmarshal(fileContents, &cons.offsets); err != nil {
			return err
		}
	}

	return nil
}

func (cons *Kinesis) storeOffsets() {
	if cons.offsetFile != "" {
		fileContents, err := json.Marshal(cons.offsets)
		if err == nil {
			ioutil.WriteFile(cons.offsetFile, fileContents, 0644)
		}
	}
}

func (cons *Kinesis) processShard(shardID string) {
	iteratorConfig := kinesis.GetShardIteratorInput{
		ShardId:                aws.String(shardID),
		ShardIteratorType:      aws.String(cons.offsetType),
		StreamName:             aws.String(cons.stream),
		StartingSequenceNumber: aws.String(cons.offsets[shardID]),
	}
	if *iteratorConfig.StartingSequenceNumber == "" {
		iteratorConfig.StartingSequenceNumber = nil
	} else {
		// starting sequence number requires ShardIteratorTypeAfterSequenceNumber
		iteratorConfig.ShardIteratorType = aws.String(kinesis.ShardIteratorTypeAfterSequenceNumber)
	}

	iterator, err := cons.client.GetShardIterator(&iteratorConfig)
	if err != nil || iterator.ShardIterator == nil {
		Log.Error.Printf("Failed to iterate shard %s:%s - %s", *iteratorConfig.StreamName, *iteratorConfig.ShardId, err.Error())
		if cons.running {
			time.AfterFunc(cons.retryTime, func() { cons.processShard(shardID) })
		}
		return // ### return, retry ###
	}

	recordConfig := kinesis.GetRecordsInput{
		ShardIterator: iterator.ShardIterator,
		Limit:         aws.Int64(cons.recordsPerQuery),
	}

	cons.AddWorker()
	defer cons.WorkerDone()

	for cons.running {
		cons.WaitOnFuse()
		result, err := cons.client.GetRecords(&recordConfig)
		if err != nil {
			Log.Error.Printf("Failed to get records from shard %s:%s - %s", *iteratorConfig.StreamName, *iteratorConfig.ShardId, err.Error())
			// Check if we reached throughput limit
			if AWSerr, isAWSerr := err.(awserr.Error); isAWSerr {
				if AWSerr.Code() == "ProvisionedThroughputExceededException" {
					time.Sleep(5 * time.Second)
				}
			}
		} else {
			if result.NextShardIterator == nil {
				Log.Warning.Printf("Shard %s:%s has been closed", *iteratorConfig.StreamName, *iteratorConfig.ShardId)
				return // ### return, closed ###
			}

			for _, record := range result.Records {
				if record == nil {
					continue // ### continue ###
				}

				seq, _ := strconv.ParseInt(*record.SequenceNumber, 10, 64)
				if len(cons.delimiter) > 0 {
					messages := bytes.Split(record.Data, cons.delimiter)
					for idx, msg := range messages {
						cons.Enqueue([]byte(msg), uint64(seq) + uint64(idx))
					}
				} else {
					cons.Enqueue(record.Data, uint64(seq))
				}
				cons.offsets[*iteratorConfig.ShardId] = *record.SequenceNumber
				cons.storeOffsets()
			}

			recordConfig.ShardIterator = result.NextShardIterator
		}
		time.Sleep(cons.sleepTime)
	}
}

func (cons *Kinesis) connect() error {
	cons.client = kinesis.New(session.New(cons.config))

	// Get shard ids for stream
	streamQuery := &kinesis.DescribeStreamInput{
		StreamName: aws.String(cons.stream),
	}

	streamInfo, err := cons.client.DescribeStream(streamQuery)
	if err != nil {
		return err
	}

	if streamInfo.StreamDescription == nil {
		return fmt.Errorf("StreamDescription could not be retrieved.")
	}

	cons.running = true
	for _, shard := range streamInfo.StreamDescription.Shards {
		if shard.ShardId == nil {
			return fmt.Errorf("ShardId could not be retrieved.")
		}

		if _, offsetStored := cons.offsets[*shard.ShardId]; !offsetStored {
			cons.offsets[*shard.ShardId] = cons.defaultOffset
		}
	}

	for shardID := range cons.offsets {
		Log.Debug.Printf("Starting kinesis consumer for %s:%s", cons.stream, shardID)
		go cons.processShard(shardID)
	}

	if cons.shardTime > 0 {
		time.AfterFunc(cons.shardTime, cons.updateShards)
	}

	return nil
}

func (cons *Kinesis) updateShards() {
	streamQuery := &kinesis.DescribeStreamInput{
		StreamName: aws.String(cons.stream),
	}

	for cons.running {
		streamInfo, err := cons.client.DescribeStream(streamQuery)
		if err != nil {
			Log.Warning.Printf("StreamInfo could not be retrieved.")
		}

		if streamInfo.StreamDescription == nil {
			Log.Warning.Printf("StreamDescription could not be retrieved.")
			continue
		}

		cons.running = true
		for _, shard := range streamInfo.StreamDescription.Shards {
			if shard.ShardId == nil {
				Log.Warning.Printf("ShardId could not be retrieved.")
				continue
			}

			if _, offsetStored := cons.offsets[*shard.ShardId]; !offsetStored {
				cons.offsets[*shard.ShardId] = cons.defaultOffset
				Log.Debug.Printf("Starting kinesis consumer for %s:%s", cons.stream, *shard.ShardId)
				go cons.processShard(*shard.ShardId)
			}
		}
		time.Sleep(cons.shardTime)
	}
}

func (cons *Kinesis) close() {
	cons.running = false
	cons.WorkerDone()
}

// Consume listens to stdin.
func (cons *Kinesis) Consume(workers *sync.WaitGroup) {
	cons.AddMainWorker(workers)
	defer cons.close()

	if err := cons.connect(); err != nil {
		Log.Error.Print("Kinesis connection error: ", err)
	} else {
		cons.ControlLoop()
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Kubernetes 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 cloudup

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"k8s.io/kops/cloudmock/aws/mockec2"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
)

func Test_TemplateFunctions_CloudControllerConfigArgv(t *testing.T) {
	tests := []struct {
		desc          string
		cluster       *kops.Cluster
		expectedArgv  []string
		expectedError error
	}{
		{
			desc: "Default Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider:                  string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{},
			}},
			expectedArgv: []string{
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Log Level Configuration",
			cluster: &kops.Cluster{
				Spec: kops.ClusterSpec{
					CloudProvider: string(kops.CloudProviderOpenstack),
					ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
						LogLevel: 3,
					},
				},
			},
			expectedArgv: []string{
				"--v=3",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "ExternalCloudControllerManager CloudProvider Configuration",
			cluster: &kops.Cluster{
				Spec: kops.ClusterSpec{
					ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
						CloudProvider: string(kops.CloudProviderOpenstack),
						LogLevel:      3,
					},
				},
			},
			expectedArgv: []string{
				"--cloud-provider=openstack",
				"--v=3",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "No CloudProvider Configuration",
			cluster: &kops.Cluster{
				Spec: kops.ClusterSpec{
					ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
						LogLevel: 3,
					},
				},
			},
			expectedError: fmt.Errorf("Cloud Provider is not set"),
		},
		{
			desc: "k8s cluster name",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					ClusterName: "k8s",
				},
			}},
			expectedArgv: []string{
				"--cluster-name=k8s",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Default Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					Master: "127.0.0.1",
				},
			}},
			expectedArgv: []string{
				"--master=127.0.0.1",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Cluster-cidr Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					ClusterCIDR: "10.0.0.0/24",
				},
			}},
			expectedArgv: []string{
				"--cluster-cidr=10.0.0.0/24",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "AllocateNodeCIDRs Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					AllocateNodeCIDRs: fi.Bool(true),
				},
			}},
			expectedArgv: []string{
				"--allocate-node-cidrs=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "ConfigureCloudRoutes Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					ConfigureCloudRoutes: fi.Bool(true),
				},
			}},
			expectedArgv: []string{
				"--configure-cloud-routes=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "CIDRAllocatorType Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					CIDRAllocatorType: fi.String("RangeAllocator"),
				},
			}},
			expectedArgv: []string{
				"--cidr-allocator-type=RangeAllocator",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "CIDRAllocatorType Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					UseServiceAccountCredentials: fi.Bool(false),
				},
			}},
			expectedArgv: []string{
				"--use-service-account-credentials=false",
				"--v=2",
				"--cloud-provider=openstack",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Leader Election",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					LeaderElection: &kops.LeaderElectionConfiguration{LeaderElect: fi.Bool(true)},
				},
			}},
			expectedArgv: []string{
				"--leader-elect=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Leader Migration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					LeaderElection:        &kops.LeaderElectionConfiguration{LeaderElect: fi.Bool(true)},
					EnableLeaderMigration: fi.Bool(true),
				},
			}},
			expectedArgv: []string{
				"--enable-leader-migration=true",
				"--leader-elect=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
	}
	for _, testCase := range tests {
		t.Run(testCase.desc, func(t *testing.T) {
			tf := &TemplateFunctions{}
			tf.Cluster = testCase.cluster

			actual, error := tf.CloudControllerConfigArgv()
			if !reflect.DeepEqual(error, testCase.expectedError) {
				t.Errorf("Error differs: %+v instead of %+v", error, testCase.expectedError)
			}
			if !reflect.DeepEqual(actual, testCase.expectedArgv) {
				t.Errorf("Argv differs: %+v instead of %+v", actual, testCase.expectedArgv)
			}
		})
	}
}

func Test_KarpenterInstanceTypes(t *testing.T) {
	amiId := "ami-073c8c0760395aab8"
	ec2Client := &mockec2.MockEC2{}
	ec2Client.Images = append(ec2Client.Images, &ec2.Image{
		CreationDate:   aws.String("2016-10-21T20:07:19.000Z"),
		ImageId:        &amiId,
		Name:           aws.String("focal"),
		OwnerId:        aws.String(awsup.WellKnownAccountUbuntu),
		RootDeviceName: aws.String("/dev/xvda"),
		Architecture:   aws.String("x86_64"),
	})
	ig := kops.InstanceGroupSpec{
		Image: amiId,
	}
	cloud := &awsup.MockAWSCloud{MockCloud: awsup.MockCloud{
		MockEC2: ec2Client,
	}}
	_, err := karpenterInstanceTypes(cloud, ig)
	if err != nil {
		t.Errorf("failed to fetch instance types: %v", err)
	}
}
<commit_msg>test case for disabling a controller.<commit_after>/*
Copyright 2019 The Kubernetes 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 cloudup

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/ec2"
	"k8s.io/kops/cloudmock/aws/mockec2"
	"k8s.io/kops/pkg/apis/kops"
	"k8s.io/kops/upup/pkg/fi"
	"k8s.io/kops/upup/pkg/fi/cloudup/awsup"
)

func Test_TemplateFunctions_CloudControllerConfigArgv(t *testing.T) {
	tests := []struct {
		desc          string
		cluster       *kops.Cluster
		expectedArgv  []string
		expectedError error
	}{
		{
			desc: "Default Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider:                  string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{},
			}},
			expectedArgv: []string{
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Log Level Configuration",
			cluster: &kops.Cluster{
				Spec: kops.ClusterSpec{
					CloudProvider: string(kops.CloudProviderOpenstack),
					ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
						LogLevel: 3,
					},
				},
			},
			expectedArgv: []string{
				"--v=3",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "ExternalCloudControllerManager CloudProvider Configuration",
			cluster: &kops.Cluster{
				Spec: kops.ClusterSpec{
					ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
						CloudProvider: string(kops.CloudProviderOpenstack),
						LogLevel:      3,
					},
				},
			},
			expectedArgv: []string{
				"--cloud-provider=openstack",
				"--v=3",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "No CloudProvider Configuration",
			cluster: &kops.Cluster{
				Spec: kops.ClusterSpec{
					ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
						LogLevel: 3,
					},
				},
			},
			expectedError: fmt.Errorf("Cloud Provider is not set"),
		},
		{
			desc: "k8s cluster name",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					ClusterName: "k8s",
				},
			}},
			expectedArgv: []string{
				"--cluster-name=k8s",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Default Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					Master: "127.0.0.1",
				},
			}},
			expectedArgv: []string{
				"--master=127.0.0.1",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Cluster-cidr Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					ClusterCIDR: "10.0.0.0/24",
				},
			}},
			expectedArgv: []string{
				"--cluster-cidr=10.0.0.0/24",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "AllocateNodeCIDRs Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					AllocateNodeCIDRs: fi.Bool(true),
				},
			}},
			expectedArgv: []string{
				"--allocate-node-cidrs=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "ConfigureCloudRoutes Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					ConfigureCloudRoutes: fi.Bool(true),
				},
			}},
			expectedArgv: []string{
				"--configure-cloud-routes=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "CIDRAllocatorType Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					CIDRAllocatorType: fi.String("RangeAllocator"),
				},
			}},
			expectedArgv: []string{
				"--cidr-allocator-type=RangeAllocator",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "CIDRAllocatorType Configuration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					UseServiceAccountCredentials: fi.Bool(false),
				},
			}},
			expectedArgv: []string{
				"--use-service-account-credentials=false",
				"--v=2",
				"--cloud-provider=openstack",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Leader Election",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					LeaderElection: &kops.LeaderElectionConfiguration{LeaderElect: fi.Bool(true)},
				},
			}},
			expectedArgv: []string{
				"--leader-elect=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Leader Migration",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					LeaderElection:        &kops.LeaderElectionConfiguration{LeaderElect: fi.Bool(true)},
					EnableLeaderMigration: fi.Bool(true),
				},
			}},
			expectedArgv: []string{
				"--enable-leader-migration=true",
				"--leader-elect=true",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
		{
			desc: "Disable Controller",
			cluster: &kops.Cluster{Spec: kops.ClusterSpec{
				CloudProvider: string(kops.CloudProviderOpenstack),
				ExternalCloudControllerManager: &kops.CloudControllerManagerConfig{
					Controllers: []string{"*", "-nodeipam"},
				},
			}},
			expectedArgv: []string{
				"--controllers=*,-nodeipam",
				"--v=2",
				"--cloud-provider=openstack",
				"--use-service-account-credentials=true",
				"--cloud-config=/etc/kubernetes/cloud.config",
			},
		},
	}
	for _, testCase := range tests {
		t.Run(testCase.desc, func(t *testing.T) {
			tf := &TemplateFunctions{}
			tf.Cluster = testCase.cluster

			actual, error := tf.CloudControllerConfigArgv()
			if !reflect.DeepEqual(error, testCase.expectedError) {
				t.Errorf("Error differs: %+v instead of %+v", error, testCase.expectedError)
			}
			if !reflect.DeepEqual(actual, testCase.expectedArgv) {
				t.Errorf("Argv differs: %+v instead of %+v", actual, testCase.expectedArgv)
			}
		})
	}
}

func Test_KarpenterInstanceTypes(t *testing.T) {
	amiId := "ami-073c8c0760395aab8"
	ec2Client := &mockec2.MockEC2{}
	ec2Client.Images = append(ec2Client.Images, &ec2.Image{
		CreationDate:   aws.String("2016-10-21T20:07:19.000Z"),
		ImageId:        &amiId,
		Name:           aws.String("focal"),
		OwnerId:        aws.String(awsup.WellKnownAccountUbuntu),
		RootDeviceName: aws.String("/dev/xvda"),
		Architecture:   aws.String("x86_64"),
	})
	ig := kops.InstanceGroupSpec{
		Image: amiId,
	}
	cloud := &awsup.MockAWSCloud{MockCloud: awsup.MockCloud{
		MockEC2: ec2Client,
	}}
	_, err := karpenterInstanceTypes(cloud, ig)
	if err != nil {
		t.Errorf("failed to fetch instance types: %v", err)
	}
}
<|endoftext|>


<commit_before>package graph

import (
	"os"
	"testing"

	"github.com/gyuho/goraph/graph/testdata"
)

func TestNewDefaultGraphInterface(t *testing.T) {
	g := NewDefaultGraph()
	t.Log(g.String())
}

func TestNewDefaultGraphFromJSON(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	g, err := newDefaultGraphFromJSON(f, "graph_00")
	if err != nil {
		t.Error(err)
	}
	// fmt.Println(g.String())
	if g.VertexToChildren["C"]["S"] != 9.0 {
		t.Errorf("weight from C to S must be 9.0 but %f", g.VertexToChildren["C"]["S"])
	}
	for _, graph := range testdata.GraphSlice {
		f, err := os.Open("testdata/graph.json")
		if err != nil {
			t.Error(err)
		}
		defer f.Close()
		g, err := newDefaultGraphFromJSON(f, graph.Name)
		if err != nil {
			t.Error(err)
		}
		if len(g.Vertices) != graph.TotalVertexCount {
			t.Errorf("%s | Expected %d but %d", graph.Name, graph.TotalVertexCount, len(g.Vertices))
		}
		for _, elem := range graph.EdgeToWeight {
			weight1, err := g.GetWeight(elem.Nodes[0], elem.Nodes[1])
			if err != nil {
				t.Error(err)
			}
			weight2 := elem.Weight
			if weight1 != weight2 {
				t.Errorf("Expected %f but %f", weight2, weight1)
			}
		}
	}
}

func TestDefaultGraph_GetVertices(t *testing.T) {
	for _, graph := range testdata.GraphSlice {
		f, err := os.Open("testdata/graph.json")
		if err != nil {
			t.Error(err)
		}
		defer f.Close()
		g, err := newDefaultGraphFromJSON(f, graph.Name)
		if err != nil {
			t.Error(err)
		}
		if len(g.GetVertices()) != graph.TotalVertexCount {
			t.Errorf("wrong number of vertices: %s", g)
		}
	}
}

func TestDefaultGraph_Init(t *testing.T) {
	for _, graph := range testdata.GraphSlice {
		f, err := os.Open("testdata/graph.json")
		if err != nil {
			t.Error(err)
		}
		defer f.Close()
		g, err := newDefaultGraphFromJSON(f, graph.Name)
		if err != nil {
			t.Error(err)
		}
		g.Init()
		if len(g.Vertices) != 0 {
			t.Errorf("not initialized: %s", g)
		}
	}
}

func TestDefaultGraph_DeleteVertex(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	g, err := newDefaultGraphFromJSON(f, "graph_01")
	if err != nil {
		t.Error(err)
	}
	if !g.DeleteVertex("D") {
		t.Error("D does not exist in the graph")
	}
	if g.FindVertex("D") {
		t.Errorf("Expected false but %s", g)
	}
	if v, err := g.GetParents("C"); err != nil || len(v) != 1 {
		t.Fatalf("Expected 1 edge incoming to C but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("C"); err != nil || len(v) != 2 {
		t.Fatalf("Expected 2 edges outgoing from C but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("F"); err != nil || len(v) != 2 {
		t.Fatalf("Expected 2 edges outgoing from F but %v\n\n%s", err, g)
	}
	if v, err := g.GetParents("F"); err != nil || len(v) != 2 {
		t.Fatalf("Expected 2 edges incoming to F but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("B"); err != nil || len(v) != 3 {
		t.Fatalf("Expected 3 edges outgoing from B but %v\n\n%s", err, g)
	}
	if v, err := g.GetParents("E"); err != nil || len(v) != 4 {
		t.Fatalf("Expected 4 edges incoming to E but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("E"); err != nil || len(v) != 3 {
		t.Fatalf("Expected 3 edges outgoing from E but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("T"); err != nil || len(v) != 3 {
		t.Fatalf("Expected 3 edges outgoing from T but %v\n\n%s", err, g)
	}
}

func TestDefaultGraph_DeleteEdge(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()

	g, err := newDefaultGraphFromJSON(f, "graph_01")
	if err != nil {
		t.Error(err)
	}

	if err := g.DeleteEdge("B", "D"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetParents("D"); err != nil || len(v) != 4 {
		t.Errorf("Expected 4 edges incoming to D but %v\n\n%s", err, g)
	}

	if err := g.DeleteEdge("B", "C"); err != nil {
		t.Error(err)
	}
	if err := g.DeleteEdge("S", "C"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetChildren("S"); err != nil || len(v) != 2 {
		t.Errorf("Expected 2 edges outgoing from S but %v\n\n%s", err, g)
	}

	if err := g.DeleteEdge("C", "E"); err != nil {
		t.Error(err)
	}
	if err := g.DeleteEdge("E", "D"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetChildren("E"); err != nil || len(v) != 3 {
		t.Errorf("Expected 3 edges outgoing from E but %v\n\n%s", err, g)
	}
	if v, err := g.GetParents("E"); err != nil || len(v) != 3 {
		t.Errorf("Expected 3 edges incoming to E but %v\n\n%s", err, g)
	}

	if err := g.DeleteEdge("F", "E"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetParents("E"); err != nil || len(v) != 2 {
		t.Errorf("Expected 2 edges incoming to E but %v\n\n%s", err, g)
	}
}

func TestDefaultGraph_ReplaceEdge(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	g, err := newDefaultGraphFromJSON(f, "graph_00")
	if err != nil {
		t.Error(err)
	}
	if err := g.ReplaceEdge("C", "S", 1.0); err != nil {
		t.Error(err)
	}
	if v, err := g.GetWeight("C", "S"); err != nil || v != 1.0 {
		t.Errorf("weight from C to S must be 1.0 but %v\n\n%v", err, g)
	}
}
<commit_msg>add more test cases for internal interface<commit_after>package graph

import (
	"os"
	"testing"

	"github.com/gyuho/goraph/graph/testdata"
)

func TestNewDefaultGraphInterface(t *testing.T) {
	g := NewDefaultGraph()
	t.Log(g.String())
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	_, err = NewDefaultGraphFromJSON(f, "graph_00")
	if err != nil {
		t.Error(err)
	}
}

func TestNewDefaultGraphFromJSON(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	g, err := newDefaultGraphFromJSON(f, "graph_00")
	if err != nil {
		t.Error(err)
	}
	// fmt.Println(g.String())
	if g.VertexToChildren["C"]["S"] != 9.0 {
		t.Errorf("weight from C to S must be 9.0 but %f", g.VertexToChildren["C"]["S"])
	}
	for _, graph := range testdata.GraphSlice {
		f, err := os.Open("testdata/graph.json")
		if err != nil {
			t.Error(err)
		}
		defer f.Close()
		g, err := newDefaultGraphFromJSON(f, graph.Name)
		if err != nil {
			t.Error(err)
		}
		if len(g.Vertices) != graph.TotalVertexCount {
			t.Errorf("%s | Expected %d but %d", graph.Name, graph.TotalVertexCount, len(g.Vertices))
		}
		for _, elem := range graph.EdgeToWeight {
			weight1, err := g.GetWeight(elem.Nodes[0], elem.Nodes[1])
			if err != nil {
				t.Error(err)
			}
			weight2 := elem.Weight
			if weight1 != weight2 {
				t.Errorf("Expected %f but %f", weight2, weight1)
			}
		}
	}
}

func TestDefaultGraph_GetVertices(t *testing.T) {
	for _, graph := range testdata.GraphSlice {
		f, err := os.Open("testdata/graph.json")
		if err != nil {
			t.Error(err)
		}
		defer f.Close()
		g, err := newDefaultGraphFromJSON(f, graph.Name)
		if err != nil {
			t.Error(err)
		}
		if len(g.GetVertices()) != graph.TotalVertexCount {
			t.Errorf("wrong number of vertices: %s", g)
		}
	}
}

func TestDefaultGraph_Init(t *testing.T) {
	for _, graph := range testdata.GraphSlice {
		f, err := os.Open("testdata/graph.json")
		if err != nil {
			t.Error(err)
		}
		defer f.Close()
		g, err := newDefaultGraphFromJSON(f, graph.Name)
		if err != nil {
			t.Error(err)
		}
		g.Init()
		if len(g.Vertices) != 0 {
			t.Errorf("not initialized: %s", g)
		}
	}
}

func TestDefaultGraph_DeleteVertex(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	g, err := newDefaultGraphFromJSON(f, "graph_01")
	if err != nil {
		t.Error(err)
	}
	if !g.DeleteVertex("D") {
		t.Error("D does not exist in the graph")
	}
	if g.FindVertex("D") {
		t.Errorf("Expected false but %s", g)
	}
	if v, err := g.GetParents("C"); err != nil || len(v) != 1 {
		t.Fatalf("Expected 1 edge incoming to C but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("C"); err != nil || len(v) != 2 {
		t.Fatalf("Expected 2 edges outgoing from C but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("F"); err != nil || len(v) != 2 {
		t.Fatalf("Expected 2 edges outgoing from F but %v\n\n%s", err, g)
	}
	if v, err := g.GetParents("F"); err != nil || len(v) != 2 {
		t.Fatalf("Expected 2 edges incoming to F but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("B"); err != nil || len(v) != 3 {
		t.Fatalf("Expected 3 edges outgoing from B but %v\n\n%s", err, g)
	}
	if v, err := g.GetParents("E"); err != nil || len(v) != 4 {
		t.Fatalf("Expected 4 edges incoming to E but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("E"); err != nil || len(v) != 3 {
		t.Fatalf("Expected 3 edges outgoing from E but %v\n\n%s", err, g)
	}
	if v, err := g.GetChildren("T"); err != nil || len(v) != 3 {
		t.Fatalf("Expected 3 edges outgoing from T but %v\n\n%s", err, g)
	}
}

func TestDefaultGraph_DeleteEdge(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()

	g, err := newDefaultGraphFromJSON(f, "graph_01")
	if err != nil {
		t.Error(err)
	}

	if err := g.DeleteEdge("B", "D"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetParents("D"); err != nil || len(v) != 4 {
		t.Errorf("Expected 4 edges incoming to D but %v\n\n%s", err, g)
	}

	if err := g.DeleteEdge("B", "C"); err != nil {
		t.Error(err)
	}
	if err := g.DeleteEdge("S", "C"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetChildren("S"); err != nil || len(v) != 2 {
		t.Errorf("Expected 2 edges outgoing from S but %v\n\n%s", err, g)
	}

	if err := g.DeleteEdge("C", "E"); err != nil {
		t.Error(err)
	}
	if err := g.DeleteEdge("E", "D"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetChildren("E"); err != nil || len(v) != 3 {
		t.Errorf("Expected 3 edges outgoing from E but %v\n\n%s", err, g)
	}
	if v, err := g.GetParents("E"); err != nil || len(v) != 3 {
		t.Errorf("Expected 3 edges incoming to E but %v\n\n%s", err, g)
	}

	if err := g.DeleteEdge("F", "E"); err != nil {
		t.Error(err)
	}
	if v, err := g.GetParents("E"); err != nil || len(v) != 2 {
		t.Errorf("Expected 2 edges incoming to E but %v\n\n%s", err, g)
	}
}

func TestDefaultGraph_ReplaceEdge(t *testing.T) {
	f, err := os.Open("testdata/graph.json")
	if err != nil {
		t.Error(err)
	}
	defer f.Close()
	g, err := newDefaultGraphFromJSON(f, "graph_00")
	if err != nil {
		t.Error(err)
	}
	if err := g.ReplaceEdge("C", "S", 1.0); err != nil {
		t.Error(err)
	}
	if v, err := g.GetWeight("C", "S"); err != nil || v != 1.0 {
		t.Errorf("weight from C to S must be 1.0 but %v\n\n%v", err, g)
	}
}
<|endoftext|>


<commit_before>package gexpect

import (
	"bufio"
	"github.com/shavac/gexpect/pty"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"syscall"
	"time"
)

var (
	err error
)

type SubProcess struct {
	Term            *pty.Terminal
	cmd             *exec.Cmd
	DelayBeforeSend time.Duration
	CheckInterval   time.Duration
	Before          []byte
	After           []byte
	Match           []byte
	echo            bool
}

func (sp *SubProcess) Start() (err error) {
	return sp.Term.Start(sp.cmd)
}

func (sp *SubProcess) Close() (err error) {
	sp.Terminate()
	return sp.Term.Close()
}

func (sp *SubProcess) WaitTimeout(d time.Duration) (err error) {
	if sp.echo {
		go func() {
			io.Copy(os.Stdout, sp)
		}()
	}
	execerr := make(chan error, 1)
	if d > 0 {
		go func() {
			time.Sleep(d)
			execerr <- TIMEOUT
		}()
	}
	go func() {
		execerr <- sp.cmd.Wait()
	}()
	return <-execerr
}

func (sp *SubProcess) Wait() error {
	return sp.WaitTimeout(0)
}

func (sp *SubProcess) Terminate() error {
	return sp.cmd.Process.Kill()
}

func (sp *SubProcess) Expect(expreg ...*regexp.Regexp) (matchIndex int, err error) {
	return sp.ExpectTimeout(0, expreg...)
}

func (sp *SubProcess) ExpectTimeout(timeout time.Duration, expreg ...*regexp.Regexp) (matchIndex int, err error) {
	buf := make([]byte, 2048)
	c := make(chan byte, 1)
	checkpoint := make(chan int, 1)
	rerr := make(chan error, 1)
	go func() {
		for {
			if _, err := io.ReadAtLeast(sp, buf, 1); err != nil {
				rerr <- err
				continue
			}
			for _, b := range buf {
				c <- b
			}
			if sp.echo {
				os.Stdout.Write(buf)
				os.Stdout.Sync()
			}
		}
	}()
	if timeout > 0 {
		go func() {
			time.Sleep(timeout)
			rerr <- TIMEOUT
		}()
	}
	go func() {
		for i := 1; ; i++ {
			time.Sleep(sp.CheckInterval)
			checkpoint <- i
		}
	}()
	for {
		select {
		case c1 := <-c:
			buf = append(buf, c1)
		case e := <-rerr:
			sp.Before = append(sp.Before, buf...)
			return -1, e
		case <-checkpoint:
			for idx, re := range expreg {
				if loc := re.FindIndex(buf); loc != nil {
					sp.Match = buf[loc[0]:loc[1]]
					sp.Before = append(sp.Before, buf[0:loc[0]]...)
					buf = make([]byte, 2048)
					return idx, nil
				}
			} // no match
		}
	}
	return -1, nil
}

func (sp *SubProcess) Read(b []byte) (n int, err error) {
	return sp.Term.Read(b)
}

func (sp *SubProcess) Write(b []byte) (n int, err error) {
	time.Sleep(sp.DelayBeforeSend)
	return sp.Term.Write(b)
}

func (sp *SubProcess) Writeln(b []byte) (n int, err error) {
	bn := append(b, []byte("\r\n")...)
	return sp.Write(bn)
}

func (sp *SubProcess) Send(response string) (err error) {
	_, err = sp.Write([]byte(response))
	return
}

func (sp *SubProcess) SendLine(response string) (err error) {
	return sp.Send(response + "\r\n")
}

func (sp *SubProcess) Interact() (err error) {
	return sp.InteractTimeout(0)
}

func (sp *SubProcess) InteractTimeout(d time.Duration) (err error) {
	sp.Write(sp.After)
	sp.After = []byte{}
	oldState, _ := pty.Tcgetattr(os.Stdin)
	pty.SetRaw(os.Stdin)
	defer pty.Tcsetattr(os.Stdin, oldState)
	s := make(chan os.Signal, 1)
	signal.Notify(s, syscall.SIGINT, syscall.SIGWINCH, syscall.SIGTSTP)
	go func() {
		for sig := range s {
			switch sig {
			case syscall.SIGINT:
				sp.Term.SendIntr()
			case syscall.SIGWINCH:
				if x, y, err := pty.GetWinSize(os.Stdout); err == nil {
					sp.Term.SetWinSize(x, y)
				}
			default:
				continue
			}
		}
	}()
	execerr := make(chan error, 1)
	if d > 0 {
		go func() {
			time.Sleep(d)
			execerr <- TIMEOUT
		}()
	}
	go func() {
		execerr <- sp.cmd.Wait()
	}()
	in := make(chan byte, 1)
	stdin := bufio.NewReader(os.Stdin)
	go func() error {
		var b byte
		for {
			if b, err = stdin.ReadByte(); err != nil {
				if err == io.EOF {
					sp.Term.SendEOF()
					continue
				} else {
					return err
				}
			}
			in <- b
		}
	}()
	go func() {
		io.Copy(os.Stdout, sp)
		return
	}()
	for {
		select {
		case err := <-execerr:
			return err
		case b := <-in:
			_, err = sp.Write([]byte{b})
		}
	}
	return
}

func (sp *SubProcess) Echo() {
	sp.echo = true
}

func (sp *SubProcess) NoEcho() {
	sp.echo = false
}

func NewSubProcess(name string, arg ...string) (sp *SubProcess, err error) {
	sp = new(SubProcess)
	if sp.Term, err = pty.NewTerminal(); err != nil {
		return
	}
	sp.cmd = exec.Command(name, arg...)
	sp.DelayBeforeSend = 50 * time.Microsecond
	sp.CheckInterval = time.Microsecond
	if x, y, err := pty.GetWinSize(os.Stdout); err == nil {
		sp.Term.SetWinSize(x, y)

	}
	return
}
<commit_msg>Fixed Expect<commit_after>package gexpect

import (
	"bufio"
	"io"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"syscall"
	"time"

	"github.com/shavac/gexpect/pty"
)

var (
	err error
)

type SubProcess struct {
	Term            *pty.Terminal
	cmd             *exec.Cmd
	DelayBeforeSend time.Duration
	CheckInterval   time.Duration
	Before          []byte
	After           []byte
	Match           []byte
	echo            bool
}

func (sp *SubProcess) Start() (err error) {
	return sp.Term.Start(sp.cmd)
}

func (sp *SubProcess) Close() (err error) {
	sp.Terminate()
	return sp.Term.Close()
}

func (sp *SubProcess) WaitTimeout(d time.Duration) (err error) {
	if sp.echo {
		go func() {
			io.Copy(os.Stdout, sp)
		}()
	}
	execerr := make(chan error, 1)
	if d > 0 {
		go func() {
			time.Sleep(d)
			execerr <- TIMEOUT
		}()
	}
	go func() {
		execerr <- sp.cmd.Wait()
	}()
	return <-execerr
}

func (sp *SubProcess) Wait() error {
	return sp.WaitTimeout(0)
}

func (sp *SubProcess) Terminate() error {
	return sp.cmd.Process.Kill()
}

func (sp *SubProcess) Expect(expreg ...*regexp.Regexp) (matchIndex int, err error) {
	return sp.ExpectTimeout(0, expreg...)
}

func (sp *SubProcess) ExpectTimeout(timeout time.Duration, expreg ...*regexp.Regexp) (matchIndex int, err error) {
	c := make(chan byte, 1)
	checkpoint := make(chan int, 1)
	rerr := make(chan error, 1)
	go func() {
		buf := make([]byte, 2048)
		for {
			if _, err := io.ReadAtLeast(sp, buf, 1); err != nil {
				rerr <- err
				continue
			}
			for _, b := range buf {
				c <- b
			}
			if sp.echo {
				os.Stdout.Write(buf)
				os.Stdout.Sync()
			}
		}
	}()
	if timeout > 0 {
		go func() {
			time.Sleep(timeout)
			rerr <- TIMEOUT
		}()
	}
	go func() {
		for i := 1; ; i++ {
			time.Sleep(sp.CheckInterval)
			checkpoint <- i
		}
	}()
	buf := make([]byte, 0, 2048)
	for {
		select {
		case c1 := <-c:
			buf = append(buf, c1)
		case e := <-rerr:
			sp.Before = append(sp.Before, buf...)
			return -1, e
		case <-checkpoint:
			for idx, re := range expreg {
				if loc := re.FindIndex(buf); loc != nil {
					sp.Match = buf[loc[0]:loc[1]]
					sp.Before = append(sp.Before, buf[0:loc[0]]...)
					buf = make([]byte, 0, 2048)
					return idx, nil
				}
			} // no match
		}
	}
	return -1, nil
}

func (sp *SubProcess) Read(b []byte) (n int, err error) {
	return sp.Term.Read(b)
}

func (sp *SubProcess) Write(b []byte) (n int, err error) {
	time.Sleep(sp.DelayBeforeSend)
	return sp.Term.Write(b)
}

func (sp *SubProcess) Writeln(b []byte) (n int, err error) {
	bn := append(b, []byte("\r\n")...)
	return sp.Write(bn)
}

func (sp *SubProcess) Send(response string) (err error) {
	_, err = sp.Write([]byte(response))
	return
}

func (sp *SubProcess) SendLine(response string) (err error) {
	return sp.Send(response + "\r\n")
}

func (sp *SubProcess) Interact() (err error) {
	return sp.InteractTimeout(0)
}

func (sp *SubProcess) InteractTimeout(d time.Duration) (err error) {
	sp.Write(sp.After)
	sp.After = []byte{}
	oldState, _ := pty.Tcgetattr(os.Stdin)
	pty.SetRaw(os.Stdin)
	defer pty.Tcsetattr(os.Stdin, oldState)
	s := make(chan os.Signal, 1)
	signal.Notify(s, syscall.SIGINT, syscall.SIGWINCH, syscall.SIGTSTP)
	go func() {
		for sig := range s {
			switch sig {
			case syscall.SIGINT:
				sp.Term.SendIntr()
			case syscall.SIGWINCH:
				if x, y, err := pty.GetWinSize(os.Stdout); err == nil {
					sp.Term.SetWinSize(x, y)
				}
			default:
				continue
			}
		}
	}()
	execerr := make(chan error, 1)
	if d > 0 {
		go func() {
			time.Sleep(d)
			execerr <- TIMEOUT
		}()
	}
	go func() {
		execerr <- sp.cmd.Wait()
	}()
	in := make(chan byte, 1)
	stdin := bufio.NewReader(os.Stdin)
	go func() error {
		var b byte
		for {
			if b, err = stdin.ReadByte(); err != nil {
				if err == io.EOF {
					sp.Term.SendEOF()
					continue
				} else {
					return err
				}
			}
			in <- b
		}
	}()
	go func() {
		io.Copy(os.Stdout, sp)
		return
	}()
	for {
		select {
		case err := <-execerr:
			return err
		case b := <-in:
			_, err = sp.Write([]byte{b})
		}
	}
	return
}

func (sp *SubProcess) Echo() {
	sp.echo = true
}

func (sp *SubProcess) NoEcho() {
	sp.echo = false
}

func NewSubProcess(name string, arg ...string) (sp *SubProcess, err error) {
	sp = new(SubProcess)
	if sp.Term, err = pty.NewTerminal(); err != nil {
		return
	}
	sp.cmd = exec.Command(name, arg...)
	sp.DelayBeforeSend = 50 * time.Microsecond
	sp.CheckInterval = time.Microsecond
	if x, y, err := pty.GetWinSize(os.Stdout); err == nil {
		sp.Term.SetWinSize(x, y)

	}
	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2014-2017 Ludovic Fauvet
// Licensed under the MIT license

package daemon

import (
	"fmt"
	"github.com/etix/mirrorbits/database"
	"github.com/etix/mirrorbits/mirrors"
	"github.com/etix/mirrorbits/utils"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	clusterAnnounce = "HELLO"
)

type cluster struct {
	redis *database.Redis

	nodes         []Node
	nodeIndex     int
	nodeTotal     int
	nodesLock     sync.RWMutex
	mirrorsIndex  []string
	stop          chan bool
	wg            sync.WaitGroup
	running       bool
	StartStopLock sync.Mutex
}

type Node struct {
	ID           string
	LastAnnounce int64
}

type ByNodeID []Node

func (n ByNodeID) Len() int           { return len(n) }
func (n ByNodeID) Swap(i, j int)      { n[i], n[j] = n[j], n[i] }
func (n ByNodeID) Less(i, j int) bool { return n[i].ID < n[j].ID }

func NewCluster(r *database.Redis) *cluster {
	c := &cluster{
		redis: r,
		nodes: make([]Node, 0),
		stop:  make(chan bool),
	}
	return c
}

func (c *cluster) Start() {
	c.StartStopLock.Lock()
	defer c.StartStopLock.Unlock()

	if c.running == true {
		return
	}
	log.Debug("Cluster starting...")
	c.running = true
	c.wg.Add(1)
	c.stop = make(chan bool)
	go c.clusterLoop()
}

func (c *cluster) Stop() {
	c.StartStopLock.Lock()
	defer c.StartStopLock.Unlock()

	select {
	case _, _ = <-c.stop:
		return
	default:
		close(c.stop)
		c.wg.Wait()
		c.running = false
		log.Debug("Cluster stopped")
	}
}

func (c *cluster) clusterLoop() {
	clusterChan := make(chan string, 10)
	announceTicker := time.NewTicker(1 * time.Second)

	hostname := utils.Hostname()
	if len(hostname) == 0 {
		hostname = "unknown"
	}
	nodeID := fmt.Sprintf("%s-%05d", hostname, rand.Intn(32000))

	c.refreshNodeList(nodeID, nodeID)
	c.redis.Pubsub.SubscribeEvent(database.CLUSTER, clusterChan)

	for {
		select {
		case <-c.stop:
			c.wg.Done()
			return
		case <-announceTicker.C:
			c.announce()
		case data := <-clusterChan:
			if !strings.HasPrefix(data, clusterAnnounce+" ") {
				// Garbage
				continue
			}
			c.refreshNodeList(data[len(clusterAnnounce)+1:], nodeID)
		}
	}
}

func (c *cluster) announce() {
	r := c.redis.Get()
	database.Publish(r, database.CLUSTER, fmt.Sprintf("%s %s", clusterAnnounce, c.nodeID))
	r.Close()
}

func (c *cluster) refreshNodeList(nodeID, self string) {
	found := false

	c.nodesLock.Lock()

	// Expire unreachable nodes
	for i := 0; i < len(c.nodes); i++ {
		if utils.ElapsedSec(c.nodes[i].LastAnnounce, 5) && c.nodes[i].ID != nodeID && c.nodes[i].ID != self {
			log.Noticef("<- Node %s left the cluster", c.nodes[i].ID)
			c.nodes = append(c.nodes[:i], c.nodes[i+1:]...)
			i--
		} else if c.nodes[i].ID == nodeID {
			found = true
			c.nodes[i].LastAnnounce = time.Now().UTC().Unix()
		}
	}

	// Join new node
	if !found {
		if nodeID != self {
			log.Noticef("-> Node %s joined the cluster", nodeID)
		}
		n := Node{
			ID:           nodeID,
			LastAnnounce: time.Now().UTC().Unix(),
		}
		// TODO use binary search here
		// See https://golang.org/pkg/sort/#Search
		c.nodes = append(c.nodes, n)
		sort.Sort(ByNodeID(c.nodes))
	}

	c.nodeTotal = len(c.nodes)

	// TODO use binary search here
	// See https://golang.org/pkg/sort/#Search
	for i, n := range c.nodes {
		if n.ID == self {
			c.nodeIndex = i
			break
		}
	}

	c.nodesLock.Unlock()
}

func (c *cluster) AddMirror(mirror *mirrors.Mirror) {
	c.nodesLock.Lock()
	c.mirrorsIndex = addMirrorIDToSlice(c.mirrorsIndex, mirror.ID)
	c.nodesLock.Unlock()
}

func (c *cluster) RemoveMirror(mirror *mirrors.Mirror) {
	c.nodesLock.Lock()
	c.mirrorsIndex = removeMirrorIDFromSlice(c.mirrorsIndex, mirror.ID)
	c.nodesLock.Unlock()
}

func (c *cluster) IsHandled(mirrorID string) bool {
	c.nodesLock.RLock()
	index := sort.SearchStrings(c.mirrorsIndex, mirrorID)

	mRange := int(float32(len(c.mirrorsIndex))/float32(c.nodeTotal) + 0.5)
	start := mRange * c.nodeIndex
	c.nodesLock.RUnlock()

	// Check bounding to see if this mirror must be handled by this node.
	// The distribution of the nodes should be balanced except for the last node
	// that could contain one more node.
	if index >= start && (index < start+mRange || c.nodeIndex == c.nodeTotal-1) {
		return true
	}
	return false
}

func removeMirrorIDFromSlice(slice []string, mirrorID string) []string {
	// See https://golang.org/pkg/sort/#SearchStrings
	idx := sort.SearchStrings(slice, mirrorID)
	if idx < len(slice) && slice[idx] == mirrorID {
		slice = append(slice[:idx], slice[idx+1:]...)
	}
	return slice
}

func addMirrorIDToSlice(slice []string, mirrorID string) []string {
	// See https://golang.org/pkg/sort/#SearchStrings
	idx := sort.SearchStrings(slice, mirrorID)
	if idx >= len(slice) || slice[idx] != mirrorID {
		slice = append(slice[:idx], append([]string{mirrorID}, slice[idx:]...)...)
	}
	return slice
}
<commit_msg>cluster: generate the node ID in the constructor<commit_after>// Copyright (c) 2014-2017 Ludovic Fauvet
// Licensed under the MIT license

package daemon

import (
	"fmt"
	"github.com/etix/mirrorbits/database"
	"github.com/etix/mirrorbits/mirrors"
	"github.com/etix/mirrorbits/utils"
	"math/rand"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	clusterAnnounce = "HELLO"
)

type cluster struct {
	redis *database.Redis

	nodeID        string
	nodes         []Node
	nodeIndex     int
	nodeTotal     int
	nodesLock     sync.RWMutex
	mirrorsIndex  []string
	stop          chan bool
	wg            sync.WaitGroup
	running       bool
	StartStopLock sync.Mutex
}

type Node struct {
	ID           string
	LastAnnounce int64
}

type ByNodeID []Node

func (n ByNodeID) Len() int           { return len(n) }
func (n ByNodeID) Swap(i, j int)      { n[i], n[j] = n[j], n[i] }
func (n ByNodeID) Less(i, j int) bool { return n[i].ID < n[j].ID }

func NewCluster(r *database.Redis) *cluster {
	c := &cluster{
		redis: r,
		nodes: make([]Node, 0),
		stop:  make(chan bool),
	}

	hostname := utils.Hostname()
	if len(hostname) == 0 {
		hostname = "unknown"
	}
	c.nodeID = fmt.Sprintf("%s-%05d", hostname, rand.Intn(32000))
	return c
}

func (c *cluster) Start() {
	c.StartStopLock.Lock()
	defer c.StartStopLock.Unlock()

	if c.running == true {
		return
	}
	log.Debug("Cluster starting...")
	c.running = true
	c.wg.Add(1)
	c.stop = make(chan bool)
	go c.clusterLoop()
}

func (c *cluster) Stop() {
	c.StartStopLock.Lock()
	defer c.StartStopLock.Unlock()

	select {
	case _, _ = <-c.stop:
		return
	default:
		close(c.stop)
		c.wg.Wait()
		c.running = false
		log.Debug("Cluster stopped")
	}
}

func (c *cluster) clusterLoop() {
	clusterChan := make(chan string, 10)
	announceTicker := time.NewTicker(1 * time.Second)

	c.refreshNodeList(c.nodeID, c.nodeID)
	c.redis.Pubsub.SubscribeEvent(database.CLUSTER, clusterChan)

	for {
		select {
		case <-c.stop:
			c.wg.Done()
			return
		case <-announceTicker.C:
			c.announce()
		case data := <-clusterChan:
			if !strings.HasPrefix(data, clusterAnnounce+" ") {
				// Garbage
				continue
			}
			c.refreshNodeList(data[len(clusterAnnounce)+1:], c.nodeID)
		}
	}
}

func (c *cluster) announce() {
	r := c.redis.Get()
	database.Publish(r, database.CLUSTER, fmt.Sprintf("%s %s", clusterAnnounce, c.nodeID))
	r.Close()
}

func (c *cluster) refreshNodeList(nodeID, self string) {
	found := false

	c.nodesLock.Lock()

	// Expire unreachable nodes
	for i := 0; i < len(c.nodes); i++ {
		if utils.ElapsedSec(c.nodes[i].LastAnnounce, 5) && c.nodes[i].ID != nodeID && c.nodes[i].ID != self {
			log.Noticef("<- Node %s left the cluster", c.nodes[i].ID)
			c.nodes = append(c.nodes[:i], c.nodes[i+1:]...)
			i--
		} else if c.nodes[i].ID == nodeID {
			found = true
			c.nodes[i].LastAnnounce = time.Now().UTC().Unix()
		}
	}

	// Join new node
	if !found {
		if nodeID != self {
			log.Noticef("-> Node %s joined the cluster", nodeID)
		}
		n := Node{
			ID:           nodeID,
			LastAnnounce: time.Now().UTC().Unix(),
		}
		// TODO use binary search here
		// See https://golang.org/pkg/sort/#Search
		c.nodes = append(c.nodes, n)
		sort.Sort(ByNodeID(c.nodes))
	}

	c.nodeTotal = len(c.nodes)

	// TODO use binary search here
	// See https://golang.org/pkg/sort/#Search
	for i, n := range c.nodes {
		if n.ID == self {
			c.nodeIndex = i
			break
		}
	}

	c.nodesLock.Unlock()
}

func (c *cluster) AddMirror(mirror *mirrors.Mirror) {
	c.nodesLock.Lock()
	c.mirrorsIndex = addMirrorIDToSlice(c.mirrorsIndex, mirror.ID)
	c.nodesLock.Unlock()
}

func (c *cluster) RemoveMirror(mirror *mirrors.Mirror) {
	c.nodesLock.Lock()
	c.mirrorsIndex = removeMirrorIDFromSlice(c.mirrorsIndex, mirror.ID)
	c.nodesLock.Unlock()
}

func (c *cluster) IsHandled(mirrorID string) bool {
	c.nodesLock.RLock()
	index := sort.SearchStrings(c.mirrorsIndex, mirrorID)

	mRange := int(float32(len(c.mirrorsIndex))/float32(c.nodeTotal) + 0.5)
	start := mRange * c.nodeIndex
	c.nodesLock.RUnlock()

	// Check bounding to see if this mirror must be handled by this node.
	// The distribution of the nodes should be balanced except for the last node
	// that could contain one more node.
	if index >= start && (index < start+mRange || c.nodeIndex == c.nodeTotal-1) {
		return true
	}
	return false
}

func removeMirrorIDFromSlice(slice []string, mirrorID string) []string {
	// See https://golang.org/pkg/sort/#SearchStrings
	idx := sort.SearchStrings(slice, mirrorID)
	if idx < len(slice) && slice[idx] == mirrorID {
		slice = append(slice[:idx], slice[idx+1:]...)
	}
	return slice
}

func addMirrorIDToSlice(slice []string, mirrorID string) []string {
	// See https://golang.org/pkg/sort/#SearchStrings
	idx := sort.SearchStrings(slice, mirrorID)
	if idx >= len(slice) || slice[idx] != mirrorID {
		slice = append(slice[:idx], append([]string{mirrorID}, slice[idx:]...)...)
	}
	return slice
}
<|endoftext|>


<commit_before><commit_msg>Fix incorrect package name<commit_after><|endoftext|>


<commit_before>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 ogletest

import (
	"fmt"
	"reflect"

	"github.com/jacobsa/ogletest/srcutil"
)

// Test suites that implement this interface have special meaning to
// RegisterTestSuite.
type SetUpTestSuiteInterface interface {
	// This method will be called exactly once, before the first test method is
	// run. The receiver of this method will be a zero value of the test suite
	// type, and is not shared with any other methods. Use this method to set up
	// any necessary global state shared by all of the test methods.
	SetUpTestSuite()
}

// Test suites that implement this interface have special meaning to
// RegisterTestSuite.
type TearDownTestSuiteInterface interface {
	// This method will be called exactly once, after the last test method is
	// run. The receiver of this method will be a zero value of the test suite
	// type, and is not shared with any other methods. Use this method to clean
	// up after any necessary global state shared by all of the test methods.
	TearDownTestSuite()
}

// Test suites that implement this interface have special meaning to
// Register.
type SetUpInterface interface {
	// This method is called before each test method is invoked, with the same
	// receiver as that test method. At the time this method is invoked, the
	// receiver is a zero value for the test suite type. Use this method for
	// common setup code that works on data not shared across tests.
	SetUp(*TestInfo)
}

// Test suites that implement this interface have special meaning to
// Register.
type TearDownInterface interface {
	// This method is called after each test method is invoked, with the same
	// receiver as that test method. Use this method for common cleanup code that
	// works on data not shared across tests.
	TearDown()
}

// RegisterTestSuite tells ogletest about a test suite containing tests that it
// should run. Any exported method on the type pointed to by the supplied
// prototype value will be treated as test methods, with the exception of the
// methods defined by the following interfaces, which when present are treated
// as described in the documentation for those interfaces:
//
//  *  SetUpTestSuiteInterface
//  *  SetUpInterface
//  *  TearDownInterface
//  *  TearDownTestSuiteInterface
//
// Each test method is invoked on a different receiver, which is initially a
// zero value of the test suite type.
//
// Example:
//
//     // Some value that is needed by the tests but is expensive to compute.
//     var someExpensiveThing uint
//
//     type FooTest struct {
//       // Path to a temporary file used by the tests. Each test gets a
//       // different temporary file.
//       tempFile string
//     }
//     func init() { ogletest.RegisterTestSuite(&FooTest{}) }
//
//     func (t *FooTest) SetUpTestSuite() {
//       someExpensiveThing = ComputeSomeExpensiveThing()
//     }
//
//     func (t *FooTest) SetUp(ti *ogletest.TestInfo) {
//       t.tempFile = CreateTempFile()
//     }
//
//     func (t *FooTest) TearDown() {
//       DeleteTempFile(t.tempFile)
//     }
//
//     func (t *FooTest) FrobinicatorIsSuccessfullyTweaked() {
//       res := DoSomethingWithExpensiveThing(someExpensiveThing, t.tempFile)
//       ExpectThat(res, Equals(true))
//     }
//
func RegisterTestSuite(p interface{}) {
	if p == nil {
		panic("RegisterTestSuite called with nil suite.")
	}

	val := reflect.ValueOf(p)
	typ := val.Type()
	var zeroInstance reflect.Value

	// We will transform to a TestSuite struct.
	suite := TestSuite{}
	suite.Name = typ.Elem().Name()

	zeroInstance = reflect.New(typ.Elem())
	if i, ok := zeroInstance.Interface().(SetUpTestSuiteInterface); ok {
		suite.SetUp = func() { i.SetUpTestSuite() }
	}

	zeroInstance = reflect.New(typ.Elem())
	if i, ok := zeroInstance.Interface().(TearDownTestSuiteInterface); ok {
		suite.TearDown = func() { i.TearDownTestSuite() }
	}

	// Transform a list of test methods for the suite, filtering them to just the
	// ones that we don't need to skip.
	for _, method := range filterMethods(suite.Name, srcutil.GetMethodsInSourceOrder(typ)) {
		var tf TestFunction
		tf.Name = method.Name

		// Create an instance to be operated on by all of the TestFunction's
		// internal functions.
		instance := reflect.New(typ.Elem())

		// Bind the functions to the instance.
		if i, ok := instance.Interface().(SetUpInterface); ok {
			tf.SetUp = func(ti *TestInfo) { i.SetUp(ti) }
		}

		methodCopy := method
		tf.Run = func() { runTestMethod(instance, methodCopy) }

		if i, ok := instance.Interface().(TearDownInterface); ok {
			tf.TearDown = func() { i.TearDown() }
		}

		// Save the TestFunction.
		suite.TestFunctions = append(suite.TestFunctions, tf)
	}

	// Register the suite.
	Register(suite)
}

func runTestMethod(suite reflect.Value, method reflect.Method) {
	if method.Func.Type().NumIn() != 1 {
		panic(fmt.Sprintf(
			"%s: expected 1 args, actually %d.",
			method.Name,
			method.Func.Type().NumIn()))
	}

	method.Func.Call([]reflect.Value{suite})
}

func filterMethods(suiteName string, in []reflect.Method) (out []reflect.Method) {
	for _, m := range in {
		// Skip set up, tear down, and unexported methods.
		if isSpecialMethod(m.Name) || !isExportedMethod(m.Name) {
			continue
		}
	}

	return
}

func isSpecialMethod(name string) bool {
	return (name == "SetUpTestSuite") ||
		(name == "TearDownTestSuite") ||
		(name == "SetUp") ||
		(name == "TearDown")
}

func isExportedMethod(name string) bool {
	return len(name) > 0 && name[0] >= 'A' && name[0] <= 'Z'
}
<commit_msg>Fix missing the tests<commit_after>// Copyright 2011 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// 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 ogletest

import (
	"fmt"
	"reflect"

	"github.com/jacobsa/ogletest/srcutil"
)

// Test suites that implement this interface have special meaning to
// RegisterTestSuite.
type SetUpTestSuiteInterface interface {
	// This method will be called exactly once, before the first test method is
	// run. The receiver of this method will be a zero value of the test suite
	// type, and is not shared with any other methods. Use this method to set up
	// any necessary global state shared by all of the test methods.
	SetUpTestSuite()
}

// Test suites that implement this interface have special meaning to
// RegisterTestSuite.
type TearDownTestSuiteInterface interface {
	// This method will be called exactly once, after the last test method is
	// run. The receiver of this method will be a zero value of the test suite
	// type, and is not shared with any other methods. Use this method to clean
	// up after any necessary global state shared by all of the test methods.
	TearDownTestSuite()
}

// Test suites that implement this interface have special meaning to
// Register.
type SetUpInterface interface {
	// This method is called before each test method is invoked, with the same
	// receiver as that test method. At the time this method is invoked, the
	// receiver is a zero value for the test suite type. Use this method for
	// common setup code that works on data not shared across tests.
	SetUp(*TestInfo)
}

// Test suites that implement this interface have special meaning to
// Register.
type TearDownInterface interface {
	// This method is called after each test method is invoked, with the same
	// receiver as that test method. Use this method for common cleanup code that
	// works on data not shared across tests.
	TearDown()
}

// RegisterTestSuite tells ogletest about a test suite containing tests that it
// should run. Any exported method on the type pointed to by the supplied
// prototype value will be treated as test methods, with the exception of the
// methods defined by the following interfaces, which when present are treated
// as described in the documentation for those interfaces:
//
//  *  SetUpTestSuiteInterface
//  *  SetUpInterface
//  *  TearDownInterface
//  *  TearDownTestSuiteInterface
//
// Each test method is invoked on a different receiver, which is initially a
// zero value of the test suite type.
//
// Example:
//
//     // Some value that is needed by the tests but is expensive to compute.
//     var someExpensiveThing uint
//
//     type FooTest struct {
//       // Path to a temporary file used by the tests. Each test gets a
//       // different temporary file.
//       tempFile string
//     }
//     func init() { ogletest.RegisterTestSuite(&FooTest{}) }
//
//     func (t *FooTest) SetUpTestSuite() {
//       someExpensiveThing = ComputeSomeExpensiveThing()
//     }
//
//     func (t *FooTest) SetUp(ti *ogletest.TestInfo) {
//       t.tempFile = CreateTempFile()
//     }
//
//     func (t *FooTest) TearDown() {
//       DeleteTempFile(t.tempFile)
//     }
//
//     func (t *FooTest) FrobinicatorIsSuccessfullyTweaked() {
//       res := DoSomethingWithExpensiveThing(someExpensiveThing, t.tempFile)
//       ExpectThat(res, Equals(true))
//     }
//
func RegisterTestSuite(p interface{}) {
	if p == nil {
		panic("RegisterTestSuite called with nil suite.")
	}

	val := reflect.ValueOf(p)
	typ := val.Type()
	var zeroInstance reflect.Value

	// We will transform to a TestSuite struct.
	suite := TestSuite{}
	suite.Name = typ.Elem().Name()

	zeroInstance = reflect.New(typ.Elem())
	if i, ok := zeroInstance.Interface().(SetUpTestSuiteInterface); ok {
		suite.SetUp = func() { i.SetUpTestSuite() }
	}

	zeroInstance = reflect.New(typ.Elem())
	if i, ok := zeroInstance.Interface().(TearDownTestSuiteInterface); ok {
		suite.TearDown = func() { i.TearDownTestSuite() }
	}

	// Transform a list of test methods for the suite, filtering them to just the
	// ones that we don't need to skip.
	for _, method := range filterMethods(suite.Name, srcutil.GetMethodsInSourceOrder(typ)) {
		var tf TestFunction
		tf.Name = method.Name

		// Create an instance to be operated on by all of the TestFunction's
		// internal functions.
		instance := reflect.New(typ.Elem())

		// Bind the functions to the instance.
		if i, ok := instance.Interface().(SetUpInterface); ok {
			tf.SetUp = func(ti *TestInfo) { i.SetUp(ti) }
		}

		methodCopy := method
		tf.Run = func() { runTestMethod(instance, methodCopy) }

		if i, ok := instance.Interface().(TearDownInterface); ok {
			tf.TearDown = func() { i.TearDown() }
		}

		// Save the TestFunction.
		suite.TestFunctions = append(suite.TestFunctions, tf)
	}

	// Register the suite.
	Register(suite)
}

func runTestMethod(suite reflect.Value, method reflect.Method) {
	if method.Func.Type().NumIn() != 1 {
		panic(fmt.Sprintf(
			"%s: expected 1 args, actually %d.",
			method.Name,
			method.Func.Type().NumIn()))
	}

	method.Func.Call([]reflect.Value{suite})
}

func filterMethods(suiteName string, in []reflect.Method) (out []reflect.Method) {
	for _, m := range in {
		// Skip set up, tear down, and unexported methods.
		if isSpecialMethod(m.Name) || !isExportedMethod(m.Name) {
			continue
		}

		out = append(out, m)
	}

	return
}

func isSpecialMethod(name string) bool {
	return (name == "SetUpTestSuite") ||
		(name == "TearDownTestSuite") ||
		(name == "SetUp") ||
		(name == "TearDown")
}

func isExportedMethod(name string) bool {
	return len(name) > 0 && name[0] >= 'A' && name[0] <= 'Z'
}
<|endoftext|>


<commit_before>// Copyright 2016 The LUCI 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 coordinator

import (
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/gcloud/pubsub"
	log "go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/logdog/api/endpoints/coordinator/services/v1"

	gcps "cloud.google.com/go/pubsub"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc/metadata"

	"golang.org/x/net/context"
)

// ArchivalPublisher is capable of publishing archival requests.
type ArchivalPublisher interface {
	// Close shutdowns this publisher instance, releasing all its resources.
	Close() error

	// Publish publishes the supplied ArchiveTask.
	Publish(context.Context, *logdog.ArchiveTask) error

	// NewPublishIndex returns a new publish index. Each publish index is unique
	// within its request.
	NewPublishIndex() uint64
}

// PubsubArchivalPublisher is an ArchivalPublisher that uses Cloud Pub/Sub
// for task archival.
type PubsubArchivalPublisher struct {
	// Publisher is the client used to publish messages.
	Publisher pubsub.Publisher

	// AECtx is the AppEngine Context to use for publish operations.
	AECtx context.Context

	// PublishIndexFunc is a function that will return a unique publish index
	// for this request.
	PublishIndexFunc func() uint64
}

// Close implements ArchivalPublisher.
func (p *PubsubArchivalPublisher) Close() error { return nil }

// Publish implements ArchivalPublisher.
func (p *PubsubArchivalPublisher) Publish(c context.Context, t *logdog.ArchiveTask) error {
	d, err := proto.Marshal(t)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to marshal task.")
		return err
	}

	msg := gcps.Message{
		Data: d,
	}

	// We want to bind our cancellation to "c", but need to use our AppEngine
	// Context for the actual Publish call. Handle this with a goroutine that will
	// cancel our AppEngine Context if "c" is cancelled.
	//
	// The goroutine will naturally exit at the end of this function when "AECtx"
	// is cancelled via defer.
	aeCtx, cancelFunc := context.WithCancel(p.AECtx)
	defer cancelFunc()

	go func(aeCtx context.Context) {
		select {
		case <-aeCtx.Done():
		case <-c.Done():
			cancelFunc()
		}
	}(aeCtx)

	// Create a new AppEngine context. Don't pass gRPC metadata to PubSub, since
	// we don't want any caller RPC to be forwarded to the backend service.
	aeCtx = metadata.NewOutgoingContext(aeCtx, nil)

	log.Fields{
		"project": t.Project,
		"hash":    t.Id,
		"key":     t.Key,
	}.Infof(c, "Publishing archival message for stream.")

	// Publishing usually happens immediately.
	// If it's taken more than 15s, something has already gone horribly wrong,
	// so just kill it and try again.
	nc, _ := clock.WithTimeout(c, time.Second*15)
	_, err = p.Publisher.Publish(nc, &msg)
	return err
}

// NewPublishIndex implements ArchivalPublisher.
func (p *PubsubArchivalPublisher) NewPublishIndex() uint64 {
	return p.PublishIndexFunc()
}
<commit_msg>[logdog] Clean up context usage in archival publisher<commit_after>// Copyright 2016 The LUCI 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 coordinator

import (
	"time"

	"go.chromium.org/luci/common/clock"
	"go.chromium.org/luci/common/gcloud/pubsub"
	log "go.chromium.org/luci/common/logging"
	"go.chromium.org/luci/logdog/api/endpoints/coordinator/services/v1"

	gcps "cloud.google.com/go/pubsub"
	"github.com/golang/protobuf/proto"
	"google.golang.org/grpc/metadata"

	"golang.org/x/net/context"
)

// ArchivalPublisher is capable of publishing archival requests.
type ArchivalPublisher interface {
	// Close shutdowns this publisher instance, releasing all its resources.
	Close() error

	// Publish publishes the supplied ArchiveTask.
	Publish(context.Context, *logdog.ArchiveTask) error

	// NewPublishIndex returns a new publish index. Each publish index is unique
	// within its request.
	NewPublishIndex() uint64
}

// PubsubArchivalPublisher is an ArchivalPublisher that uses Cloud Pub/Sub
// for task archival.
type PubsubArchivalPublisher struct {
	// Publisher is the client used to publish messages.
	Publisher pubsub.Publisher

	// AECtx is the AppEngine Context to use for publish operations.
	AECtx context.Context

	// PublishIndexFunc is a function that will return a unique publish index
	// for this request.
	PublishIndexFunc func() uint64
}

// Close implements ArchivalPublisher.
func (p *PubsubArchivalPublisher) Close() error { return nil }

// Publish implements ArchivalPublisher.
func (p *PubsubArchivalPublisher) Publish(c context.Context, t *logdog.ArchiveTask) error {
	d, err := proto.Marshal(t)
	if err != nil {
		log.WithError(err).Errorf(c, "Failed to marshal task.")
		return err
	}

	msg := gcps.Message{
		Data: d,
	}

	// We want to bind our cancellation to "c", but need to use our AppEngine
	// Context for the actual Publish call. Handle this with a goroutine that will
	// cancel our AppEngine Context if "c" is cancelled.
	//
	// The goroutine will naturally exit at the end of this function when "AECtx"
	// is cancelled via defer.
	//
	// Publishing usually happens immediately.
	// If it's taken more than 15s, something has already gone horribly wrong,
	// so just kill it and try again.
	aeCtx, cancelFunc := clock.WithTimeout(p.AECtx, time.Second*15)
	defer cancelFunc()

	go func(aeCtx context.Context) {
		select {
		case <-aeCtx.Done():
		case <-c.Done():
			cancelFunc()
		}
	}(aeCtx)

	// Create a new AppEngine context. Don't pass gRPC metadata to PubSub, since
	// we don't want any caller RPC to be forwarded to the backend service.
	aeCtx = metadata.NewOutgoingContext(aeCtx, nil)

	log.Fields{
		"project": t.Project,
		"hash":    t.Id,
		"key":     t.Key,
	}.Infof(c, "Publishing archival message for stream.")

	_, err = p.Publisher.Publish(aeCtx, &msg)
	return err
}

// NewPublishIndex implements ArchivalPublisher.
func (p *PubsubArchivalPublisher) NewPublishIndex() uint64 {
	return p.PublishIndexFunc()
}
<|endoftext|>


<commit_before>package controller

import (
	_ "crypto/sha1"
	"fmt"
	"github.com/eaciit/acl"
	_ "github.com/eaciit/colony-core/v0"
	"github.com/eaciit/colony-manager/helper"
	"github.com/eaciit/dbox"
	_ "github.com/eaciit/dbox/dbc/csv"
	"github.com/eaciit/knot/knot.v1"
	"github.com/eaciit/toolkit"
	//"gopkg.in/gomail.v2"
	_ "time"
)

type LoginController struct {
	App
	SenderEmail   string
	HostEmail     string
	PortEmail     int
	UserEmail     string
	PasswordEmail string
}

func init() {
	conn, err := dbox.NewConnection("mongo",
		&dbox.ConnectionInfo{"localhost:27017", "valegrab", "", "", toolkit.M{}.Set("timeout", 3)})
	if err != nil {
		fmt.Println(err.Error())
	}
	err = conn.Connect()
	if err != nil {
		fmt.Println(err.Error())
	}

	err = acl.SetDb(conn)
	if err != nil {
		fmt.Println(err.Error())
	}
}
func (l *LoginController) ProcessLogin(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := toolkit.M{}
	err := r.GetPayload(&payload)

	switch {
	case err != nil:
		return helper.CreateResult(false, nil, err.Error())
	case !payload.Has("username") || !payload.Has("password"):
		return helper.CreateResult(false, nil, "username or password not found")
	case payload.Has("username") && len(toolkit.ToString(payload["username"])) == 0:
		return helper.CreateResult(false, nil, "username cannot empty")
	case payload.Has("password") && len(toolkit.ToString(payload["password"])) == 0:
		return helper.CreateResult(false, nil, "password cannot empty")
	}

	sessid, err := acl.Login(toolkit.ToString(payload["username"]), toolkit.ToString(payload["password"]))
	if err != nil {
		return helper.CreateResult(true, "", err.Error())
	}
	r.SetSession("sessionid", sessid)
	return helper.CreateResult(true, toolkit.M{}.Set("sessionid", sessid), "")

}
func (l *LoginController) ResetPassword(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := toolkit.M{}
	err := r.GetPayload(&payload)
	if err != nil {
		fmt.Println(err.Error())
	}
	email := payload.Has("email")
	bashUrl := payload.Has("bashUrl")
	fmt.Println(email)
	fmt.Println(bashUrl)

	/*
		l.SenderEmail = "admin.support@eaciit.com"
		l.HostEmail = "smtp.office365.com"
		l.PortEmail = 587
		l.UserEmail = "admin.support@eaciit.com"
		l.PasswordEmail = "******"

		m := gomail.NewMessage()

		m.SetHeader("From", l.SenderEmail)
		m.SetHeader("To", email)

		m.SetHeader("Subject", "subject")
		m.SetBody("text/html", "message")

		d := gomail.NewPlainDialer(l.HostEmail, l.PortEmail, l.UserEmail, l.PasswordEmail)

		err := d.DialAndSend(m)

		if err != "" {
			return helper.CreateResult(false, nil, err.Error())
		}

	*/
	return helper.CreateResult(false, nil, "ac")
}

/* ==========================================
var payload = {
sessionid:"t7AuS0YIE9w8gOWY22HPJaj1pSxEjBNU",
accesscheck:[""],
accessid:""
};

app.ajaxPost("/acl/authenticate", payload)
============================================= */

func (l *LoginController) Authenticate(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	var iaccenum acl.AccessTypeEnum

	payload := toolkit.M{}
	result := toolkit.M{}
	result.Set("hasaccess", false)

	err := r.GetPayload(&payload)
	switch {
	case err != nil:
		return helper.CreateResult(false, nil, err.Error())
	}

	switch toolkit.TypeName(payload["accesscheck"]) {
	case "[]interface {}":
		for _, val := range payload["accesscheck"].([]interface{}) {
			tacc := acl.GetAccessEnum(toolkit.ToString(val))
			if !acl.Matchaccess(int(tacc), int(iaccenum)) {
				iaccenum += tacc
			}
		}
	default:
		iaccenum = acl.GetAccessEnum(toolkit.ToString(payload["accesscheck"]))
	}
	// toolkit.Println("Type name : ", toolkit.TypeName(payload["accesscheck"]))

	found := acl.HasAccess(toolkit.ToString(payload["sessionid"]),
		acl.IDTypeSession,
		toolkit.ToString(payload["accessid"]),
		iaccenum)

	if found {
		result.Set("hasaccess", found)
	}

	return helper.CreateResult(true, result, "")
}
func CreateLoginController(l *knot.Server) *LoginController {
	var controller = new(LoginController)
	controller.Server = l
	return controller
}
<commit_msg>add func getmenuaccess<commit_after>package controller

import (
	"fmt"
	"github.com/eaciit/acl"
	"github.com/eaciit/colony-core/v0"
	"github.com/eaciit/colony-manager/helper"
	"github.com/eaciit/dbox"
	// _ "github.com/eaciit/dbox/dbc/csv"
	"github.com/eaciit/knot/knot.v1"
	"github.com/eaciit/toolkit"
	//"gopkg.in/gomail.v2"
	_ "reflect"
	_ "time"
)

type LoginController struct {
	App
	SenderEmail   string
	HostEmail     string
	PortEmail     int
	UserEmail     string
	PasswordEmail string
}

func (l *LoginController) prepareconnection() (conn dbox.IConnection, err error) {
	conn, err = dbox.NewConnection("mongo",
		&dbox.ConnectionInfo{"localhost:27017", "valegrab", "", "", toolkit.M{}.Set("timeout", 3)})
	if err != nil {
		return
	}
	err = conn.Connect()
	return
}

func (l *LoginController) InitialSetDatabase() error {
	conn, err := l.prepareconnection()

	if err != nil {
		return err
	}

	err = acl.SetDb(conn)
	if err != nil {
		return err
	}
	return nil
}

func (l *LoginController) GetAccessMenu(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson
	sessionId := r.Session("sessionid", "")
	fmt.Println(sessionId)
	if sessionId != "" {

		sesssion := new(acl.Session)
		err := acl.FindByID(sesssion, sessionId)
		if err != nil {
			return helper.CreateResult(true, nil, err.Error())
		}

		user := new(acl.User)
		err = acl.FindByID(user, sesssion.UserID)
		if err != nil {
			return helper.CreateResult(true, nil, err.Error())
		}

		results := make([]toolkit.M, 0, 0)
		if len(user.Grants) > 0 {
			for _, v := range user.Grants {
				result := toolkit.M{}
				menu := GetMenu(v.AccessID)
				result.Set(v.AccessID, menu)

				results = append(results, result)
			}
			// fmt.Println(result)
			// return helper.CreateResult(true, result, "")
		}
		fmt.Println(results)
		return helper.CreateResult(true, results, "")
	}
	return helper.CreateResult(false, nil, "Please Login !")
}
func GetMenu(accesId string) interface{} {
	var query *dbox.Filter
	query = dbox.Contains("_id", accesId)
	cursor, err := colonycore.Find(new(colonycore.Menu), query)
	data := []colonycore.Menu{}
	err = cursor.Fetch(&data, 0, false)
	if err != nil {
		return helper.CreateResult(false, nil, err.Error())
	}
	return data
}
func (l *LoginController) ProcessLogin(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := toolkit.M{}
	err := r.GetPayload(&payload)

	switch {
	case err != nil:
		return helper.CreateResult(false, nil, err.Error())
	case !payload.Has("username") || !payload.Has("password"):
		return helper.CreateResult(false, nil, "username or password not found")
	case payload.Has("username") && len(toolkit.ToString(payload["username"])) == 0:
		return helper.CreateResult(false, nil, "username cannot empty")
	case payload.Has("password") && len(toolkit.ToString(payload["password"])) == 0:
		return helper.CreateResult(false, nil, "password cannot empty")
	}

	sessid, err := acl.Login(toolkit.ToString(payload["username"]), toolkit.ToString(payload["password"]))
	if err != nil {
		return helper.CreateResult(true, "", err.Error())
	}
	r.SetSession("sessionid", sessid)
	return helper.CreateResult(true, toolkit.M{}.Set("sessionid", sessid), "Login Success")

}

func (l *LoginController) ResetPassword(r *knot.WebContext) interface{} {
	r.Config.OutputType = knot.OutputJson

	payload := toolkit.M{}
	err := r.GetPayload(&payload)
	if err != nil {
		fmt.Println(err.Error())
	}
	email := payload.Has("email")
	url := payload.Has("url")
	fmt.Println(email)
	fmt.Println(url)

	/*
		l.SenderEmail = "admin.support@eaciit.com"
		l.HostEmail = "smtp.office365.com"
		l.PortEmail = 587
		l.UserEmail = "admin.support@eaciit.com"
		l.PasswordEmail = "******"

		m := gomail.NewMessage()

		m.SetHeader("From", l.SenderEmail)
		m.SetHeader("To", email)

		m.SetHeader("Subject", "subject")
		m.SetBody("text/html", "message")

		d := gomail.NewPlainDialer(l.HostEmail, l.PortEmail, l.UserEmail, l.PasswordEmail)

		err := d.DialAndSend(m)

		if err != "" {
			return helper.CreateResult(false, nil, err.Error())
		}

	*/
	return helper.CreateResult(false, nil, "ac")
}
func CreateLoginController(l *knot.Server) *LoginController {
	var controller = new(LoginController)
	controller.Server = l
	return controller
}
<|endoftext|>