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

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

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

package go_kafka_client

import (
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/stealthly/siesta"
	"time"
)

// LowLevelClient is a low-level Kafka client that manages broker connections, responsible to fetch metadata and is able
// to handle Fetch and Offset requests.
//TODO not sure that's a good name for this interface
type LowLevelClient interface {
	// This will be called right after connecting to ConsumerCoordinator so this client can initialize itself
	// with bootstrap broker list for example. May return an error to signal this client is unable to work with given configuration.
	Initialize() error

	// This will be called each time the fetch request to Kafka should be issued. Topic, partition and offset are self-explanatory.
	// Should return a slice of Messages and an error if a fetch error occurred.
	// Note that for performance reasons it makes sense to keep open broker connections and reuse them on every fetch call.
	Fetch(topic string, partition int32, offset int64) ([]*Message, error)

	// This will be called when call to Fetch returns an error. As every client has different error mapping we ask here explicitly
	// whether the returned error is an OffsetOutOfRange error. Should return true if it is, false otherwise.
	IsOffsetOutOfRange(error) bool

	// This will be called to handle OffsetOutOfRange error. OffsetTime will be either "smallest" or "largest".
	// Should return a corresponding offset value and an error if it occurred.
	GetAvailableOffset(topic string, partition int32, offsetTime string) (int64, error)

	// This will be called to gracefully shutdown this client.
	Close()
}

// SaramaClient implements LowLevelClient and uses github.com/Shopify/sarama as underlying implementation.
type SaramaClient struct {
	config *ConsumerConfig
	client sarama.Client
}

// Creates a new SaramaClient using a given ConsumerConfig.
func NewSaramaClient(config *ConsumerConfig) *SaramaClient {
	return &SaramaClient{
		config: config,
	}
}

// Returns a string representation of this SaramaClient.
func (this *SaramaClient) String() string {
	return "Sarama client"
}

// This will be called right after connecting to ConsumerCoordinator so this client can initialize itself
// with bootstrap broker list for example. May return an error to signal this client is unable to work with given configuration.
func (this *SaramaClient) Initialize() error {
	bootstrapBrokers, err := BootstrapBrokers(this.config.Coordinator)
	if err != nil {
		return err
	}

	client, err := sarama.NewClient(bootstrapBrokers, nil)
	if err != nil {
		return err
	}
	this.client = client

	return nil
}

// This will be called each time the fetch request to Kafka should be issued. Topic, partition and offset are self-explanatory.
// Returns slice of Messages and an error if a fetch error occurred.
func (this *SaramaClient) Fetch(topic string, partition int32, offset int64) ([]*Message, error) {
	leader, err := this.client.Leader(topic, partition)
	if err != nil {
		this.client.RefreshMetadata(topic)
		return nil, err
	}

	fetchRequest := new(sarama.FetchRequest)
	fetchRequest.MinBytes = this.config.FetchMinBytes
	fetchRequest.MaxWaitTime = this.config.FetchWaitMaxMs
	Debugf(this, "Adding block: topic=%s, partition=%d, offset=%d, fetchsize=%d", topic, partition, offset, this.config.FetchMessageMaxBytes)
	fetchRequest.AddBlock(topic, partition, offset, this.config.FetchMessageMaxBytes)

	response, err := leader.Fetch(fetchRequest)
	if err != nil {
		this.client.RefreshMetadata(topic)
		return nil, err
	}

	messages := make([]*Message, 0)
	if response != nil {
		Debug(this, "Processing fetch response")
		for topic, partitionAndData := range response.Blocks {
			for partition, data := range partitionAndData {
				switch data.Err {
				case sarama.ErrNoError:
					{
						if len(data.MsgSet.Messages) > 0 {
							this.filterPartitionData(data, offset)
							messages = this.collectMessages(data, topic, partition)
							if this.config.Debug {
								timestamp := time.Now().UnixNano() / int64(time.Millisecond)
								for _, message := range messages {
									message.DecodedKey = []int64{timestamp}
								}
							}
						} else {
							Debugf(this, "No messages in %s:%d at offset %d", topic, partition, offset)
						}
					}
				default:
					{
						this.client.RefreshMetadata(topic)
						return nil, data.Err
					}
				}
			}
		}
	}

	return messages, nil
}

// Checks whether the given error indicates an OffsetOutOfRange error.
func (this *SaramaClient) IsOffsetOutOfRange(err error) bool {
	return err == sarama.ErrOffsetOutOfRange
}

// This will be called to handle OffsetOutOfRange error. OffsetTime will be either "smallest" or "largest".
func (this *SaramaClient) GetAvailableOffset(topic string, partition int32, offsetTime string) (int64, error) {
	time := sarama.OffsetNewest
	if offsetTime == "smallest" {
		time = sarama.OffsetOldest
	}
	offset, err := this.client.GetOffset(topic, partition, time)
	if err != nil {
		return -1, nil
	}

	return offset, nil
}

// Gracefully shuts down this client.
func (this *SaramaClient) Close() {
	this.client.Close()
}

func (this *SaramaClient) filterPartitionData(partitionData *sarama.FetchResponseBlock, requestedOffset int64) {
	lowestCorrectIndex := 0
	for i, v := range partitionData.MsgSet.Messages {
		if v.Offset < requestedOffset {
			lowestCorrectIndex = i + 1
		} else {
			break
		}
	}
	partitionData.MsgSet.Messages = partitionData.MsgSet.Messages[lowestCorrectIndex:]
}

func (this *SaramaClient) collectMessages(partitionData *sarama.FetchResponseBlock, topic string, partition int32) []*Message {
	messages := make([]*Message, 0)

	for _, message := range partitionData.MsgSet.Messages {
		if message.Msg.Set != nil {
			for _, wrapped := range message.Msg.Set.Messages {
				decodedKey, err := this.config.KeyDecoder.Decode(wrapped.Msg.Key)
				if err != nil {
					//TODO: what if we fail to decode the key: fail-fast or fail-safe strategy?
					Error(this, err.Error())
				}
				decodedValue, err := this.config.ValueDecoder.Decode(wrapped.Msg.Value)
				if err != nil {
					//TODO: what if we fail to decode the value: fail-fast or fail-safe strategy?
					Error(this, err.Error())
				}
				messages = append(messages, &Message{
					Key:          wrapped.Msg.Key,
					Value:        wrapped.Msg.Value,
					DecodedKey:   decodedKey,
					DecodedValue: decodedValue,
					Topic:        topic,
					Partition:    partition,
					Offset:       wrapped.Offset,
				})
			}
		} else {
			decodedKey, err := this.config.KeyDecoder.Decode(message.Msg.Key)
			if err != nil {
				//TODO: what if we fail to decode the key: fail-fast or fail-safe strategy?
				Error(this, err.Error())
			}
			decodedValue, err := this.config.ValueDecoder.Decode(message.Msg.Value)
			if err != nil {
				//TODO: what if we fail to decode the value: fail-fast or fail-safe strategy?
				Error(this, err.Error())
			}
			messages = append(messages, &Message{
				Key:          message.Msg.Key,
				Value:        message.Msg.Value,
				DecodedKey:   decodedKey,
				DecodedValue: decodedValue,
				Topic:        topic,
				Partition:    partition,
				Offset:       message.Offset,
			})
		}
	}

	return messages
}

// SiestaClient implements LowLevelClient and OffsetStorage and uses github.com/stealthly/siesta as underlying implementation.
type SiestaClient struct {
	config    *ConsumerConfig
	connector siesta.Connector
}

// Creates a new SiestaClient using a given ConsumerConfig.
func NewSiestaClient(config *ConsumerConfig) *SiestaClient {
	return &SiestaClient{
		config: config,
	}
}

// Returns a string representation of this SaramaClient.
func (this *SiestaClient) String() string {
	return "Siesta client"
}

// This will be called right after connecting to ConsumerCoordinator so this client can initialize itself
// with bootstrap broker list for example. May return an error to signal this client is unable to work with given configuration.
func (this *SiestaClient) Initialize() error {
	bootstrapBrokers, err := BootstrapBrokers(this.config.Coordinator)
	if err != nil {
		return err
	}

	connectorConfig := siesta.NewConnectorConfig()
	connectorConfig.BrokerList = bootstrapBrokers
	connectorConfig.ReadTimeout = this.config.SocketTimeout
	connectorConfig.WriteTimeout = this.config.SocketTimeout
	connectorConfig.ConnectTimeout = this.config.SocketTimeout
	connectorConfig.FetchSize = this.config.FetchMessageMaxBytes
	connectorConfig.ClientId = this.config.Clientid

	this.connector, err = siesta.NewDefaultConnector(connectorConfig)
	if err != nil {
		return err
	}

	return nil
}

// This will be called each time the fetch request to Kafka should be issued. Topic, partition and offset are self-explanatory.
// Returns slice of Messages and an error if a fetch error occurred.
func (this *SiestaClient) Fetch(topic string, partition int32, offset int64) ([]*Message, error) {
	Tracef(this, "Fetching %s %d from %d", topic, partition, offset)
	response, err := this.connector.Fetch(topic, partition, offset)
	if err != nil {
		return nil, err
	}

	messages := make([]*Message, 0)

	timestamp := time.Now().UnixNano() / int64(time.Millisecond)
	collector := func(topic string, partition int32, offset int64, key []byte, value []byte) {
		decodedKey, err := this.config.KeyDecoder.Decode(key)
		if err != nil {
			//TODO: what if we fail to decode the key: fail-fast or fail-safe strategy?
			Error(this, err.Error())
		}
		decodedValue, err := this.config.ValueDecoder.Decode(value)
		if err != nil {
			//TODO: what if we fail to decode the value: fail-fast or fail-safe strategy?
			Error(this, err.Error())
		}

		if this.config.Debug {
			decodedKey = []int64{timestamp}
		}

		messages = append(messages, &Message{
			Key:          key,
			Value:        value,
			DecodedKey:   decodedKey,
			DecodedValue: decodedValue,
			Topic:        topic,
			Partition:    partition,
			Offset:       offset,
		})
	}

	return messages, response.CollectMessages(collector)
}

// Checks whether the given error indicates an OffsetOutOfRange error.
func (this *SiestaClient) IsOffsetOutOfRange(err error) bool {
	return err == siesta.OffsetOutOfRange
}

// This will be called to handle OffsetOutOfRange error. OffsetTime will be either "smallest" or "largest".
func (this *SiestaClient) GetAvailableOffset(topic string, partition int32, offsetTime string) (int64, error) {
	time := siesta.LatestTime
	if offsetTime == "smallest" {
		time = siesta.EarliestTime
	}
	return this.connector.GetAvailableOffset(topic, partition, time)
}

// Gets the offset for a given group, topic and partition.
// May return an error if fails to retrieve the offset.
func (this *SiestaClient) GetOffset(group string, topic string, partition int32) (int64, error) {
	return this.connector.GetOffset(group, topic, partition)
}

// Commits the given offset for a given group, topic and partition.
// May return an error if fails to commit the offset.
func (this *SiestaClient) CommitOffset(group string, topic string, partition int32, offset int64) error {
	return this.connector.CommitOffset(group, topic, partition, offset)
}

// Gracefully shuts down this client.
func (this *SiestaClient) Close() {
	<-this.connector.Close()
}

// BootstrapBrokers queries the ConsumerCoordinator for all known brokers in the cluster to be used later as a bootstrap list for the LowLevelClient.
func BootstrapBrokers(coordinator ConsumerCoordinator) ([]string, error) {
	bootstrapBrokers := make([]string, 0)
	brokers, err := coordinator.GetAllBrokers()
	if err != nil {
		return nil, err
	}
	for _, broker := range brokers {
		bootstrapBrokers = append(bootstrapBrokers, fmt.Sprintf("%s:%d", broker.Host, broker.Port))
	}

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

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

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

package go_kafka_client

import (
	"fmt"
	"time"

	"github.com/Shopify/sarama"
	"github.com/stealthly/siesta"
)

// LowLevelClient is a low-level Kafka client that manages broker connections, responsible to fetch metadata and is able
// to handle Fetch and Offset requests.
//TODO not sure that's a good name for this interface
type LowLevelClient interface {
	// This will be called right after connecting to ConsumerCoordinator so this client can initialize itself
	// with bootstrap broker list for example. May return an error to signal this client is unable to work with given configuration.
	Initialize() error

	// This will be called each time the fetch request to Kafka should be issued. Topic, partition and offset are self-explanatory.
	// Should return a slice of Messages and an error if a fetch error occurred.
	// Note that for performance reasons it makes sense to keep open broker connections and reuse them on every fetch call.
	Fetch(topic string, partition int32, offset int64) ([]*Message, error)

	// This will be called when call to Fetch returns an error. As every client has different error mapping we ask here explicitly
	// whether the returned error is an OffsetOutOfRange error. Should return true if it is, false otherwise.
	IsOffsetOutOfRange(error) bool

	// This will be called to handle OffsetOutOfRange error. OffsetTime will be either "smallest" or "largest".
	// Should return a corresponding offset value and an error if it occurred.
	GetAvailableOffset(topic string, partition int32, offsetTime string) (int64, error)

	// This will be called to gracefully shutdown this client.
	Close()
}

// SaramaClient implements LowLevelClient and uses github.com/Shopify/sarama as underlying implementation.
type SaramaClient struct {
	config *ConsumerConfig
	client sarama.Client
}

// Creates a new SaramaClient using a given ConsumerConfig.
func NewSaramaClient(config *ConsumerConfig) *SaramaClient {
	return &SaramaClient{
		config: config,
	}
}

// Returns a string representation of this SaramaClient.
func (this *SaramaClient) String() string {
	return "Sarama client"
}

// This will be called right after connecting to ConsumerCoordinator so this client can initialize itself
// with bootstrap broker list for example. May return an error to signal this client is unable to work with given configuration.
func (this *SaramaClient) Initialize() error {
	bootstrapBrokers, err := BootstrapBrokers(this.config.Coordinator)
	if err != nil {
		return err
	}

	client, err := sarama.NewClient(bootstrapBrokers, nil)
	if err != nil {
		return err
	}
	this.client = client

	return nil
}

// This will be called each time the fetch request to Kafka should be issued. Topic, partition and offset are self-explanatory.
// Returns slice of Messages and an error if a fetch error occurred.
func (this *SaramaClient) Fetch(topic string, partition int32, offset int64) ([]*Message, error) {
	leader, err := this.client.Leader(topic, partition)
	if err != nil {
		this.client.RefreshMetadata(topic)
		return nil, err
	}

	fetchRequest := new(sarama.FetchRequest)
	fetchRequest.MinBytes = this.config.FetchMinBytes
	fetchRequest.MaxWaitTime = this.config.FetchWaitMaxMs
	Debugf(this, "Adding block: topic=%s, partition=%d, offset=%d, fetchsize=%d", topic, partition, offset, this.config.FetchMessageMaxBytes)
	fetchRequest.AddBlock(topic, partition, offset, this.config.FetchMessageMaxBytes)

	response, err := leader.Fetch(fetchRequest)
	if err != nil {
		this.client.RefreshMetadata(topic)
		return nil, err
	}

	messages := make([]*Message, 0)
	if response != nil {
		Debug(this, "Processing fetch response")
		for topic, partitionAndData := range response.Blocks {
			for partition, data := range partitionAndData {
				switch data.Err {
				case sarama.ErrNoError:
					{
						if len(data.MsgSet.Messages) > 0 {
							this.filterPartitionData(data, offset)
							messages = this.collectMessages(data, topic, partition)
							if this.config.Debug {
								timestamp := time.Now().UnixNano() / int64(time.Millisecond)
								for _, message := range messages {
									message.DecodedKey = []int64{timestamp}
								}
							}
						} else {
							Debugf(this, "No messages in %s:%d at offset %d", topic, partition, offset)
						}
					}
				default:
					{
						this.client.RefreshMetadata(topic)
						return nil, data.Err
					}
				}
			}
		}
	}

	return messages, nil
}

// Checks whether the given error indicates an OffsetOutOfRange error.
func (this *SaramaClient) IsOffsetOutOfRange(err error) bool {
	return err == sarama.ErrOffsetOutOfRange
}

// This will be called to handle OffsetOutOfRange error. OffsetTime will be either "smallest" or "largest".
func (this *SaramaClient) GetAvailableOffset(topic string, partition int32, offsetTime string) (int64, error) {
	time := sarama.OffsetNewest
	if offsetTime == "smallest" {
		time = sarama.OffsetOldest
	}
	offset, err := this.client.GetOffset(topic, partition, time)
	if err != nil {
		return -1, nil
	}

	return offset, nil
}

// Gracefully shuts down this client.
func (this *SaramaClient) Close() {
	this.client.Close()
}

func (this *SaramaClient) filterPartitionData(partitionData *sarama.FetchResponseBlock, requestedOffset int64) {
	lowestCorrectIndex := 0
	for i, v := range partitionData.MsgSet.Messages {
		if v.Offset < requestedOffset {
			lowestCorrectIndex = i + 1
		} else {
			break
		}
	}
	partitionData.MsgSet.Messages = partitionData.MsgSet.Messages[lowestCorrectIndex:]
}

func (this *SaramaClient) collectMessages(partitionData *sarama.FetchResponseBlock, topic string, partition int32) []*Message {
	messages := make([]*Message, 0)

	for _, message := range partitionData.MsgSet.Messages {
		if message.Msg.Set != nil {
			for _, wrapped := range message.Msg.Set.Messages {
				decodedKey, err := this.config.KeyDecoder.Decode(wrapped.Msg.Key)
				if err != nil {
					//TODO: what if we fail to decode the key: fail-fast or fail-safe strategy?
					Error(this, err.Error())
				}
				decodedValue, err := this.config.ValueDecoder.Decode(wrapped.Msg.Value)
				if err != nil {
					//TODO: what if we fail to decode the value: fail-fast or fail-safe strategy?
					Error(this, err.Error())
				}
				messages = append(messages, &Message{
					Key:          wrapped.Msg.Key,
					Value:        wrapped.Msg.Value,
					DecodedKey:   decodedKey,
					DecodedValue: decodedValue,
					Topic:        topic,
					Partition:    partition,
					Offset:       wrapped.Offset,
				})
			}
		} else {
			decodedKey, err := this.config.KeyDecoder.Decode(message.Msg.Key)
			if err != nil {
				//TODO: what if we fail to decode the key: fail-fast or fail-safe strategy?
				Error(this, err.Error())
			}
			decodedValue, err := this.config.ValueDecoder.Decode(message.Msg.Value)
			if err != nil {
				//TODO: what if we fail to decode the value: fail-fast or fail-safe strategy?
				Error(this, err.Error())
			}
			messages = append(messages, &Message{
				Key:          message.Msg.Key,
				Value:        message.Msg.Value,
				DecodedKey:   decodedKey,
				DecodedValue: decodedValue,
				Topic:        topic,
				Partition:    partition,
				Offset:       message.Offset,
			})
		}
	}

	return messages
}

// SiestaClient implements LowLevelClient and OffsetStorage and uses github.com/stealthly/siesta as underlying implementation.
type SiestaClient struct {
	config    *ConsumerConfig
	connector siesta.Connector
}

// Creates a new SiestaClient using a given ConsumerConfig.
func NewSiestaClient(config *ConsumerConfig) *SiestaClient {
	return &SiestaClient{
		config: config,
	}
}

// Returns a string representation of this SaramaClient.
func (this *SiestaClient) String() string {
	return "Siesta client"
}

// This will be called right after connecting to ConsumerCoordinator so this client can initialize itself
// with bootstrap broker list for example. May return an error to signal this client is unable to work with given configuration.
func (this *SiestaClient) Initialize() error {
	bootstrapBrokers, err := BootstrapBrokers(this.config.Coordinator)
	if err != nil {
		return err
	}

	connectorConfig := siesta.NewConnectorConfig()
	connectorConfig.BrokerList = bootstrapBrokers
	connectorConfig.ReadTimeout = this.config.SocketTimeout
	connectorConfig.WriteTimeout = this.config.SocketTimeout
	connectorConfig.ConnectTimeout = this.config.SocketTimeout
	connectorConfig.FetchSize = this.config.FetchMessageMaxBytes
	connectorConfig.ClientID = this.config.Clientid

	this.connector, err = siesta.NewDefaultConnector(connectorConfig)
	if err != nil {
		return err
	}

	return nil
}

// This will be called each time the fetch request to Kafka should be issued. Topic, partition and offset are self-explanatory.
// Returns slice of Messages and an error if a fetch error occurred.
func (this *SiestaClient) Fetch(topic string, partition int32, offset int64) ([]*Message, error) {
	Tracef(this, "Fetching %s %d from %d", topic, partition, offset)
	response, err := this.connector.Fetch(topic, partition, offset)
	if err != nil {
		return nil, err
	}

	messages := make([]*Message, 0)

	timestamp := time.Now().UnixNano() / int64(time.Millisecond)
	collector := func(topic string, partition int32, offset int64, key []byte, value []byte) {
		decodedKey, err := this.config.KeyDecoder.Decode(key)
		if err != nil {
			//TODO: what if we fail to decode the key: fail-fast or fail-safe strategy?
			Error(this, err.Error())
		}
		decodedValue, err := this.config.ValueDecoder.Decode(value)
		if err != nil {
			//TODO: what if we fail to decode the value: fail-fast or fail-safe strategy?
			Error(this, err.Error())
		}

		if this.config.Debug {
			decodedKey = []int64{timestamp}
		}

		messages = append(messages, &Message{
			Key:          key,
			Value:        value,
			DecodedKey:   decodedKey,
			DecodedValue: decodedValue,
			Topic:        topic,
			Partition:    partition,
			Offset:       offset,
		})
	}

	return messages, response.CollectMessages(collector)
}

// Checks whether the given error indicates an OffsetOutOfRange error.
func (this *SiestaClient) IsOffsetOutOfRange(err error) bool {
	return err == siesta.ErrOffsetOutOfRange
}

// This will be called to handle OffsetOutOfRange error. OffsetTime will be either "smallest" or "largest".
func (this *SiestaClient) GetAvailableOffset(topic string, partition int32, offsetTime string) (int64, error) {
	time := siesta.LatestTime
	if offsetTime == "smallest" {
		time = siesta.EarliestTime
	}
	return this.connector.GetAvailableOffset(topic, partition, time)
}

// Gets the offset for a given group, topic and partition.
// May return an error if fails to retrieve the offset.
func (this *SiestaClient) GetOffset(group string, topic string, partition int32) (int64, error) {
	return this.connector.GetOffset(group, topic, partition)
}

// Commits the given offset for a given group, topic and partition.
// May return an error if fails to commit the offset.
func (this *SiestaClient) CommitOffset(group string, topic string, partition int32, offset int64) error {
	return this.connector.CommitOffset(group, topic, partition, offset)
}

// Gracefully shuts down this client.
func (this *SiestaClient) Close() {
	<-this.connector.Close()
}

// BootstrapBrokers queries the ConsumerCoordinator for all known brokers in the cluster to be used later as a bootstrap list for the LowLevelClient.
func BootstrapBrokers(coordinator ConsumerCoordinator) ([]string, error) {
	bootstrapBrokers := make([]string, 0)
	brokers, err := coordinator.GetAllBrokers()
	if err != nil {
		return nil, err
	}
	for _, broker := range brokers {
		bootstrapBrokers = append(bootstrapBrokers, fmt.Sprintf("%s:%d", broker.Host, broker.Port))
	}

	return bootstrapBrokers, nil
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/ecopony/pggameday"
	"github.com/ecopony/gamedayapi"
	"fmt"
	"os"
	"sync"
	"strconv"
)

var commands = []string{
	"import-pitches-for-team-and-year",
	"import-pitches-for-year",
}

func main() {
	args := os.Args[1:]
	command := args[0]
	if !isCommandValid(command) {
		fmt.Println(fmt.Sprintf("%s is not a valid command. Valid commands:", command))
		printValidCommands()
		os.Exit(1)
	}

	if command == "import-pitches-for-team-and-year" {
		validateArgLength(args, 2)
		teamCode := args[1]

		yearArgs := args[2:]
		var years []int
		for i := 0; i < len(yearArgs); i++ {
			year, err := strconv.Atoi(yearArgs[i])
			if err != nil {
				fmt.Println("Year is not valid")
			}
			years = append(years, year)
		}

		pggameday.ImportPitchesForTeamAndYears(teamCode, years)
	} else {	// import-pitches-for-year
		validateArgLength(args, 1)
		yearArg := args[1]
		year, err := strconv.Atoi(yearArg)
		if err != nil {
			fmt.Println("Year is not valid")
		}
		years := []int{year}
		teams := gamedayapi.TeamsForYear(year)
		var wg sync.WaitGroup
		for _, team := range teams {
			wg.Add(1)
			go func(team string) {
				defer wg.Done()
				pggameday.ImportPitchesForTeamAndYears(team, years)
			}(team)
		}
		wg.Wait()
	}
}

func isCommandValid(command string) bool {
	for _, validCommand := range commands {
		if command == validCommand {
			return true
		}
	}
	return false
}

func printUsage() {
	fmt.Println("Usage: pgmlbgd <command> [<team code>] <years>")
}

func printValidCommands() {
	printUsage()
	fmt.Println("Valid commands:")
	for _, validCommand := range commands {
		fmt.Println(fmt.Sprintf("\t%s", validCommand))
	}
}

func validateArgLength(args []string, validLength int) {
	if len(args) <= validLength {
		printUsage()
		os.Exit(1)
	}
}

<commit_msg>Refactoring the command line bits to use branching using a map to simplify command validation and break out the command implementations into their own functions.<commit_after>package main

import (
	"github.com/ecopony/pggameday"
	"github.com/ecopony/gamedayapi"
	"fmt"
	"os"
	"sync"
	"strconv"
)

var commands = map[string]func([]string) {
	"import-pitches-for-year": importPitchesForYear,
	"import-pitches-for-team-and-year": importPitchesForTeamAndYear,
}

func main() {
	args := os.Args[1:]
	command := args[0]

	if function, ok := commands[command]; ok {
		function(args)
	} else {
		fmt.Println(fmt.Sprintf("%s is not a valid command. Valid commands:", command))
		printValidCommands()
		os.Exit(1)
	}
}

func printUsage() {
	fmt.Println("Usage: pgmlbgd <command> [<team code>] <years>")
}

func printValidCommands() {
	printUsage()
	fmt.Println("Valid commands:")
	for validCommand, _ := range commands {
		fmt.Println(fmt.Sprintf("\t%s", validCommand))
	}
}

func validateArgLength(args []string, validLength int) {
	if len(args) <= validLength {
		printUsage()
		os.Exit(1)
	}
}

func importPitchesForYear(args []string) {
	validateArgLength(args, 1)
	yearArg := args[1]
	year, err := strconv.Atoi(yearArg)
	if err != nil {
		fmt.Println("Year is not valid")
	}
	years := []int{year}
	teams := gamedayapi.TeamsForYear(year)
	var wg sync.WaitGroup
	for _, team := range teams {
		wg.Add(1)
		go func(team string) {
			defer wg.Done()
			pggameday.ImportPitchesForTeamAndYears(team, years)
		}(team)
	}
	wg.Wait()
}

func importPitchesForTeamAndYear(args []string) {
	validateArgLength(args, 2)
	teamCode := args[1]
	yearArgs := args[2:]
	var years []int
	for i := 0; i < len(yearArgs); i++ {
		year, err := strconv.Atoi(yearArgs[i])
		if err != nil {
			fmt.Println("Year is not valid")
		}
		years = append(years, year)
	}

	pggameday.ImportPitchesForTeamAndYears(teamCode, years)
}

<|endoftext|>


<commit_before>//
// DISCLAIMER
//
// Copyright 2017 ArangoDB GmbH, Cologne, Germany
//
// 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.
//
// Copyright holder is ArangoDB GmbH, Cologne, Germany
//
// Author Ewout Prangsma
//

package test

import (
	"context"
	"crypto/tls"
	httplib "net/http"
	"os"
	"strings"
	"sync"
	"testing"
	"time"

	driver "github.com/arangodb/go-driver"
	"github.com/arangodb/go-driver/http"
	"github.com/arangodb/go-driver/vst"
	"github.com/arangodb/go-driver/vst/protocol"
)

var (
	logEndpointsOnce sync.Once
)

// skipBelowVersion skips the test if the current server version is less than
// the given version.
func skipBelowVersion(c driver.Client, version driver.Version, t *testing.T) {
	x, err := c.Version(nil)
	if err != nil {
		t.Fatalf("Failed to get version info: %s", describe(err))
	}
	if x.Version.CompareTo(version) < 0 {
		t.Skipf("Skipping below version '%s', got version '%s'", version, x.Version)
	}
}

// getEndpointsFromEnv returns the endpoints specified in the TEST_ENDPOINTS
// environment variable.
func getEndpointsFromEnv(t testEnv) []string {
	eps := strings.Split(os.Getenv("TEST_ENDPOINTS"), ",")
	if len(eps) == 0 {
		t.Fatal("No endpoints found in environment variable TEST_ENDPOINTS")
	}
	return eps
}

// getContentTypeFromEnv returns the content-type specified in the TEST_CONTENT_TYPE
// environment variable (json|vpack).
func getContentTypeFromEnv(t testEnv) driver.ContentType {
	switch ct := os.Getenv("TEST_CONTENT_TYPE"); ct {
	case "vpack":
		return driver.ContentTypeVelocypack
	case "json", "":
		return driver.ContentTypeJSON
	default:
		t.Fatalf("Unknown content type '%s'", ct)
		return 0
	}
}

// createAuthenticationFromEnv initializes an authentication specified in the TEST_AUTHENTICATION
// environment variable.
func createAuthenticationFromEnv(t testEnv) driver.Authentication {
	authSpec := os.Getenv("TEST_AUTHENTICATION")
	if authSpec == "" {
		return nil
	}
	parts := strings.Split(authSpec, ":")
	switch parts[0] {
	case "basic":
		if len(parts) != 3 {
			t.Fatalf("Expected username & password for basic authentication")
		}
		return driver.BasicAuthentication(parts[1], parts[2])
	case "jwt":
		if len(parts) != 3 {
			t.Fatalf("Expected username & password for jwt authentication")
		}
		return driver.JWTAuthentication(parts[1], parts[2])
	default:
		t.Fatalf("Unknown authentication: '%s'", parts[0])
		return nil
	}
}

// createConnectionFromEnv initializes a Connection from information specified in environment variables.
func createConnectionFromEnv(t testEnv) driver.Connection {
	connSpec := os.Getenv("TEST_CONNECTION")
	connVer := os.Getenv("TEST_CVERSION")
	switch connSpec {
	case "vst":
		var version protocol.Version
		switch connVer {
		case "1.0", "":
			version = protocol.Version1_0
		case "1.1":
			version = protocol.Version1_1
		default:
			t.Fatalf("Unknown connection version '%s'", connVer)
		}
		config := vst.ConnectionConfig{
			Endpoints: getEndpointsFromEnv(t),
			TLSConfig: &tls.Config{InsecureSkipVerify: true},
			Transport: protocol.TransportConfig{
				Version: version,
			},
		}
		conn, err := vst.NewConnection(config)
		if err != nil {
			t.Fatalf("Failed to create new vst connection: %s", describe(err))
		}
		return conn

	case "http", "":
		config := http.ConnectionConfig{
			Endpoints:   getEndpointsFromEnv(t),
			TLSConfig:   &tls.Config{InsecureSkipVerify: true},
			ContentType: getContentTypeFromEnv(t),
		}
		conn, err := http.NewConnection(config)
		if err != nil {
			t.Fatalf("Failed to create new http connection: %s", describe(err))
		}
		return conn

	default:
		t.Fatalf("Unknown connection type: '%s'", connSpec)
		return nil
	}
}

// createClientFromEnv initializes a Client from information specified in environment variables.
func createClientFromEnv(t testEnv, waitUntilReady bool, connection ...*driver.Connection) driver.Client {
	conn := createConnectionFromEnv(t)
	if len(connection) == 1 {
		*connection[0] = conn
	}
	c, err := driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: createAuthenticationFromEnv(t),
	})
	if err != nil {
		t.Fatalf("Failed to create new client: %s", describe(err))
	}
	if waitUntilReady {
		timeout := 3 * time.Minute
		ctx, cancel := context.WithTimeout(context.Background(), timeout)
		defer cancel()
		if up := waitUntilServerAvailable(ctx, c, t); !up {
			t.Fatalf("Connection is not available in %s", timeout)
		}
		// Synchronize endpoints
		if err := c.SynchronizeEndpoints(context.Background()); err != nil {
			t.Errorf("Failed to synchronize endpoints: %s", describe(err))
		} else {
			logEndpointsOnce.Do(func() {
				t.Logf("Found endpoints: %v", conn.Endpoints())
			})
		}
	}
	return c
}

// waitUntilServerAvailable keeps waiting until the server/cluster that the client is addressing is available.
func waitUntilServerAvailable(ctx context.Context, c driver.Client, t testEnv) bool {
	instanceUp := make(chan bool)
	go func() {
		for {
			verCtx, cancel := context.WithTimeout(ctx, time.Second*5)
			if _, err := c.Version(verCtx); err == nil {
				//t.Logf("Found version %s", v.Version)
				cancel()
				instanceUp <- true
				return
			} else {
				cancel()
				//t.Logf("Version failed: %s %#v", describe(err), err)
				time.Sleep(time.Second)
			}
		}
	}()
	select {
	case up := <-instanceUp:
		return up
	case <-ctx.Done():
		return false
	}
}

// TestCreateClientHttpConnection creates an HTTP connection to the environment specified
// endpoints and creates a client for that.
func TestCreateClientHttpConnection(t *testing.T) {
	conn, err := http.NewConnection(http.ConnectionConfig{
		Endpoints: getEndpointsFromEnv(t),
		TLSConfig: &tls.Config{InsecureSkipVerify: true},
	})
	if err != nil {
		t.Fatalf("Failed to create new http connection: %s", describe(err))
	}
	_, err = driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: createAuthenticationFromEnv(t),
	})
	if err != nil {
		t.Fatalf("Failed to create new client: %s", describe(err))
	}
}

// TestCreateClientHttpConnectionCustomTransport creates an HTTP connection to the environment specified
// endpoints with a custom HTTP roundtripper and creates a client for that.
func TestCreateClientHttpConnectionCustomTransport(t *testing.T) {
	conn, err := http.NewConnection(http.ConnectionConfig{
		Endpoints: getEndpointsFromEnv(t),
		Transport: &httplib.Transport{},
		TLSConfig: &tls.Config{InsecureSkipVerify: true},
	})
	if err != nil {
		t.Fatalf("Failed to create new http connection: %s", describe(err))
	}
	c, err := driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: createAuthenticationFromEnv(t),
	})
	if err != nil {
		t.Fatalf("Failed to create new client: %s", describe(err))
	}
	timeout := 3 * time.Minute
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	if up := waitUntilServerAvailable(ctx, c, t); !up {
		t.Fatalf("Connection is not available in %s", timeout)
	}
	if info, err := c.Version(driver.WithDetails(ctx)); err != nil {
		t.Errorf("Version failed: %s", describe(err))
	} else {
		t.Logf("Got server version %s", info)
	}
}

// TestResponseHeader checks the Response.Header function.
func TestResponseHeader(t *testing.T) {
	c := createClientFromEnv(t, true)
	ctx := context.Background()
	var resp driver.Response

	db := ensureDatabase(ctx, c, "_system", nil, t)
	col := ensureCollection(ctx, db, "response_header_test", nil, t)

	// `ETag` header must contain the `_id` of the new document
	// `ERev` header must contain the `_rev` of the new document
	doc := map[string]string{
		"Test":   "TestResponseHeader",
		"Intent": "Check Response.Header",
	}
	meta, err := col.CreateDocument(driver.WithResponse(ctx, &resp), doc)
	if err != nil {
		t.Fatalf("CreateDocument failed: %s", describe(err))
	}
	if x := resp.Header("ETag"); x != meta.Rev {
		t.Errorf("Unexpected result from Header('ETag'), got '%s', expected '%s'", x, meta.Rev)
	}
	if x := resp.Header("etag"); x != meta.Rev {
		t.Errorf("Unexpected result from Header('etag'), got '%s', expected '%s'", x, meta.Rev)
	}
	if x := resp.Header("ETAG"); x != meta.Rev {
		t.Errorf("Unexpected result from Header('ETAG'), got '%s', expected '%s'", x, meta.Rev)
	}
}
<commit_msg>Updated test<commit_after>//
// DISCLAIMER
//
// Copyright 2017 ArangoDB GmbH, Cologne, Germany
//
// 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.
//
// Copyright holder is ArangoDB GmbH, Cologne, Germany
//
// Author Ewout Prangsma
//

package test

import (
	"context"
	"crypto/tls"
	httplib "net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"

	driver "github.com/arangodb/go-driver"
	"github.com/arangodb/go-driver/http"
	"github.com/arangodb/go-driver/vst"
	"github.com/arangodb/go-driver/vst/protocol"
)

var (
	logEndpointsOnce sync.Once
)

// skipBelowVersion skips the test if the current server version is less than
// the given version.
func skipBelowVersion(c driver.Client, version driver.Version, t *testing.T) {
	x, err := c.Version(nil)
	if err != nil {
		t.Fatalf("Failed to get version info: %s", describe(err))
	}
	if x.Version.CompareTo(version) < 0 {
		t.Skipf("Skipping below version '%s', got version '%s'", version, x.Version)
	}
}

// getEndpointsFromEnv returns the endpoints specified in the TEST_ENDPOINTS
// environment variable.
func getEndpointsFromEnv(t testEnv) []string {
	eps := strings.Split(os.Getenv("TEST_ENDPOINTS"), ",")
	if len(eps) == 0 {
		t.Fatal("No endpoints found in environment variable TEST_ENDPOINTS")
	}
	return eps
}

// getContentTypeFromEnv returns the content-type specified in the TEST_CONTENT_TYPE
// environment variable (json|vpack).
func getContentTypeFromEnv(t testEnv) driver.ContentType {
	switch ct := os.Getenv("TEST_CONTENT_TYPE"); ct {
	case "vpack":
		return driver.ContentTypeVelocypack
	case "json", "":
		return driver.ContentTypeJSON
	default:
		t.Fatalf("Unknown content type '%s'", ct)
		return 0
	}
}

// createAuthenticationFromEnv initializes an authentication specified in the TEST_AUTHENTICATION
// environment variable.
func createAuthenticationFromEnv(t testEnv) driver.Authentication {
	authSpec := os.Getenv("TEST_AUTHENTICATION")
	if authSpec == "" {
		return nil
	}
	parts := strings.Split(authSpec, ":")
	switch parts[0] {
	case "basic":
		if len(parts) != 3 {
			t.Fatalf("Expected username & password for basic authentication")
		}
		return driver.BasicAuthentication(parts[1], parts[2])
	case "jwt":
		if len(parts) != 3 {
			t.Fatalf("Expected username & password for jwt authentication")
		}
		return driver.JWTAuthentication(parts[1], parts[2])
	default:
		t.Fatalf("Unknown authentication: '%s'", parts[0])
		return nil
	}
}

// createConnectionFromEnv initializes a Connection from information specified in environment variables.
func createConnectionFromEnv(t testEnv) driver.Connection {
	connSpec := os.Getenv("TEST_CONNECTION")
	connVer := os.Getenv("TEST_CVERSION")
	switch connSpec {
	case "vst":
		var version protocol.Version
		switch connVer {
		case "1.0", "":
			version = protocol.Version1_0
		case "1.1":
			version = protocol.Version1_1
		default:
			t.Fatalf("Unknown connection version '%s'", connVer)
		}
		config := vst.ConnectionConfig{
			Endpoints: getEndpointsFromEnv(t),
			TLSConfig: &tls.Config{InsecureSkipVerify: true},
			Transport: protocol.TransportConfig{
				Version: version,
			},
		}
		conn, err := vst.NewConnection(config)
		if err != nil {
			t.Fatalf("Failed to create new vst connection: %s", describe(err))
		}
		return conn

	case "http", "":
		config := http.ConnectionConfig{
			Endpoints:   getEndpointsFromEnv(t),
			TLSConfig:   &tls.Config{InsecureSkipVerify: true},
			ContentType: getContentTypeFromEnv(t),
		}
		conn, err := http.NewConnection(config)
		if err != nil {
			t.Fatalf("Failed to create new http connection: %s", describe(err))
		}
		return conn

	default:
		t.Fatalf("Unknown connection type: '%s'", connSpec)
		return nil
	}
}

// createClientFromEnv initializes a Client from information specified in environment variables.
func createClientFromEnv(t testEnv, waitUntilReady bool, connection ...*driver.Connection) driver.Client {
	conn := createConnectionFromEnv(t)
	if len(connection) == 1 {
		*connection[0] = conn
	}
	c, err := driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: createAuthenticationFromEnv(t),
	})
	if err != nil {
		t.Fatalf("Failed to create new client: %s", describe(err))
	}
	if waitUntilReady {
		timeout := 3 * time.Minute
		ctx, cancel := context.WithTimeout(context.Background(), timeout)
		defer cancel()
		if up := waitUntilServerAvailable(ctx, c, t); !up {
			t.Fatalf("Connection is not available in %s", timeout)
		}
		// Synchronize endpoints
		if err := c.SynchronizeEndpoints(context.Background()); err != nil {
			t.Errorf("Failed to synchronize endpoints: %s", describe(err))
		} else {
			logEndpointsOnce.Do(func() {
				t.Logf("Found endpoints: %v", conn.Endpoints())
			})
		}
	}
	return c
}

// waitUntilServerAvailable keeps waiting until the server/cluster that the client is addressing is available.
func waitUntilServerAvailable(ctx context.Context, c driver.Client, t testEnv) bool {
	instanceUp := make(chan bool)
	go func() {
		for {
			verCtx, cancel := context.WithTimeout(ctx, time.Second*5)
			if _, err := c.Version(verCtx); err == nil {
				//t.Logf("Found version %s", v.Version)
				cancel()
				instanceUp <- true
				return
			} else {
				cancel()
				//t.Logf("Version failed: %s %#v", describe(err), err)
				time.Sleep(time.Second)
			}
		}
	}()
	select {
	case up := <-instanceUp:
		return up
	case <-ctx.Done():
		return false
	}
}

// TestCreateClientHttpConnection creates an HTTP connection to the environment specified
// endpoints and creates a client for that.
func TestCreateClientHttpConnection(t *testing.T) {
	conn, err := http.NewConnection(http.ConnectionConfig{
		Endpoints: getEndpointsFromEnv(t),
		TLSConfig: &tls.Config{InsecureSkipVerify: true},
	})
	if err != nil {
		t.Fatalf("Failed to create new http connection: %s", describe(err))
	}
	_, err = driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: createAuthenticationFromEnv(t),
	})
	if err != nil {
		t.Fatalf("Failed to create new client: %s", describe(err))
	}
}

// TestCreateClientHttpConnectionCustomTransport creates an HTTP connection to the environment specified
// endpoints with a custom HTTP roundtripper and creates a client for that.
func TestCreateClientHttpConnectionCustomTransport(t *testing.T) {
	conn, err := http.NewConnection(http.ConnectionConfig{
		Endpoints: getEndpointsFromEnv(t),
		Transport: &httplib.Transport{},
		TLSConfig: &tls.Config{InsecureSkipVerify: true},
	})
	if err != nil {
		t.Fatalf("Failed to create new http connection: %s", describe(err))
	}
	c, err := driver.NewClient(driver.ClientConfig{
		Connection:     conn,
		Authentication: createAuthenticationFromEnv(t),
	})
	if err != nil {
		t.Fatalf("Failed to create new client: %s", describe(err))
	}
	timeout := 3 * time.Minute
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	if up := waitUntilServerAvailable(ctx, c, t); !up {
		t.Fatalf("Connection is not available in %s", timeout)
	}
	if info, err := c.Version(driver.WithDetails(ctx)); err != nil {
		t.Errorf("Version failed: %s", describe(err))
	} else {
		t.Logf("Got server version %s", info)
	}
}

// TestResponseHeader checks the Response.Header function.
func TestResponseHeader(t *testing.T) {
	c := createClientFromEnv(t, true)
	ctx := context.Background()
	var resp driver.Response

	db := ensureDatabase(ctx, c, "_system", nil, t)
	col := ensureCollection(ctx, db, "response_header_test", nil, t)

	// `ETag` header must contain the `_rev` of the new document in quotes.
	doc := map[string]string{
		"Test":   "TestResponseHeader",
		"Intent": "Check Response.Header",
	}
	meta, err := col.CreateDocument(driver.WithResponse(ctx, &resp), doc)
	if err != nil {
		t.Fatalf("CreateDocument failed: %s", describe(err))
	}
	expectedETag := strconv.Quote(meta.Rev)
	if x := resp.Header("ETag"); x != expectedETag {
		t.Errorf("Unexpected result from Header('ETag'), got '%s', expected '%s'", x, expectedETag)
	}
	if x := resp.Header("etag"); x != expectedETag {
		t.Errorf("Unexpected result from Header('etag'), got '%s', expected '%s'", x, expectedETag)
	}
	if x := resp.Header("ETAG"); x != expectedETag {
		t.Errorf("Unexpected result from Header('ETAG'), got '%s', expected '%s'", x, expectedETag)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2015 Google Inc. All rights reserved.

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 e2e

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util/wait"

	"github.com/golang/glog"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

// Delete a Replication Controller and all pods it spawned
func DeleteRC(c *client.Client, ns, name string) error {
	rc, err := c.ReplicationControllers(ns).Get(name)
	if err != nil {
		return fmt.Errorf("Failed to find replication controller %s in namespace %s: %v", name, ns, err)
	}

	rc.Spec.Replicas = 0

	if _, err := c.ReplicationControllers(ns).Update(rc); err != nil {
		return fmt.Errorf("Failed to resize replication controller %s to zero: %v", name, err)
	}

	if err := wait.Poll(time.Second, time.Minute*20, client.ControllerHasDesiredReplicas(c, rc)); err != nil {
		return fmt.Errorf("Error waiting for replication controller %s replicas to reach 0: %v", name, err)
	}

	// Delete the replication controller.
	if err := c.ReplicationControllers(ns).Delete(name); err != nil {
		return fmt.Errorf("Failed to delete replication controller %s: %v", name, err)
	}
	return nil
}

// Launch a Replication Controller and wait for all pods it spawns
// to become running
func RunRC(c *client.Client, name string, ns, image string, replicas int) {
	defer GinkgoRecover()

	var last int
	current := 0
	same := 0

	defer func() {
		By("Cleaning up the replication controller")
		err := DeleteRC(c, ns, name)
		Expect(err).NotTo(HaveOccurred())
	}()

	By(fmt.Sprintf("Creating replication controller %s", name))
	_, err := c.ReplicationControllers(ns).Create(&api.ReplicationController{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: replicas,
			Selector: map[string]string{
				"name": name,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  name,
							Image: image,
							Ports: []api.ContainerPort{{ContainerPort: 80}},
						},
					},
				},
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())

	By(fmt.Sprintf("Making sure all %d replicas exist", replicas))
	label := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := c.Pods(ns).List(label)
	Expect(err).NotTo(HaveOccurred())
	current = len(pods.Items)
	failCount := 5
	for same < failCount && current < replicas {
		glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Controller %s: Number of submitted pods dropped from %d to %d", last, current)
		}

		if same >= failCount {
			glog.Infof("No pods submitted for the last %d checks", failCount)
		}

		last = current
		time.Sleep(5 * time.Second)
		pods, err = c.Pods(ns).List(label)
		Expect(err).NotTo(HaveOccurred())
		current = len(pods.Items)
	}
	Expect(current).To(Equal(replicas))
	glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)

	By("Waiting for each pod to be running")
	same = 0
	last = 0
	failCount = 10
	current = 0
	for same < failCount && current < replicas {
		current = 0
		waiting := 0
		pending := 0
		unknown := 0
		time.Sleep(10 * time.Second)

		currentPods, listErr := c.Pods(ns).List(label)
		Expect(listErr).NotTo(HaveOccurred())
		if len(currentPods.Items) != len(pods.Items) {
			Failf("Number of reported pods changed: %d vs %d", len(currentPods.Items), len(pods.Items))
		}
		for _, p := range currentPods.Items {
			if p.Status.Phase == api.PodRunning {
				current++
			} else if p.Status.Phase == api.PodPending {
				if p.Spec.Host == "" {
					waiting++
				} else {
					pending++
				}
			} else if p.Status.Phase == api.PodUnknown {
				unknown++
			}
		}
		glog.Infof("Pod States: %d running, %d pending, %d waiting, %d unknown ", current, pending, waiting, unknown)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Number of running pods dropped from %d to %d", last, current)
		}
		if same >= failCount {
			glog.Infof("No pods started for the last %d checks", failCount)
		}
		last = current
	}
	Expect(current).To(Equal(replicas))
}

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = Describe("Density", func() {
	var c *client.Client
	var minionCount int
	var RCName string
	var ns string

	BeforeEach(func() {
		var err error
		c, err = loadClient()
		expectNoError(err)
		minions, err := c.Nodes().List()
		expectNoError(err)
		minionCount = len(minions.Items)
		Expect(minionCount).NotTo(BeZero())
		ns = api.NamespaceDefault
	})

	AfterEach(func() {
		// Remove any remaining pods from this test if the
		// replication controller still exists and the replica count
		// isn't 0.  This means the controller wasn't cleaned up
		// during the test so clean it up here
		rc, err := c.ReplicationControllers(ns).Get(RCName)
		if err == nil && rc.Spec.Replicas != 0 {
			DeleteRC(c, ns, RCName)
		}
	})

	// Tests with "Skipped" substring in their name will be skipped when running
	// e2e test suite without --ginkgo.focus & --ginkgo.skip flags.

	type Density struct {
		skip          bool
		totalPods     int
		podsPerMinion int
		rcsPerThread  int
	}

	//This test should always run, even if larger densities are skipped.
	d3 := Density{totalPods: 3, podsPerMinion: 0, rcsPerThread: 1, skip: false}

	//These tests are varied and customizable.
	//TODO (wojtek-t):don't skip d30 after #6059
	d30 := Density{totalPods: 30, podsPerMinion: 0, rcsPerThread: 1, skip: true}
	d50 := Density{totalPods: 50, podsPerMinion: 0, rcsPerThread: 1, skip: true}
	d100 := Density{totalPods: 100, podsPerMinion: 0, rcsPerThread: 1, skip: true}
	d500t5 := Density{totalPods: 500, podsPerMinion: 10, rcsPerThread: 5, skip: true}
	d500t25 := Density{totalPods: 500, podsPerMinion: 10, rcsPerThread: 25, skip: true}

	dtests := []Density{d3, d30, d50, d100, d500t5, d500t25}

	//Run each test in the array which isn't skipped.
	for i := range dtests {

		//cannot do a range iterator over structs.
		dtest := dtests[i]

		glog.Info("Density test parameters: %v", dtest)

		//if ppm==0, its a raw density test.
		//otherwise, we continue launching n nodes per pod in threads till we meet the totalPods #.
		if dtest.podsPerMinion == 0 {
			//basic density tests
			name := fmt.Sprintf("should allow starting %d pods per node", dtest.totalPods)

			if dtest.skip {
				name = "[Skipped] " + name
			}
			It(name, func() {
				RCName = "my-hostname-density" + strconv.Itoa(dtest.totalPods) + "-" + string(util.NewUUID())
				RunRC(c, RCName, ns, "gcr.io/google_containers/pause:go", dtest.totalPods)
			})
			glog.Info("moving on, test already finished....")
		} else {
			// # of threads calibrate to totalPods
			threads := (dtest.totalPods / (dtest.podsPerMinion * dtest.rcsPerThread))

			name := fmt.Sprintf(
				"[Skipped] should be able to launch %v pods, %v per minion, in %v rcs/thread.",
				dtest.totalPods, dtest.podsPerMinion, dtest.rcsPerThread)

			if dtest.skip {
				name = "[Skipped] " + name
			}

			podsLaunched := 0
			It(name, func() {

				var wg sync.WaitGroup

				//count down latch.., once all threads are launched, we wait for
				//it to decrement down to zero.
				wg.Add(threads)

				//create queue of pending requests on the api server.
				for i := 0; i < threads; i++ {
					go func() {
						// call to wg.Done will serve as a count down latch.
						defer wg.Done()
						for i := 0; i < dtest.rcsPerThread; i++ {
							name := "my-short-lived-pod" + string(util.NewUUID())
							n := dtest.podsPerMinion * minionCount
							RunRC(c, name, ns, "gcr.io/google_containers/pause:go", n)
							podsLaunched += n
							glog.Info("Launched %v pods so far...", podsLaunched)
						}
					}()
				}
				//Wait for all the pods from all the RC's to return.
				wg.Wait()
				glog.Info("%v pods out of %v launched", podsLaunched, dtest.totalPods)
			})
		}
	}
})
<commit_msg>E2E: parameter logs to inside of It declarations.<commit_after>/*
Copyright 2015 Google Inc. All rights reserved.

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 e2e

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util/wait"

	"github.com/golang/glog"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

// Delete a Replication Controller and all pods it spawned
func DeleteRC(c *client.Client, ns, name string) error {
	rc, err := c.ReplicationControllers(ns).Get(name)
	if err != nil {
		return fmt.Errorf("Failed to find replication controller %s in namespace %s: %v", name, ns, err)
	}

	rc.Spec.Replicas = 0

	if _, err := c.ReplicationControllers(ns).Update(rc); err != nil {
		return fmt.Errorf("Failed to resize replication controller %s to zero: %v", name, err)
	}

	if err := wait.Poll(time.Second, time.Minute*20, client.ControllerHasDesiredReplicas(c, rc)); err != nil {
		return fmt.Errorf("Error waiting for replication controller %s replicas to reach 0: %v", name, err)
	}

	// Delete the replication controller.
	if err := c.ReplicationControllers(ns).Delete(name); err != nil {
		return fmt.Errorf("Failed to delete replication controller %s: %v", name, err)
	}
	return nil
}

// Launch a Replication Controller and wait for all pods it spawns
// to become running
func RunRC(c *client.Client, name string, ns, image string, replicas int) {
	defer GinkgoRecover()

	var last int
	current := 0
	same := 0

	defer func() {
		By("Cleaning up the replication controller")
		err := DeleteRC(c, ns, name)
		Expect(err).NotTo(HaveOccurred())
	}()

	By(fmt.Sprintf("Creating replication controller %s", name))
	_, err := c.ReplicationControllers(ns).Create(&api.ReplicationController{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: replicas,
			Selector: map[string]string{
				"name": name,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  name,
							Image: image,
							Ports: []api.ContainerPort{{ContainerPort: 80}},
						},
					},
				},
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())

	By(fmt.Sprintf("Making sure all %d replicas exist", replicas))
	label := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := c.Pods(ns).List(label)
	Expect(err).NotTo(HaveOccurred())
	current = len(pods.Items)
	failCount := 5
	for same < failCount && current < replicas {
		glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Controller %s: Number of submitted pods dropped from %d to %d", last, current)
		}

		if same >= failCount {
			glog.Infof("No pods submitted for the last %d checks", failCount)
		}

		last = current
		time.Sleep(5 * time.Second)
		pods, err = c.Pods(ns).List(label)
		Expect(err).NotTo(HaveOccurred())
		current = len(pods.Items)
	}
	Expect(current).To(Equal(replicas))
	glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)

	By("Waiting for each pod to be running")
	same = 0
	last = 0
	failCount = 10
	current = 0
	for same < failCount && current < replicas {
		current = 0
		waiting := 0
		pending := 0
		unknown := 0
		time.Sleep(10 * time.Second)

		currentPods, listErr := c.Pods(ns).List(label)
		Expect(listErr).NotTo(HaveOccurred())
		if len(currentPods.Items) != len(pods.Items) {
			Failf("Number of reported pods changed: %d vs %d", len(currentPods.Items), len(pods.Items))
		}
		for _, p := range currentPods.Items {
			if p.Status.Phase == api.PodRunning {
				current++
			} else if p.Status.Phase == api.PodPending {
				if p.Spec.Host == "" {
					waiting++
				} else {
					pending++
				}
			} else if p.Status.Phase == api.PodUnknown {
				unknown++
			}
		}
		glog.Infof("Pod States: %d running, %d pending, %d waiting, %d unknown ", current, pending, waiting, unknown)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Number of running pods dropped from %d to %d", last, current)
		}
		if same >= failCount {
			glog.Infof("No pods started for the last %d checks", failCount)
		}
		last = current
	}
	Expect(current).To(Equal(replicas))
}

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = Describe("Density", func() {
	var c *client.Client
	var minionCount int
	var RCName string
	var ns string

	BeforeEach(func() {
		var err error
		c, err = loadClient()
		expectNoError(err)
		minions, err := c.Nodes().List()
		expectNoError(err)
		minionCount = len(minions.Items)
		Expect(minionCount).NotTo(BeZero())
		ns = api.NamespaceDefault
	})

	AfterEach(func() {
		// Remove any remaining pods from this test if the
		// replication controller still exists and the replica count
		// isn't 0.  This means the controller wasn't cleaned up
		// during the test so clean it up here
		rc, err := c.ReplicationControllers(ns).Get(RCName)
		if err == nil && rc.Spec.Replicas != 0 {
			DeleteRC(c, ns, RCName)
		}
	})

	// Tests with "Skipped" substring in their name will be skipped when running
	// e2e test suite without --ginkgo.focus & --ginkgo.skip flags.

	type Density struct {
		skip          bool
		totalPods     int
		podsPerMinion int
		rcsPerThread  int
	}

	//This test should always run, even if larger densities are skipped.
	d3 := Density{totalPods: 3, podsPerMinion: 0, rcsPerThread: 1, skip: false}

	//These tests are varied and customizable.
	//TODO (wojtek-t):don't skip d30 after #6059
	d30 := Density{totalPods: 30, podsPerMinion: 0, rcsPerThread: 1, skip: true}
	d50 := Density{totalPods: 50, podsPerMinion: 0, rcsPerThread: 1, skip: true}
	d100 := Density{totalPods: 100, podsPerMinion: 0, rcsPerThread: 1, skip: true}
	d500t5 := Density{totalPods: 500, podsPerMinion: 10, rcsPerThread: 5, skip: true}
	d500t25 := Density{totalPods: 500, podsPerMinion: 10, rcsPerThread: 25, skip: true}

	dtests := []Density{d3, d30, d50, d100, d500t5, d500t25}

	//Run each test in the array which isn't skipped.
	for i := range dtests {

		//cannot do a range iterator over structs.
		dtest := dtests[i]

		//if ppm==0, its a raw density test.
		//otherwise, we continue launching n nodes per pod in threads till we meet the totalPods #.
		if dtest.podsPerMinion == 0 {
			//basic density tests
			name := fmt.Sprintf("should allow starting %d pods per node", dtest.totalPods)

			if dtest.skip {
				name = "[Skipped] " + name
			}
			It(name, func() {
				glog.Info("Density test parameters: %v", dtest)
				RCName = "my-hostname-density" + strconv.Itoa(dtest.totalPods) + "-" + string(util.NewUUID())
				RunRC(c, RCName, ns, "gcr.io/google_containers/pause:go", dtest.totalPods)
			})
		} else {
			// # of threads calibrate to totalPods
			threads := (dtest.totalPods / (dtest.podsPerMinion * dtest.rcsPerThread))

			name := fmt.Sprintf(
				"[Skipped] should be able to launch %v pods, %v per minion, in %v rcs/thread.",
				dtest.totalPods, dtest.podsPerMinion, dtest.rcsPerThread)

			if dtest.skip {
				name = "[Skipped] " + name
			}

			podsLaunched := 0
			It(name, func() {

				var wg sync.WaitGroup

				//count down latch.., once all threads are launched, we wait for
				//it to decrement down to zero.
				wg.Add(threads)

				//create queue of pending requests on the api server.
				for i := 0; i < threads; i++ {
					go func() {
						// call to wg.Done will serve as a count down latch.
						defer wg.Done()
						for i := 0; i < dtest.rcsPerThread; i++ {
							name := "my-short-lived-pod" + string(util.NewUUID())
							n := dtest.podsPerMinion * minionCount
							RunRC(c, name, ns, "gcr.io/google_containers/pause:go", n)
							podsLaunched += n
							glog.Info("Launched %v pods so far...", podsLaunched)
						}
					}()
				}
				//Wait for all the pods from all the RC's to return.
				wg.Wait()
				glog.Info("%v pods out of %v launched", podsLaunched, dtest.totalPods)
			})
		}
	}
})
<|endoftext|>


<commit_before>/*
Copyright 2015 Google Inc. All rights reserved.

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 e2e

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util/wait"

	"github.com/golang/glog"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

// Convenient wrapper around listing pods supporting retries.
func listPods(c *client.Client, namespace string, label labels.Selector) (*api.PodList, error) {
	maxRetries := 4
	pods, err := c.Pods(namespace).List(label)
	for i := 0; i < maxRetries; i++ {
		if err == nil {
			return pods, nil
		}
		pods, err = c.Pods(namespace).List(label)
	}
	return pods, err
}

// Delete a Replication Controller and all pods it spawned
func DeleteRC(c *client.Client, ns, name string) error {
	rc, err := c.ReplicationControllers(ns).Get(name)
	if err != nil {
		return fmt.Errorf("Failed to find replication controller %s in namespace %s: %v", name, ns, err)
	}

	rc.Spec.Replicas = 0

	if _, err := c.ReplicationControllers(ns).Update(rc); err != nil {
		return fmt.Errorf("Failed to resize replication controller %s to zero: %v", name, err)
	}

	// Wait up to 20 minutes until all replicas are killed.
	endTime := time.Now().Add(time.Minute * 20)
	for {
		if time.Now().After(endTime) {
			return fmt.Errorf("Timeout while waiting for replication controller %s replicas to 0", name)
		}
		remainingTime := endTime.Sub(time.Now())
		err := wait.Poll(time.Second, remainingTime, client.ControllerHasDesiredReplicas(c, rc))
		if err != nil {
			glog.Errorf("Error while waiting for replication controller %s replicas to read 0: %v", name, err)
		} else {
			break
		}
	}

	// Delete the replication controller.
	if err := c.ReplicationControllers(ns).Delete(name); err != nil {
		return fmt.Errorf("Failed to delete replication controller %s: %v", name, err)
	}
	return nil
}

// Launch a Replication Controller and wait for all pods it spawns
// to become running
func RunRC(c *client.Client, name string, ns, image string, replicas int) {
	defer GinkgoRecover()

	var last int
	current := 0
	same := 0

	defer func() {
		By("Cleaning up the replication controller")
		err := DeleteRC(c, ns, name)
		Expect(err).NotTo(HaveOccurred())
	}()

	By(fmt.Sprintf("Creating replication controller %s", name))
	_, err := c.ReplicationControllers(ns).Create(&api.ReplicationController{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: replicas,
			Selector: map[string]string{
				"name": name,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  name,
							Image: image,
							Ports: []api.ContainerPort{{ContainerPort: 80}},
						},
					},
				},
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())

	By(fmt.Sprintf("Making sure all %d replicas exist", replicas))
	label := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := listPods(c, ns, label)
	Expect(err).NotTo(HaveOccurred())
	current = len(pods.Items)
	failCount := 5
	for same < failCount && current < replicas {
		glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Controller %s: Number of submitted pods dropped from %d to %d", last, current)
		}

		if same >= failCount {
			glog.Infof("No pods submitted for the last %d checks", failCount)
		}

		last = current
		time.Sleep(5 * time.Second)
		pods, err = listPods(c, ns, label)
		Expect(err).NotTo(HaveOccurred())
		current = len(pods.Items)
	}
	Expect(current).To(Equal(replicas))
	glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)

	By("Waiting for each pod to be running")
	same = 0
	last = 0
	failCount = 10
	current = 0
	for same < failCount && current < replicas {
		current = 0
		waiting := 0
		pending := 0
		unknown := 0
		time.Sleep(10 * time.Second)

		currentPods, listErr := listPods(c, ns, label)
		Expect(listErr).NotTo(HaveOccurred())
		if len(currentPods.Items) != len(pods.Items) {
			Failf("Number of reported pods changed: %d vs %d", len(currentPods.Items), len(pods.Items))
		}
		for _, p := range currentPods.Items {
			if p.Status.Phase == api.PodRunning {
				current++
			} else if p.Status.Phase == api.PodPending {
				if p.Spec.Host == "" {
					waiting++
				} else {
					pending++
				}
			} else if p.Status.Phase == api.PodUnknown {
				unknown++
			}
		}
		glog.Infof("Pod States: %d running, %d pending, %d waiting, %d unknown ", current, pending, waiting, unknown)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Number of running pods dropped from %d to %d", last, current)
		}
		if same >= failCount {
			glog.Infof("No pods started for the last %d checks", failCount)
		}
		last = current
	}
	Expect(current).To(Equal(replicas))
}

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = Describe("Density", func() {
	var c *client.Client
	var minionCount int
	var RCName string
	var ns string

	BeforeEach(func() {
		var err error
		c, err = loadClient()
		expectNoError(err)
		minions, err := c.Nodes().List(labels.Everything(), fields.Everything())
		expectNoError(err)
		minionCount = len(minions.Items)
		Expect(minionCount).NotTo(BeZero())
		ns = api.NamespaceDefault
	})

	AfterEach(func() {
		// Remove any remaining pods from this test if the
		// replication controller still exists and the replica count
		// isn't 0.  This means the controller wasn't cleaned up
		// during the test so clean it up here
		rc, err := c.ReplicationControllers(ns).Get(RCName)
		if err == nil && rc.Spec.Replicas != 0 {
			DeleteRC(c, ns, RCName)
		}
	})

	// Tests with "Skipped" substring in their name will be skipped when running
	// e2e test suite without --ginkgo.focus & --ginkgo.skip flags.

	type Density struct {
		skip          bool
		podsPerMinion int
	}

	densityTests := []Density{
		// This test should always run, even if larger densities are skipped.
		{podsPerMinion: 3, skip: false},
		{podsPerMinion: 30, skip: false},
		// More than 30 pods per node is outside our v1.0 goals.
		// We might want to enable those tests in the future.
		{podsPerMinion: 50, skip: true},
		{podsPerMinion: 100, skip: true},
	}

	for _, testArg := range densityTests {
		name := fmt.Sprintf("should allow starting %d pods per node", testArg.podsPerMinion)
		if testArg.podsPerMinion <= 30 {
			name = "[Performance suite] " + name
		}
		if testArg.skip {
			name = "[Skipped] " + name
		}
		itArg := testArg
		It(name, func() {
			totalPods := itArg.podsPerMinion * minionCount
			RCName = "my-hostname-density" + strconv.Itoa(totalPods) + "-" + string(util.NewUUID())
			RunRC(c, RCName, ns, "gcr.io/google_containers/pause:go", totalPods)
		})
	}

	type Scalability struct {
		skip          bool
		totalPods     int
		podsPerMinion int
		rcsPerThread  int
	}

	scalabilityTests := []Scalability{
		{totalPods: 500, podsPerMinion: 10, rcsPerThread: 5, skip: true},
		{totalPods: 500, podsPerMinion: 10, rcsPerThread: 25, skip: true},
	}

	for _, testArg := range scalabilityTests {
		// # of threads calibrate to totalPods
		threads := (testArg.totalPods / (testArg.podsPerMinion * testArg.rcsPerThread))

		name := fmt.Sprintf(
			"should be able to launch %v pods, %v per minion, in %v rcs/thread.",
			testArg.totalPods, testArg.podsPerMinion, testArg.rcsPerThread)
		if testArg.skip {
			name = "[Skipped] " + name
		}

		itArg := testArg
		It(name, func() {
			podsLaunched := 0
			var wg sync.WaitGroup
			wg.Add(threads)

			// Create queue of pending requests on the api server.
			for i := 0; i < threads; i++ {
				go func() {
					defer wg.Done()
					for i := 0; i < itArg.rcsPerThread; i++ {
						name := "my-short-lived-pod" + string(util.NewUUID())
						n := itArg.podsPerMinion * minionCount
						RunRC(c, name, ns, "gcr.io/google_containers/pause:go", n)
						podsLaunched += n
						glog.Info("Launched %v pods so far...", podsLaunched)
					}
				}()
			}
			// Wait for all the pods from all the RC's to return.
			wg.Wait()
			glog.Info("%v pods out of %v launched", podsLaunched, itArg.totalPods)
		})
	}
})
<commit_msg>Add checking of events after all pods started to verify no failures in density test #6637<commit_after>/*
Copyright 2015 Google Inc. All rights reserved.

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 e2e

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/labels"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util/wait"

	"github.com/golang/glog"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

// Convenient wrapper around listing pods supporting retries.
func listPods(c *client.Client, namespace string, label labels.Selector) (*api.PodList, error) {
	maxRetries := 4
	pods, err := c.Pods(namespace).List(label)
	for i := 0; i < maxRetries; i++ {
		if err == nil {
			return pods, nil
		}
		pods, err = c.Pods(namespace).List(label)
	}
	return pods, err
}

// Delete a Replication Controller and all pods it spawned
func DeleteRC(c *client.Client, ns, name string) error {
	rc, err := c.ReplicationControllers(ns).Get(name)
	if err != nil {
		return fmt.Errorf("Failed to find replication controller %s in namespace %s: %v", name, ns, err)
	}

	rc.Spec.Replicas = 0

	if _, err := c.ReplicationControllers(ns).Update(rc); err != nil {
		return fmt.Errorf("Failed to resize replication controller %s to zero: %v", name, err)
	}

	// Wait up to 20 minutes until all replicas are killed.
	endTime := time.Now().Add(time.Minute * 20)
	for {
		if time.Now().After(endTime) {
			return fmt.Errorf("Timeout while waiting for replication controller %s replicas to 0", name)
		}
		remainingTime := endTime.Sub(time.Now())
		err := wait.Poll(time.Second, remainingTime, client.ControllerHasDesiredReplicas(c, rc))
		if err != nil {
			glog.Errorf("Error while waiting for replication controller %s replicas to read 0: %v", name, err)
		} else {
			break
		}
	}

	// Delete the replication controller.
	if err := c.ReplicationControllers(ns).Delete(name); err != nil {
		return fmt.Errorf("Failed to delete replication controller %s: %v", name, err)
	}
	return nil
}

// Launch a Replication Controller and wait for all pods it spawns
// to become running.  The controller will need to be cleaned up external
// to this method
func RunRC(c *client.Client, name string, ns, image string, replicas int) {
	defer GinkgoRecover()

	var last int
	current := 0
	same := 0

	By(fmt.Sprintf("Creating replication controller %s", name))
	_, err := c.ReplicationControllers(ns).Create(&api.ReplicationController{
		ObjectMeta: api.ObjectMeta{
			Name: name,
		},
		Spec: api.ReplicationControllerSpec{
			Replicas: replicas,
			Selector: map[string]string{
				"name": name,
			},
			Template: &api.PodTemplateSpec{
				ObjectMeta: api.ObjectMeta{
					Labels: map[string]string{"name": name},
				},
				Spec: api.PodSpec{
					Containers: []api.Container{
						{
							Name:  name,
							Image: image,
							Ports: []api.ContainerPort{{ContainerPort: 80}},
						},
					},
				},
			},
		},
	})
	Expect(err).NotTo(HaveOccurred())

	By(fmt.Sprintf("Making sure all %d replicas exist", replicas))
	label := labels.SelectorFromSet(labels.Set(map[string]string{"name": name}))
	pods, err := listPods(c, ns, label)
	Expect(err).NotTo(HaveOccurred())
	current = len(pods.Items)
	failCount := 5
	for same < failCount && current < replicas {
		glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Controller %s: Number of submitted pods dropped from %d to %d", name, last, current)
		}

		if same >= failCount {
			glog.Infof("No pods submitted for the last %d checks", failCount)
		}

		last = current
		time.Sleep(5 * time.Second)
		pods, err = listPods(c, ns, label)
		Expect(err).NotTo(HaveOccurred())
		current = len(pods.Items)
	}
	Expect(current).To(Equal(replicas))
	glog.Infof("Controller %s: Found %d pods out of %d", name, current, replicas)

	By("Waiting for each pod to be running")
	same = 0
	last = 0
	failCount = 10
	current = 0
	for same < failCount && current < replicas {
		current = 0
		waiting := 0
		pending := 0
		unknown := 0
		time.Sleep(10 * time.Second)

		currentPods, listErr := listPods(c, ns, label)
		Expect(listErr).NotTo(HaveOccurred())
		if len(currentPods.Items) != len(pods.Items) {
			Failf("Number of reported pods changed: %d vs %d", len(currentPods.Items), len(pods.Items))
		}
		for _, p := range currentPods.Items {
			if p.Status.Phase == api.PodRunning {
				current++
			} else if p.Status.Phase == api.PodPending {
				if p.Spec.Host == "" {
					waiting++
				} else {
					pending++
				}
			} else if p.Status.Phase == api.PodUnknown {
				unknown++
			}
		}
		glog.Infof("Pod States: %d running, %d pending, %d waiting, %d unknown ", current, pending, waiting, unknown)
		if last < current {
			same = 0
		} else if last == current {
			same++
		} else if current < last {
			Failf("Number of running pods dropped from %d to %d", last, current)
		}
		if same >= failCount {
			glog.Infof("No pods started for the last %d checks", failCount)
		}
		last = current
	}
	Expect(current).To(Equal(replicas))
}

// This test suite can take a long time to run, so by default it is added to
// the ginkgo.skip list (see driver.go).
// To run this suite you must explicitly ask for it by setting the
// -t/--test flag or ginkgo.focus flag.
var _ = Describe("Density", func() {
	var c *client.Client
	var minionCount int
	var RCName string
	var ns string

	BeforeEach(func() {
		var err error
		c, err = loadClient()
		expectNoError(err)
		minions, err := c.Nodes().List(labels.Everything(), fields.Everything())
		expectNoError(err)
		minionCount = len(minions.Items)
		Expect(minionCount).NotTo(BeZero())
		ns = api.NamespaceDefault
	})

	AfterEach(func() {
		// Remove any remaining pods from this test if the
		// replication controller still exists and the replica count
		// isn't 0.  This means the controller wasn't cleaned up
		// during the test so clean it up here
		rc, err := c.ReplicationControllers(ns).Get(RCName)
		if err == nil && rc.Spec.Replicas != 0 {
			By("Cleaning up the replication controller")
			err := DeleteRC(c, ns, RCName)
			expectNoError(err)
		}

		// Clean up the namespace if a non-default one was used
		if ns != api.NamespaceDefault {
			By("Cleaning up the namespace")
			err := c.Namespaces().Delete(ns)
			expectNoError(err)
		}
	})

	// Tests with "Skipped" substring in their name will be skipped when running
	// e2e test suite without --ginkgo.focus & --ginkgo.skip flags.
	type Density struct {
		skip          bool
		podsPerMinion int
	}

	densityTests := []Density{
		// This test should always run, even if larger densities are skipped.
		{podsPerMinion: 3, skip: false},
		{podsPerMinion: 30, skip: false},
		// More than 30 pods per node is outside our v1.0 goals.
		// We might want to enable those tests in the future.
		{podsPerMinion: 50, skip: true},
		{podsPerMinion: 100, skip: true},
	}

	for _, testArg := range densityTests {
		name := fmt.Sprintf("should allow starting %d pods per node", testArg.podsPerMinion)
		if testArg.podsPerMinion <= 30 {
			name = "[Performance suite] " + name
		}
		if testArg.skip {
			name = "[Skipped] " + name
		}
		itArg := testArg
		It(name, func() {
			uuid := string(util.NewUUID())
			totalPods := itArg.podsPerMinion * minionCount
			nameStr := strconv.Itoa(totalPods) + "-" + uuid
			ns = "e2e-density" + nameStr
			RCName = "my-hostname-density" + nameStr
			RunRC(c, RCName, ns, "gcr.io/google_containers/pause:go", totalPods)
			By("waiting for all events to be recorded")
			last := -1
			current := 0
			var events *api.EventList
			for last < current {
				e, err := c.Events(ns).List(
					labels.Everything(),
					fields.Set{
						"involvedObject.namespace": ns,
					}.AsSelector(),
				)
				expectNoError(err)
				last = current
				current = len(e.Items)
				events = e
				time.Sleep(10 * time.Second)
			}
			Logf("Found %d events", current)

			// Verify no pods were killed or failed to start
			By("verifying no pods were killed or failed to start")
			for _, e := range events.Items {
				for _, s := range []string{"kill", "fail"} {
					Expect(e.Reason).NotTo(ContainSubstring(s), "event:' %s', reason: '%s', message: '%s', field path: '%s'", e, e.ObjectMeta.Name, e.Message, e.InvolvedObject.FieldPath)
				}
			}
		})
	}

	type Scalability struct {
		skip          bool
		totalPods     int
		podsPerMinion int
		rcsPerThread  int
	}

	scalabilityTests := []Scalability{
		{totalPods: 500, podsPerMinion: 10, rcsPerThread: 5, skip: true},
		{totalPods: 500, podsPerMinion: 10, rcsPerThread: 25, skip: true},
	}

	for _, testArg := range scalabilityTests {
		// # of threads calibrate to totalPods
		threads := (testArg.totalPods / (testArg.podsPerMinion * testArg.rcsPerThread))

		name := fmt.Sprintf(
			"should be able to launch %v pods, %v per minion, in %v rcs/thread.",
			testArg.totalPods, testArg.podsPerMinion, testArg.rcsPerThread)
		if testArg.skip {
			name = "[Skipped] " + name
		}

		itArg := testArg
		It(name, func() {
			podsLaunched := 0
			var wg sync.WaitGroup
			wg.Add(threads)

			// Create queue of pending requests on the api server.
			for i := 0; i < threads; i++ {
				go func() {
					defer wg.Done()
					for i := 0; i < itArg.rcsPerThread; i++ {
						name := "my-short-lived-pod" + string(util.NewUUID())
						n := itArg.podsPerMinion * minionCount
						RunRC(c, name, ns, "gcr.io/google_containers/pause:go", n)
						podsLaunched += n
						glog.Info("Launched %v pods so far...", podsLaunched)
						err := DeleteRC(c, ns, name)
						expectNoError(err)
					}
				}()
			}
			// Wait for all the pods from all the RC's to return.
			wg.Wait()
			glog.Info("%v pods out of %v launched", podsLaunched, itArg.totalPods)
		})
	}
})
<|endoftext|>