<commit_before>package groupcsm

import (
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/mailgun/kafka-pixy/actor"
	"github.com/mailgun/kafka-pixy/config"
	"github.com/mailgun/kafka-pixy/consumer"
	"github.com/mailgun/kafka-pixy/consumer/dispatcher"
	"github.com/mailgun/kafka-pixy/consumer/msgfetcher"
	"github.com/mailgun/kafka-pixy/consumer/multiplexer"
	"github.com/mailgun/kafka-pixy/consumer/partitioncsm"
	"github.com/mailgun/kafka-pixy/consumer/subscriber"
	"github.com/mailgun/kafka-pixy/consumer/topiccsm"
	"github.com/mailgun/kafka-pixy/offsetmgr"
	"github.com/mailgun/kazoo-go"
	"github.com/pkg/errors"
)

// groupConsumer manages a fleet of topic consumers and disposes of those that
// have been inactive for the `Config.Consumer.DisposeAfter` period of time.
//
// implements `dispatcher.Factory`.
// implements `dispatcher.Tier`.
type T struct {
	actDesc     *actor.Descriptor
	cfg         *config.Proxy
	group       string
	dispatcher  *dispatcher.T
	kafkaClt    sarama.Client
	kazooClt    *kazoo.Kazoo
	msgFetcherF msgfetcher.Factory
	offsetMgrF  offsetmgr.Factory
	subscriber  *subscriber.T
	topicCsmCh  chan *topiccsm.T
	wg          sync.WaitGroup

	multiplexersMu sync.Mutex
	multiplexers   map[string]*multiplexer.T
}

func Spawn(parentActDesc *actor.Descriptor, childSpec dispatcher.ChildSpec,
	cfg *config.Proxy, kafkaClt sarama.Client, kazooClt *kazoo.Kazoo,
	offsetMgrF offsetmgr.Factory,
) *T {
	group := string(childSpec.Key())
	actDesc := parentActDesc.NewChild(fmt.Sprintf("%s", group))
	actDesc.AddLogField("kafka.group", group)
	gc := &T{
		actDesc:      actDesc,
		cfg:          cfg,
		group:        group,
		kafkaClt:     kafkaClt,
		kazooClt:     kazooClt,
		offsetMgrF:   offsetMgrF,
		multiplexers: make(map[string]*multiplexer.T),
		topicCsmCh:   make(chan *topiccsm.T),
	}

	gc.subscriber = subscriber.Spawn(gc.actDesc, gc.group, gc.cfg, gc.kazooClt)
	gc.msgFetcherF = msgfetcher.SpawnFactory(gc.actDesc, gc.cfg, gc.kafkaClt)
	actor.Spawn(gc.actDesc, &gc.wg, gc.run)

	// Finalizer is called when all downstream topic consumers expire or if
	// the dispatcher is explicitly told to stop by the upstream dispatcher.
	finalizer := func() {
		gc.subscriber.Stop()
		// The run goroutine stops when the subscriber's channel is closed.
		gc.wg.Wait()
		// Only after run is stopped it is safe to shutdown the fetcher factory.
		gc.msgFetcherF.Stop()
	}
	gc.dispatcher = dispatcher.Spawn(gc.actDesc, gc, cfg,
		dispatcher.WithChildSpec(childSpec), dispatcher.WithFinalizer(finalizer))
	return gc
}

// implements `dispatcher.Factory`.
func (gc *T) KeyOf(req consumer.Request) dispatcher.Key {
	return dispatcher.Key(req.Topic)
}

// implements `dispatcher.Factory`.
func (gc *T) SpawnChild(childSpec dispatcher.ChildSpec) {
	topic := string(childSpec.Key())
	topiccsm.Spawn(gc.actDesc, gc.group, childSpec, gc.cfg, gc.topicCsmCh,
		func() bool { return gc.isSafe2Stop(topic) })
}

// String return string ID of this group consumer to be posted in logs.
func (gc *T) String() string {
	return gc.actDesc.String()
}

func (gc *T) isSafe2Stop(topic string) bool {
	gc.multiplexersMu.Lock()
	mux := gc.multiplexers[topic]
	gc.multiplexersMu.Unlock()
	if mux == nil {
		return true
	}
	return mux.IsSafe2Stop()
}

func (gc *T) run() {
	var (
		topicConsumers          = make(map[string]*topiccsm.T)
		topics                  []string
		subscriptions           map[string][]string
		ok                      = true
		nilOrRetryCh            <-chan time.Time
		nilOrSubscriberTopicsCh chan<- []string
		rebalancingRequired     = false
		rebalancingInProgress   = false
		retryScheduled          = false
		stopped                 = false
		rebalanceResultCh       = make(chan error, 1)
	)
	for {
		select {
		case tc := <-gc.topicCsmCh:
			// It is assumed that only one topicConsumer can exist for a
			// particular topic at a time.
			if topicConsumers[tc.Topic()] == tc {
				delete(topicConsumers, tc.Topic())
			} else {
				topicConsumers[tc.Topic()] = tc
			}
			topics = listTopics(topicConsumers)
			nilOrSubscriberTopicsCh = gc.subscriber.Topics()
			continue
		case nilOrSubscriberTopicsCh <- topics:
			nilOrSubscriberTopicsCh = nil
			continue
		case subscriptions, ok = <-gc.subscriber.Subscriptions():
			nilOrRetryCh = nil
			if !ok {
				if !rebalancingInProgress {
					goto done
				}
				stopped = true
				continue
			}
			rebalancingRequired = true
		case err := <-rebalanceResultCh:
			rebalancingInProgress = false
			if err != nil {
				gc.actDesc.Log().WithError(err).Error("rebalancing failed")
				if stopped {
					goto done
				}
				nilOrRetryCh = time.After(gc.cfg.Consumer.RetryBackoff)
				retryScheduled = true
			}
			if stopped {
				goto done
			}

		case <-nilOrRetryCh:
			retryScheduled = false
		}

		if rebalancingRequired && !rebalancingInProgress && !retryScheduled {
			rebalanceActDesc := gc.actDesc.NewChild("rebalance")
			// Copy topicConsumers to make sure `rebalance` doesn't see any
			// changes we make while it is running.
			topicConsumersCopy := make(map[string]*topiccsm.T, len(topicConsumers))
			for topic, tc := range topicConsumers {
				topicConsumersCopy[topic] = tc
			}
			subscriptions := subscriptions
			actor.Spawn(rebalanceActDesc, nil, func() {
				gc.rebalance(rebalanceActDesc, topicConsumersCopy, subscriptions, rebalanceResultCh)
			})
			rebalancingInProgress = true
			rebalancingRequired = false
		}
	}
done:
	var wg sync.WaitGroup
	gc.multiplexersMu.Lock()
	for _, mux := range gc.multiplexers {
		wg.Add(1)
		go func(mux *multiplexer.T) {
			defer wg.Done()
			mux.Stop()
		}(mux)
	}
	gc.multiplexersMu.Unlock()
	wg.Wait()
}

func (gc *T) rebalance(actDesc *actor.Descriptor, topicConsumers map[string]*topiccsm.T,
	subscriptions map[string][]string, rebalanceResultCh chan<- error,
) {
	assignedPartitions, err := gc.resolvePartitions(subscriptions, gc.kafkaClt.Partitions)
	if err != nil {
		rebalanceResultCh <- err
		return
	}
	actDesc.Log().Infof("assigned partitions: %v", assignedPartitions)
	var wg sync.WaitGroup
	// Stop consuming partitions that are no longer assigned to this group
	// and start consuming newly assigned partitions for topics that has been
	// consumed already.
	gc.multiplexersMu.Lock()
	defer gc.multiplexersMu.Unlock()
	for topic, mux := range gc.multiplexers {
		gc.rewireMuxAsync(topic, &wg, mux, topicConsumers[topic], assignedPartitions[topic])
	}
	// Start consuming partitions for topics that has not been consumed before.
	for topic, assignedTopicPartitions := range assignedPartitions {
		tc := topicConsumers[topic]
		mux := gc.multiplexers[topic]
		if tc == nil || mux != nil {
			continue
		}
		topic := topic
		spawnInFn := func(partition int32) multiplexer.In {
			return partitioncsm.Spawn(gc.actDesc, gc.group, topic, partition,
				gc.cfg, gc.subscriber, gc.msgFetcherF, gc.offsetMgrF)
		}
		mux = multiplexer.New(gc.actDesc, spawnInFn)
		gc.rewireMuxAsync(topic, &wg, mux, tc, assignedTopicPartitions)
		gc.multiplexers[topic] = mux
	}
	wg.Wait()
	// Clean up gears for topics that do not have assigned partitions anymore.
	for topic, mux := range gc.multiplexers {
		if !mux.IsRunning() {
			delete(gc.multiplexers, topic)
		}
	}
	// Notify the caller that rebalancing has completed successfully.
	rebalanceResultCh <- nil
	return
}

// rewireMuxAsync calls muxInputs in another goroutine.
func (gc *T) rewireMuxAsync(topic string, wg *sync.WaitGroup, mux *multiplexer.T, tc *topiccsm.T, assigned []int32) {
	actor.Spawn(gc.actDesc.NewChild("rewire", topic), wg, func() {
		if tc == nil {
			// Parameter output is an interface type, therefore nil should be
			// passed explicitly.
			mux.WireUp(nil, nil)
			return
		}
		mux.WireUp(tc, assigned)
	})
}

// resolvePartitions given topic subscriptions of all consumer group members,
// resolves what topic partitions are assigned to the specified group member.
func (gc *T) resolvePartitions(subscriptions map[string][]string,
	topicPartitionsFn func(string) ([]int32, error)) (map[string][]int32, error,
) {
	// Convert members->topics to topic->members map.
	topicsToMembers := make(map[string][]string)
	for groupMemberID, topics := range subscriptions {
		for _, topic := range topics {
			topicsToMembers[topic] = append(topicsToMembers[topic], groupMemberID)
		}
	}
	// Create a set of topics this consumer group member subscribed to.
	subscribedTopics := make(map[string]bool)
	for _, topic := range subscriptions[gc.cfg.ClientID] {
		subscribedTopics[topic] = true
	}
	// Resolve new partition assignments for all subscribed topics.
	assignedPartitions := make(map[string][]int32)
	for topic := range subscribedTopics {
		topicPartitions, err := topicPartitionsFn(topic)
		if err != nil {
			return nil, errors.Wrapf(err, "failed to get partition list, topic=%s", topic)
		}
		subscribersToPartitions := assignTopicPartitions(topicPartitions, topicsToMembers[topic])
		assignedTopicPartitions := subscribersToPartitions[gc.cfg.ClientID]
		if len(assignedTopicPartitions) > 0 {
			assignedPartitions[topic] = assignedTopicPartitions
		}
	}
	return assignedPartitions, nil
}

// assignTopicPartitions divides topic partitions among all consumer group
// members subscribed to the topic. The algorithm used closely resembles the
// one implemented by the standard Java High-Level consumer
// (see http://kafka.apache.org/documentation.html#distributionimpl and scroll
// down to *Consumer registration algorithm*) except it does not take in account
// how partitions are distributed among brokers.
func assignTopicPartitions(partitions []int32, subscribers []string) map[string][]int32 {
	partitionCount := len(partitions)
	subscriberCount := len(subscribers)
	if partitionCount == 0 || subscriberCount == 0 {
		return nil
	}
	sort.Slice(partitions, func(i, j int) bool { return partitions[i] < partitions[j] })
	sort.Strings(subscribers)

	subscribersToPartitions := make(map[string][]int32, subscriberCount)
	partitionsPerSubscriber := partitionCount / subscriberCount
	extra := partitionCount - subscriberCount*partitionsPerSubscriber

	begin := 0
	for _, groupMemberID := range subscribers {
		end := begin + partitionsPerSubscriber
		if extra != 0 {
			end++
			extra--
		}
		assigned := partitions[begin:end]
		if len(assigned) > 0 {
			subscribersToPartitions[groupMemberID] = partitions[begin:end]
		}
		begin = end
	}
	return subscribersToPartitions
}

func listTopics(topicConsumers map[string]*topiccsm.T) []string {
	topics := make([]string, 0, len(topicConsumers))
	for topic := range topicConsumers {
		topics = append(topics, topic)
	}
	return topics
}
<commit_msg>Make group consumer topic channel buffered<commit_after>package groupcsm

import (
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/mailgun/kafka-pixy/actor"
	"github.com/mailgun/kafka-pixy/config"
	"github.com/mailgun/kafka-pixy/consumer"
	"github.com/mailgun/kafka-pixy/consumer/dispatcher"
	"github.com/mailgun/kafka-pixy/consumer/msgfetcher"
	"github.com/mailgun/kafka-pixy/consumer/multiplexer"
	"github.com/mailgun/kafka-pixy/consumer/partitioncsm"
	"github.com/mailgun/kafka-pixy/consumer/subscriber"
	"github.com/mailgun/kafka-pixy/consumer/topiccsm"
	"github.com/mailgun/kafka-pixy/offsetmgr"
	"github.com/mailgun/kazoo-go"
	"github.com/pkg/errors"
)

// groupConsumer manages a fleet of topic consumers and disposes of those that
// have been inactive for the `Config.Consumer.DisposeAfter` period of time.
//
// implements `dispatcher.Factory`.
// implements `dispatcher.Tier`.
type T struct {
	actDesc     *actor.Descriptor
	cfg         *config.Proxy
	group       string
	dispatcher  *dispatcher.T
	kafkaClt    sarama.Client
	kazooClt    *kazoo.Kazoo
	msgFetcherF msgfetcher.Factory
	offsetMgrF  offsetmgr.Factory
	subscriber  *subscriber.T
	topicCsmCh  chan *topiccsm.T
	wg          sync.WaitGroup

	multiplexersMu sync.Mutex
	multiplexers   map[string]*multiplexer.T
}

func Spawn(parentActDesc *actor.Descriptor, childSpec dispatcher.ChildSpec,
	cfg *config.Proxy, kafkaClt sarama.Client, kazooClt *kazoo.Kazoo,
	offsetMgrF offsetmgr.Factory,
) *T {
	group := string(childSpec.Key())
	actDesc := parentActDesc.NewChild(fmt.Sprintf("%s", group))
	actDesc.AddLogField("kafka.group", group)
	gc := &T{
		actDesc:      actDesc,
		cfg:          cfg,
		group:        group,
		kafkaClt:     kafkaClt,
		kazooClt:     kazooClt,
		offsetMgrF:   offsetMgrF,
		multiplexers: make(map[string]*multiplexer.T),
		topicCsmCh:   make(chan *topiccsm.T, cfg.Consumer.ChannelBufferSize),
	}

	gc.subscriber = subscriber.Spawn(gc.actDesc, gc.group, gc.cfg, gc.kazooClt)
	gc.msgFetcherF = msgfetcher.SpawnFactory(gc.actDesc, gc.cfg, gc.kafkaClt)
	actor.Spawn(gc.actDesc, &gc.wg, gc.run)

	// Finalizer is called when all downstream topic consumers expire or if
	// the dispatcher is explicitly told to stop by the upstream dispatcher.
	finalizer := func() {
		gc.subscriber.Stop()
		// The run goroutine stops when the subscriber's channel is closed.
		gc.wg.Wait()
		// Only after run is stopped it is safe to shutdown the fetcher factory.
		gc.msgFetcherF.Stop()
	}
	gc.dispatcher = dispatcher.Spawn(gc.actDesc, gc, cfg,
		dispatcher.WithChildSpec(childSpec), dispatcher.WithFinalizer(finalizer))
	return gc
}

// implements `dispatcher.Factory`.
func (gc *T) KeyOf(req consumer.Request) dispatcher.Key {
	return dispatcher.Key(req.Topic)
}

// implements `dispatcher.Factory`.
func (gc *T) SpawnChild(childSpec dispatcher.ChildSpec) {
	topic := string(childSpec.Key())
	topiccsm.Spawn(gc.actDesc, gc.group, childSpec, gc.cfg, gc.topicCsmCh,
		func() bool { return gc.isSafe2Stop(topic) })
}

// String return string ID of this group consumer to be posted in logs.
func (gc *T) String() string {
	return gc.actDesc.String()
}

func (gc *T) isSafe2Stop(topic string) bool {
	gc.multiplexersMu.Lock()
	mux := gc.multiplexers[topic]
	gc.multiplexersMu.Unlock()
	if mux == nil {
		return true
	}
	return mux.IsSafe2Stop()
}

func (gc *T) run() {
	var (
		topicConsumers          = make(map[string]*topiccsm.T)
		topics                  []string
		subscriptions           map[string][]string
		ok                      = true
		nilOrRetryCh            <-chan time.Time
		nilOrSubscriberTopicsCh chan<- []string
		rebalanceRequired       = false
		rebalancePending        = false
		rebalanceScheduled      = false
		stopped                 = false
		rebalanceResultCh       = make(chan error, 1)
	)
	for {
		select {
		case tc := <-gc.topicCsmCh:
			// It is assumed that only one topicConsumer can exist for a
			// particular topic at a time.
			if topicConsumers[tc.Topic()] == tc {
				delete(topicConsumers, tc.Topic())
			} else {
				topicConsumers[tc.Topic()] = tc
			}
			topics = listTopics(topicConsumers)
			gc.actDesc.Log().Infof("Topics updated: %s", topics)
			nilOrSubscriberTopicsCh = gc.subscriber.Topics()
			continue

		case nilOrSubscriberTopicsCh <- topics:
			nilOrSubscriberTopicsCh = nil
			continue

		case subscriptions, ok = <-gc.subscriber.Subscriptions():
			nilOrRetryCh = nil
			if !ok {
				if !rebalancePending {
					goto done
				}
				stopped = true
				continue
			}
			rebalanceRequired = true

		case err := <-rebalanceResultCh:
			rebalancePending = false
			if err != nil {
				gc.actDesc.Log().WithError(err).Error("rebalancing failed")
				if stopped {
					goto done
				}
				nilOrRetryCh = time.After(gc.cfg.Consumer.RetryBackoff)
				rebalanceScheduled = true
			}
			if stopped {
				goto done
			}
		case <-nilOrRetryCh:
			rebalanceScheduled = false
		}

		if rebalanceRequired && !rebalancePending && !rebalanceScheduled {
			rebalanceActDesc := gc.actDesc.NewChild("rebalance")
			// Copy topicConsumers to make sure `rebalance` doesn't see any
			// changes we make while it is running.
			topicConsumersCopy := make(map[string]*topiccsm.T, len(topicConsumers))
			for topic, tc := range topicConsumers {
				topicConsumersCopy[topic] = tc
			}
			subscriptions := subscriptions
			actor.Spawn(rebalanceActDesc, nil, func() {
				gc.rebalance(rebalanceActDesc, topicConsumersCopy, subscriptions, rebalanceResultCh)
			})
			rebalancePending = true
			rebalanceRequired = false
		}
	}
done:
	var wg sync.WaitGroup
	gc.multiplexersMu.Lock()
	for _, mux := range gc.multiplexers {
		wg.Add(1)
		go func(mux *multiplexer.T) {
			defer wg.Done()
			mux.Stop()
		}(mux)
	}
	gc.multiplexersMu.Unlock()
	wg.Wait()
}

func (gc *T) rebalance(actDesc *actor.Descriptor, topicConsumers map[string]*topiccsm.T,
	subscriptions map[string][]string, rebalanceResultCh chan<- error,
) {
	assignedPartitions, err := gc.resolvePartitions(subscriptions, gc.kafkaClt.Partitions)
	if err != nil {
		rebalanceResultCh <- err
		return
	}
	actDesc.Log().Infof("assigned partitions: %v", assignedPartitions)
	var wg sync.WaitGroup
	// Stop consuming partitions that are no longer assigned to this group
	// and start consuming newly assigned partitions for topics that has been
	// consumed already.
	gc.multiplexersMu.Lock()
	defer gc.multiplexersMu.Unlock()
	for topic, mux := range gc.multiplexers {
		gc.rewireMuxAsync(topic, &wg, mux, topicConsumers[topic], assignedPartitions[topic])
	}
	// Start consuming partitions for topics that has not been consumed before.
	for topic, assignedTopicPartitions := range assignedPartitions {
		tc := topicConsumers[topic]
		mux := gc.multiplexers[topic]
		if tc == nil || mux != nil {
			continue
		}
		topic := topic
		spawnInFn := func(partition int32) multiplexer.In {
			return partitioncsm.Spawn(gc.actDesc, gc.group, topic, partition,
				gc.cfg, gc.subscriber, gc.msgFetcherF, gc.offsetMgrF)
		}
		mux = multiplexer.New(gc.actDesc, spawnInFn)
		gc.rewireMuxAsync(topic, &wg, mux, tc, assignedTopicPartitions)
		gc.multiplexers[topic] = mux
	}
	wg.Wait()
	// Clean up gears for topics that do not have assigned partitions anymore.
	for topic, mux := range gc.multiplexers {
		if !mux.IsRunning() {
			delete(gc.multiplexers, topic)
		}
	}
	// Notify the caller that rebalancing has completed successfully.
	rebalanceResultCh <- nil
	return
}

// rewireMuxAsync calls muxInputs in another goroutine.
func (gc *T) rewireMuxAsync(topic string, wg *sync.WaitGroup, mux *multiplexer.T, tc *topiccsm.T, assigned []int32) {
	actor.Spawn(gc.actDesc.NewChild("rewire", topic), wg, func() {
		if tc == nil {
			// Parameter output is an interface type, therefore nil should be
			// passed explicitly.
			mux.WireUp(nil, nil)
			return
		}
		mux.WireUp(tc, assigned)
	})
}

// resolvePartitions given topic subscriptions of all consumer group members,
// resolves what topic partitions are assigned to the specified group member.
func (gc *T) resolvePartitions(subscriptions map[string][]string,
	topicPartitionsFn func(string) ([]int32, error)) (map[string][]int32, error,
) {
	// Convert members->topics to topic->members map.
	topicsToMembers := make(map[string][]string)
	for groupMemberID, topics := range subscriptions {
		for _, topic := range topics {
			topicsToMembers[topic] = append(topicsToMembers[topic], groupMemberID)
		}
	}
	// Create a set of topics this consumer group member subscribed to.
	subscribedTopics := make(map[string]bool)
	for _, topic := range subscriptions[gc.cfg.ClientID] {
		subscribedTopics[topic] = true
	}
	// Resolve new partition assignments for all subscribed topics.
	assignedPartitions := make(map[string][]int32)
	for topic := range subscribedTopics {
		topicPartitions, err := topicPartitionsFn(topic)
		if err != nil {
			return nil, errors.Wrapf(err, "failed to get partition list, topic=%s", topic)
		}
		subscribersToPartitions := assignTopicPartitions(topicPartitions, topicsToMembers[topic])
		assignedTopicPartitions := subscribersToPartitions[gc.cfg.ClientID]
		if len(assignedTopicPartitions) > 0 {
			assignedPartitions[topic] = assignedTopicPartitions
		}
	}
	return assignedPartitions, nil
}

// assignTopicPartitions divides topic partitions among all consumer group
// members subscribed to the topic. The algorithm used closely resembles the
// one implemented by the standard Java High-Level consumer
// (see http://kafka.apache.org/documentation.html#distributionimpl and scroll
// down to *Consumer registration algorithm*) except it does not take in account
// how partitions are distributed among brokers.
func assignTopicPartitions(partitions []int32, subscribers []string) map[string][]int32 {
	partitionCount := len(partitions)
	subscriberCount := len(subscribers)
	if partitionCount == 0 || subscriberCount == 0 {
		return nil
	}
	sort.Slice(partitions, func(i, j int) bool { return partitions[i] < partitions[j] })
	sort.Strings(subscribers)

	subscribersToPartitions := make(map[string][]int32, subscriberCount)
	partitionsPerSubscriber := partitionCount / subscriberCount
	extra := partitionCount - subscriberCount*partitionsPerSubscriber

	begin := 0
	for _, groupMemberID := range subscribers {
		end := begin + partitionsPerSubscriber
		if extra != 0 {
			end++
			extra--
		}
		assigned := partitions[begin:end]
		if len(assigned) > 0 {
			subscribersToPartitions[groupMemberID] = partitions[begin:end]
		}
		begin = end
	}
	return subscribersToPartitions
}

func listTopics(topicConsumers map[string]*topiccsm.T) []string {
	topics := make([]string, 0, len(topicConsumers))
	for topic := range topicConsumers {
		topics = append(topics, topic)
	}
	return topics
}
<|endoftext|>


<commit_before>package app

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/config"
	"github.com/urfave/cli"
)

// GlobalFlags returns a new set of global flags for the client.
// This is where they are defined.
func GlobalFlags() (flags []cli.Flag) {
	return []cli.Flag{
		cli.StringFlag{
			Name:  "account",
			Usage: "account name to use when no other accounts are specified",
		},
		cli.StringFlag{
			Name:  "api-endpoint",
			Usage: "URL where the domains service can be found. Set to blank in environments without a domains service.",
		},
		cli.StringFlag{
			Name:  "auth-endpoint",
			Usage: "URL where the auth service can be found",
		},
		cli.StringFlag{
			Name:  "billing-endpoint",
			Usage: "URL where bmbilling can be found. Set to blank in environments without bmbilling",
		},
		cli.BoolFlag{
			Name:  "admin",
			Usage: "allows admin commands in the client. see bytemark --admin --help",
		},
		cli.BoolFlag{
			Name:  "yubikey",
			Usage: "use a yubikey to authenticate",
		},
		cli.StringFlag{
			Name:  "impersonate",
			Usage: "a user to request impersonation of",
		},
		cli.IntFlag{
			Name:  "debug-level",
			Usage: "how much debug output to print to the terminal",
		},
		cli.StringFlag{
			Name:  "endpoint",
			Usage: "URL of the brain",
		},
		cli.StringFlag{
			Name:  "config-dir",
			Usage: "directory in which bytemark-client's configuration resides. see bytemark help config, bytemark help profiles",
		},
		cli.StringFlag{
			Name:  "spp-endpoint",
			Usage: "URL of SPP. set to blank in environments without an SPP.",
		},
		cli.StringFlag{
			Name:  "output-format",
			Usage: "The output format to use. Currently defined output formats are human (default for most commands), json (machine readable format), table (human-readable table format)",
		},
		cli.StringFlag{
			Name:  "user",
			Usage: "user you wish to log in as",
		},
		cli.StringFlag{
			Name:  "yubikey-otp",
			Usage: "one-time password from your yubikey to use to login",
		},
		cli.IntFlag{
			Name:  "session-validity",
			Usage: "seconds until your session is automatically invalidated (max 3600)",
			Value: config.DefaultSessionValidity,
			// TODO(telyn): add more defaults to these flags
		},
	}
}
<commit_msg>Sort global flags into endpoints and config sections<commit_after>package app

import (
	"github.com/BytemarkHosting/bytemark-client/cmd/bytemark/config"
	"github.com/urfave/cli"
)

// GlobalFlags returns a new set of global flags for the client.
// This is where they are defined.
func GlobalFlags() (flags []cli.Flag) {
	return []cli.Flag{
		cli.StringFlag{
			Name:  "api-endpoint",
			Usage: "URL where the domains service can be found. Set to blank in environments without a domains service.",
		},
		cli.StringFlag{
			Name:  "auth-endpoint",
			Usage: "URL where the auth service can be found",
		},
		cli.StringFlag{
			Name:  "billing-endpoint",
			Usage: "URL where bmbilling can be found. Set to blank in environments without bmbilling",
		},
		cli.StringFlag{
			Name:  "endpoint",
			Usage: "URL of the brain",
		},
		cli.StringFlag{
			Name:  "spp-endpoint",
			Usage: "URL of SPP. set to blank in environments without an SPP.",
		},
		cli.StringFlag{
			Name:  "account",
			Usage: "account name to use when no other accounts are specified",
		},
		cli.BoolFlag{
			Name:  "admin",
			Usage: "allows admin commands in the client. see bytemark --admin --help",
		},
		cli.StringFlag{
			Name:  "config-dir",
			Usage: "directory in which bytemark-client's configuration resides. see bytemark help config, bytemark help profiles",
		},
		cli.IntFlag{
			Name:  "debug-level",
			Usage: "how much debug output to print to the terminal",
		},
		cli.StringFlag{
			Name:  "impersonate",
			Usage: "a user to request impersonation of",
		},
		cli.StringFlag{
			Name:  "output-format",
			Usage: "The output format to use. Currently defined output formats are human (default for most commands), json (machine readable format), table (human-readable table format)",
		},
		cli.IntFlag{
			Name:  "session-validity",
			Usage: "seconds until your session is automatically invalidated (max 3600)",
			Value: config.DefaultSessionValidity,
			// TODO(telyn): add more defaults to these flags
		},
		cli.StringFlag{
			Name:  "user",
			Usage: "user you wish to log in as",
		},
		cli.BoolFlag{
			Name:  "yubikey",
			Usage: "use a yubikey to authenticate",
		},
		cli.StringFlag{
			Name:  "yubikey-otp",
			Usage: "one-time password from your yubikey to use to login",
		},
	}
}
<|endoftext|>


<commit_before>// +build !windows

/*
   Copyright The containerd 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 run

import (
	gocontext "context"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cmd/ctr/commands"
	"github.com/containerd/containerd/contrib/nvidia"
	"github.com/containerd/containerd/oci"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
)

// NewContainer creates a new container
func NewContainer(ctx gocontext.Context, client *containerd.Client, context *cli.Context) (containerd.Container, error) {
	var (
		id     string
		config = context.IsSet("config")
	)
	if config {
		id = context.Args().First()
	} else {
		id = context.Args().Get(1)
	}

	if raw := context.String("checkpoint"); raw != "" {
		im, err := client.GetImage(ctx, raw)
		if err != nil {
			return nil, err
		}
		return client.NewContainer(ctx, id, containerd.WithCheckpoint(im, id), containerd.WithRuntime(context.String("runtime"), nil))
	}

	var (
		opts  []oci.SpecOpts
		cOpts []containerd.NewContainerOpts
		spec  containerd.NewContainerOpts
	)

	if config {
		opts = append(opts, oci.WithSpecFromFile(context.String("config")))
	} else {
		var (
			ref = context.Args().First()
			//for container's id is Args[1]
			args = context.Args()[2:]
		)
		opts = append(opts, oci.WithDefaultSpec(), oci.WithDefaultUnixDevices)
		opts = append(opts, oci.WithEnv(context.StringSlice("env")))
		opts = append(opts, withMounts(context))

		if context.Bool("rootfs") {
			rootfs, err := filepath.Abs(ref)
			if err != nil {
				return nil, err
			}
			opts = append(opts, oci.WithRootFSPath(rootfs))
		} else {
			snapshotter := context.String("snapshotter")
			image, err := client.GetImage(ctx, ref)
			if err != nil {
				return nil, err
			}
			unpacked, err := image.IsUnpacked(ctx, snapshotter)
			if err != nil {
				return nil, err
			}
			if !unpacked {
				if err := image.Unpack(ctx, snapshotter); err != nil {
					return nil, err
				}
			}
			opts = append(opts, oci.WithImageConfig(image))
			cOpts = append(cOpts,
				containerd.WithImage(image),
				containerd.WithSnapshotter(snapshotter),
				// Even when "readonly" is set, we don't use KindView snapshot here. (#1495)
				// We pass writable snapshot to the OCI runtime, and the runtime remounts it as read-only,
				// after creating some mount points on demand.
				containerd.WithNewSnapshot(id, image))
		}
		if context.Bool("readonly") {
			opts = append(opts, oci.WithRootFSReadonly())
		}
		if len(args) > 0 {
			opts = append(opts, oci.WithProcessArgs(args...))
		}
		if cwd := context.String("cwd"); cwd != "" {
			opts = append(opts, oci.WithProcessCwd(cwd))
		}
		if context.Bool("tty") {
			opts = append(opts, oci.WithTTY)
		}
		if context.Bool("privileged") {
			opts = append(opts, oci.WithPrivileged)
		}
		if context.Bool("net-host") {
			opts = append(opts, oci.WithHostNamespace(specs.NetworkNamespace), oci.WithHostHostsFile, oci.WithHostResolvconf)
		}

		joinNs := context.StringSlice("with-ns")
		for _, ns := range joinNs {
			parts := strings.Split(ns, ":")
			if len(parts) != 2 {
				return nil, errors.New("joining a Linux namespace using --with-ns requires the format 'nstype:path'")
			}
			if !validNamespace(parts[0]) {
				return nil, errors.New("the Linux namespace type specified in --with-ns is not valid: " + parts[0])
			}
			opts = append(opts, oci.WithLinuxNamespace(specs.LinuxNamespace{
				Type: specs.LinuxNamespaceType(parts[0]),
				Path: parts[1],
			}))
		}
		if context.IsSet("gpus") {
			opts = append(opts, nvidia.WithGPUs(nvidia.WithDevices(context.Int("gpus")), nvidia.WithAllCapabilities))
		}
		if context.IsSet("allow-new-privs") {
			opts = append(opts, oci.WithNewPrivileges)
		}
	}

	cOpts = append(cOpts, containerd.WithContainerLabels(commands.LabelArgs(context.StringSlice("label"))))
	cOpts = append(cOpts, containerd.WithRuntime(context.String("runtime"), nil))

	var s specs.Spec
	spec = containerd.WithSpec(&s, opts...)

	cOpts = append(cOpts, spec)

	// oci.WithImageConfig (WithUsername, WithUserID) depends on rootfs snapshot for resolving /etc/passwd.
	// So cOpts needs to have precedence over opts.
	// TODO: WithUsername, WithUserID should additionally support non-snapshot rootfs
	return client.NewContainer(ctx, id, cOpts...)
}

func getNewTaskOpts(context *cli.Context) []containerd.NewTaskOpts {
	if context.Bool("no-pivot") {
		return []containerd.NewTaskOpts{containerd.WithNoPivotRoot}
	}
	return nil
}

func validNamespace(ns string) bool {
	linuxNs := specs.LinuxNamespaceType(ns)
	switch linuxNs {
	case specs.PIDNamespace,
		specs.NetworkNamespace,
		specs.UTSNamespace,
		specs.MountNamespace,
		specs.UserNamespace,
		specs.IPCNamespace,
		specs.CgroupNamespace:
		return true
	default:
		return false
	}
}
<commit_msg>Remove a TODO from the code comments that is complete<commit_after>// +build !windows

/*
   Copyright The containerd 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 run

import (
	gocontext "context"
	"path/filepath"
	"strings"

	"github.com/containerd/containerd"
	"github.com/containerd/containerd/cmd/ctr/commands"
	"github.com/containerd/containerd/contrib/nvidia"
	"github.com/containerd/containerd/oci"
	specs "github.com/opencontainers/runtime-spec/specs-go"
	"github.com/pkg/errors"
	"github.com/urfave/cli"
)

// NewContainer creates a new container
func NewContainer(ctx gocontext.Context, client *containerd.Client, context *cli.Context) (containerd.Container, error) {
	var (
		id     string
		config = context.IsSet("config")
	)
	if config {
		id = context.Args().First()
	} else {
		id = context.Args().Get(1)
	}

	if raw := context.String("checkpoint"); raw != "" {
		im, err := client.GetImage(ctx, raw)
		if err != nil {
			return nil, err
		}
		return client.NewContainer(ctx, id, containerd.WithCheckpoint(im, id), containerd.WithRuntime(context.String("runtime"), nil))
	}

	var (
		opts  []oci.SpecOpts
		cOpts []containerd.NewContainerOpts
		spec  containerd.NewContainerOpts
	)

	if config {
		opts = append(opts, oci.WithSpecFromFile(context.String("config")))
	} else {
		var (
			ref = context.Args().First()
			//for container's id is Args[1]
			args = context.Args()[2:]
		)
		opts = append(opts, oci.WithDefaultSpec(), oci.WithDefaultUnixDevices)
		opts = append(opts, oci.WithEnv(context.StringSlice("env")))
		opts = append(opts, withMounts(context))

		if context.Bool("rootfs") {
			rootfs, err := filepath.Abs(ref)
			if err != nil {
				return nil, err
			}
			opts = append(opts, oci.WithRootFSPath(rootfs))
		} else {
			snapshotter := context.String("snapshotter")
			image, err := client.GetImage(ctx, ref)
			if err != nil {
				return nil, err
			}
			unpacked, err := image.IsUnpacked(ctx, snapshotter)
			if err != nil {
				return nil, err
			}
			if !unpacked {
				if err := image.Unpack(ctx, snapshotter); err != nil {
					return nil, err
				}
			}
			opts = append(opts, oci.WithImageConfig(image))
			cOpts = append(cOpts,
				containerd.WithImage(image),
				containerd.WithSnapshotter(snapshotter),
				// Even when "readonly" is set, we don't use KindView snapshot here. (#1495)
				// We pass writable snapshot to the OCI runtime, and the runtime remounts it as read-only,
				// after creating some mount points on demand.
				containerd.WithNewSnapshot(id, image))
		}
		if context.Bool("readonly") {
			opts = append(opts, oci.WithRootFSReadonly())
		}
		if len(args) > 0 {
			opts = append(opts, oci.WithProcessArgs(args...))
		}
		if cwd := context.String("cwd"); cwd != "" {
			opts = append(opts, oci.WithProcessCwd(cwd))
		}
		if context.Bool("tty") {
			opts = append(opts, oci.WithTTY)
		}
		if context.Bool("privileged") {
			opts = append(opts, oci.WithPrivileged)
		}
		if context.Bool("net-host") {
			opts = append(opts, oci.WithHostNamespace(specs.NetworkNamespace), oci.WithHostHostsFile, oci.WithHostResolvconf)
		}

		joinNs := context.StringSlice("with-ns")
		for _, ns := range joinNs {
			parts := strings.Split(ns, ":")
			if len(parts) != 2 {
				return nil, errors.New("joining a Linux namespace using --with-ns requires the format 'nstype:path'")
			}
			if !validNamespace(parts[0]) {
				return nil, errors.New("the Linux namespace type specified in --with-ns is not valid: " + parts[0])
			}
			opts = append(opts, oci.WithLinuxNamespace(specs.LinuxNamespace{
				Type: specs.LinuxNamespaceType(parts[0]),
				Path: parts[1],
			}))
		}
		if context.IsSet("gpus") {
			opts = append(opts, nvidia.WithGPUs(nvidia.WithDevices(context.Int("gpus")), nvidia.WithAllCapabilities))
		}
		if context.IsSet("allow-new-privs") {
			opts = append(opts, oci.WithNewPrivileges)
		}
	}

	cOpts = append(cOpts, containerd.WithContainerLabels(commands.LabelArgs(context.StringSlice("label"))))
	cOpts = append(cOpts, containerd.WithRuntime(context.String("runtime"), nil))

	var s specs.Spec
	spec = containerd.WithSpec(&s, opts...)

	cOpts = append(cOpts, spec)

	// oci.WithImageConfig (WithUsername, WithUserID) depends on access to rootfs for resolving via
	// the /etc/{passwd,group} files. So cOpts needs to have precedence over opts.
	return client.NewContainer(ctx, id, cOpts...)
}

func getNewTaskOpts(context *cli.Context) []containerd.NewTaskOpts {
	if context.Bool("no-pivot") {
		return []containerd.NewTaskOpts{containerd.WithNoPivotRoot}
	}
	return nil
}

func validNamespace(ns string) bool {
	linuxNs := specs.LinuxNamespaceType(ns)
	switch linuxNs {
	case specs.PIDNamespace,
		specs.NetworkNamespace,
		specs.UTSNamespace,
		specs.MountNamespace,
		specs.UserNamespace,
		specs.IPCNamespace,
		specs.CgroupNamespace:
		return true
	default:
		return false
	}
}
<|endoftext|>


<commit_before>package ingest

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"sync"

	"github.com/golang/snappy"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/schema/msg"
	"github.com/grafana/metrictank/stats"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	"github.com/grafana/metrictank/api/models"
	"github.com/grafana/metrictank/publish"
	log "github.com/sirupsen/logrus"
)

var (
	metricsValid    = stats.NewCounterRate32("metrics.http.valid")    // valid metrics received (not necessarily published)
	metricsRejected = stats.NewCounterRate32("metrics.http.rejected") // invalid metrics received

	metricsTSLock    = &sync.Mutex{}
	metricsTimestamp = make(map[int]*stats.Range32)

	discardedSamples = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "gateway",
			Name:      "invalid_samples_total",
			Help:      "The total number of samples that were discarded because they are invalid.",
		},
		[]string{"reason", "org"},
	)
)

func getMetricsTimestampStat(org int) *stats.Range32 {
	metricsTSLock.Lock()
	metricTimestamp, ok := metricsTimestamp[org]
	if !ok {
		metricTimestamp = stats.NewRange32(fmt.Sprintf("metrics.timestamp.http.%d", org)) // min/max timestamps seen in each interval
		metricsTimestamp[org] = metricTimestamp
	}
	metricsTSLock.Unlock()
	return metricTimestamp
}

func Metrics(w http.ResponseWriter, r *http.Request) {
	contentType := r.Header.Get("Content-Type")
	switch contentType {
	case "rt-metric-binary":
		metricsBinary(w, r, false)
	case "rt-metric-binary-snappy":
		metricsBinary(w, r, true)
	case "application/json":
		metricsJson(w, r)
	default:
		w.WriteHeader(400)
		fmt.Fprintf(w, "unknown content-type: %s", contentType)
	}
}

type discardsByReason map[string]int
type discardsByOrg map[int]discardsByReason

func (dbo discardsByOrg) Add(org int, reason string) {
	dbr, ok := dbo[org]
	if !ok {
		dbr = make(discardsByReason)
	}
	dbr[reason]++
	dbo[org] = dbr
}

func prepareIngest(ctx *models.Context, in []*schema.MetricData, toPublish []*schema.MetricData) ([]*schema.MetricData, MetricsResponse) {
	resp := NewMetricsResponse()
	promDiscards := make(discardsByOrg)

	var metricTimestamp *stats.Range32
	if !ctx.IsAdmin {
		metricTimestamp = getMetricsTimestampStat(ctx.ID)
	}

	for i, m := range in {
		if !ctx.IsAdmin {
			m.OrgId = ctx.ID
		}
		if m.Mtype == "" {
			m.Mtype = "gauge"
		}
		if err := m.Validate(); err != nil {
			log.Debugf("received invalid metric: %v %v %v", m.Name, m.OrgId, m.Tags)
			resp.AddInvalid(err, i)
			promDiscards.Add(m.OrgId, err.Error())
			continue
		}
		if ctx.IsAdmin {
			metricTimestamp = getMetricsTimestampStat(m.OrgId)
		} else {
			m.SetId()
		}
		metricTimestamp.ValueUint32(uint32(m.Time))
		toPublish = append(toPublish, m)
	}

	// track invalid/discards in graphite and prometheus
	metricsRejected.Add(resp.Invalid)
	metricsValid.Add(len(toPublish))
	for org, promDiscardsByOrg := range promDiscards {
		for reason, cnt := range promDiscardsByOrg {
			discardedSamples.WithLabelValues(reason, strconv.Itoa(org)).Add(float64(cnt))
		}
	}
	return toPublish, resp
}

func metricsJson(w http.ResponseWriter, r *http.Request) {
	if r.Body == nil {
		w.WriteHeader(400)
		fmt.Fprintf(w, "no data included in request.")
		return
	}
	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		select {
		case <-r.Context().Done():
			w.WriteHeader(499)
			fmt.Fprintf(w, "request canceled")
		default:
			log.Errorf("unable to read request body. %s", err)
			w.WriteHeader(500)
			fmt.Fprintf(w,"unable to read request body. %s", err)
		}
		return
	}
	metrics := make([]*schema.MetricData, 0)
	err = json.Unmarshal(body, &metrics)
	if err != nil {
		w.WriteHeader(400)
		fmt.Fprintf(w, "unable to parse request body. %s", err)
		return
	}

	toPublish := make([]*schema.MetricData, 0, len(metrics))
	toPublish, resp := prepareIngest(ctx, metrics, toPublish)

	select {
	case <-r.Context().Done():
		w.WriteHeader(499)
		fmt.Fprintf(w, "request canceled")
		return
	default:
	}

	err = publish.Publish(toPublish)
	if err != nil {
		log.Errorf("failed to publish metrics. %s", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "failed to publish metrics. %s", err)
		return
	}

	// track published in the response (it already has discards)
	resp.Published = len(toPublish)
	w.WriteHeader(200)
	json.NewEncoder(w).Encode(resp)
}

func metricsBinary(w http.ResponseWriter, r *http.Request, compressed bool) {
	if r.Body == nil {
		w.WriteHeader(400)
		fmt.Fprintf(w, "no data included in request.")
		return
	}
	var bodyReadCloser io.ReadCloser
	if compressed {
		bodyReadCloser = ioutil.NopCloser(snappy.NewReader(r.Body))
	} else {
		bodyReadCloser = r.Body
	}
	defer bodyReadCloser.Close()

	body, err := ioutil.ReadAll(bodyReadCloser)
	if err != nil {
		select {
		case <-r.Context().Done():
			w.WriteHeader(499)
			fmt.Fprintf(w, "request canceled")
		default:
			log.Errorf("unable to read request body. %s", err)
			w.WriteHeader(500)
			fmt.Fprintf(w, "unable to read request body. %s", err)
		}
		return
	}
	metricData := new(msg.MetricData)
	err = metricData.InitFromMsg(body)
	if err != nil {
		log.Errorf("payload not metricData. %s", err)
		w.WriteHeader(400)
		fmt.Fprintf(w, "payload not metricData. %s", err)
		return
	}

	err = metricData.DecodeMetricData()
	if err != nil {
		log.Errorf("failed to unmarshal metricData. %s", err)
		w.WriteHeader(400)
		fmt.Fprintf(w, "failed to unmarshal metricData. %s", err)
		return
	}

	toPublish := make([]*schema.MetricData, 0, len(metricData.Metrics))
	toPublish, resp := prepareIngest(ctx, metricData.Metrics, toPublish)

	select {
	case <-r.Context().Done():
		w.WriteHeader(499)
		fmt.Fprintf(w, "request canceled")
		return
	default:
	}

	err = publish.Publish(toPublish)
	if err != nil {
		log.Errorf("failed to publish metrics. %s", err)
		w.WriteHeader(500)
		fmt.Fprintf(w,"failed to publish metrics. %s", err)
		return
	}

	// track published in the response (it already has discards)
	resp.Published = len(toPublish)
	w.WriteHeader(200)
	json.NewEncoder(w).Encode(resp)
}
<commit_msg>There's no auth in this version, so we should assume admin = true<commit_after>package ingest

import (
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"strconv"
	"sync"

	"github.com/golang/snappy"
	"github.com/grafana/metrictank/publish"
	"github.com/grafana/metrictank/schema"
	"github.com/grafana/metrictank/schema/msg"
	"github.com/grafana/metrictank/stats"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promauto"
	log "github.com/sirupsen/logrus"
)

var (
	metricsValid    = stats.NewCounterRate32("metrics.http.valid")    // valid metrics received (not necessarily published)
	metricsRejected = stats.NewCounterRate32("metrics.http.rejected") // invalid metrics received

	metricsTSLock    = &sync.Mutex{}
	metricsTimestamp = make(map[int]*stats.Range32)

	discardedSamples = promauto.NewCounterVec(
		prometheus.CounterOpts{
			Namespace: "gateway",
			Name:      "invalid_samples_total",
			Help:      "The total number of samples that were discarded because they are invalid.",
		},
		[]string{"reason", "org"},
	)
)

func getMetricsTimestampStat(org int) *stats.Range32 {
	metricsTSLock.Lock()
	metricTimestamp, ok := metricsTimestamp[org]
	if !ok {
		metricTimestamp = stats.NewRange32(fmt.Sprintf("metrics.timestamp.http.%d", org)) // min/max timestamps seen in each interval
		metricsTimestamp[org] = metricTimestamp
	}
	metricsTSLock.Unlock()
	return metricTimestamp
}

func Metrics(w http.ResponseWriter, r *http.Request) {
	contentType := r.Header.Get("Content-Type")
	switch contentType {
	case "rt-metric-binary":
		metricsBinary(w, r, false)
	case "rt-metric-binary-snappy":
		metricsBinary(w, r, true)
	case "application/json":
		metricsJson(w, r)
	default:
		w.WriteHeader(400)
		fmt.Fprintf(w, "unknown content-type: %s", contentType)
	}
}

type discardsByReason map[string]int
type discardsByOrg map[int]discardsByReason

func (dbo discardsByOrg) Add(org int, reason string) {
	dbr, ok := dbo[org]
	if !ok {
		dbr = make(discardsByReason)
	}
	dbr[reason]++
	dbo[org] = dbr
}

func prepareIngest(in []*schema.MetricData, toPublish []*schema.MetricData) ([]*schema.MetricData, MetricsResponse) {
	resp := NewMetricsResponse()
	promDiscards := make(discardsByOrg)

	var metricTimestamp *stats.Range32

	for i, m := range in {
		if m.Mtype == "" {
			m.Mtype = "gauge"
		}
		if err := m.Validate(); err != nil {
			log.Debugf("received invalid metric: %v %v %v", m.Name, m.OrgId, m.Tags)
			resp.AddInvalid(err, i)
			promDiscards.Add(m.OrgId, err.Error())
			continue
		}
		metricTimestamp = getMetricsTimestampStat(m.OrgId)
		metricTimestamp.ValueUint32(uint32(m.Time))
		toPublish = append(toPublish, m)
	}

	// track invalid/discards in graphite and prometheus
	metricsRejected.Add(resp.Invalid)
	metricsValid.Add(len(toPublish))
	for org, promDiscardsByOrg := range promDiscards {
		for reason, cnt := range promDiscardsByOrg {
			discardedSamples.WithLabelValues(reason, strconv.Itoa(org)).Add(float64(cnt))
		}
	}
	return toPublish, resp
}

func metricsJson(w http.ResponseWriter, r *http.Request) {
	if r.Body == nil {
		w.WriteHeader(400)
		fmt.Fprintf(w, "no data included in request.")
		return
	}
	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		select {
		case <-r.Context().Done():
			w.WriteHeader(499)
			fmt.Fprintf(w, "request canceled")
		default:
			log.Errorf("unable to read request body. %s", err)
			w.WriteHeader(500)
			fmt.Fprintf(w,"unable to read request body. %s", err)
		}
		return
	}
	metrics := make([]*schema.MetricData, 0)
	err = json.Unmarshal(body, &metrics)
	if err != nil {
		w.WriteHeader(400)
		fmt.Fprintf(w, "unable to parse request body. %s", err)
		return
	}

	toPublish := make([]*schema.MetricData, 0, len(metrics))
	toPublish, resp := prepareIngest(metrics, toPublish)

	select {
	case <-r.Context().Done():
		w.WriteHeader(499)
		fmt.Fprintf(w, "request canceled")
		return
	default:
	}

	err = publish.Publish(toPublish)
	if err != nil {
		log.Errorf("failed to publish metrics. %s", err)
		w.WriteHeader(500)
		fmt.Fprintf(w, "failed to publish metrics. %s", err)
		return
	}

	// track published in the response (it already has discards)
	resp.Published = len(toPublish)
	w.WriteHeader(200)
	json.NewEncoder(w).Encode(resp)
}

func metricsBinary(w http.ResponseWriter, r *http.Request, compressed bool) {
	if r.Body == nil {
		w.WriteHeader(400)
		fmt.Fprintf(w, "no data included in request.")
		return
	}
	var bodyReadCloser io.ReadCloser
	if compressed {
		bodyReadCloser = ioutil.NopCloser(snappy.NewReader(r.Body))
	} else {
		bodyReadCloser = r.Body
	}
	defer bodyReadCloser.Close()

	body, err := ioutil.ReadAll(bodyReadCloser)
	if err != nil {
		select {
		case <-r.Context().Done():
			w.WriteHeader(499)
			fmt.Fprintf(w, "request canceled")
		default:
			log.Errorf("unable to read request body. %s", err)
			w.WriteHeader(500)
			fmt.Fprintf(w, "unable to read request body. %s", err)
		}
		return
	}
	metricData := new(msg.MetricData)
	err = metricData.InitFromMsg(body)
	if err != nil {
		log.Errorf("payload not metricData. %s", err)
		w.WriteHeader(400)
		fmt.Fprintf(w, "payload not metricData. %s", err)
		return
	}

	err = metricData.DecodeMetricData()
	if err != nil {
		log.Errorf("failed to unmarshal metricData. %s", err)
		w.WriteHeader(400)
		fmt.Fprintf(w, "failed to unmarshal metricData. %s", err)
		return
	}

	toPublish := make([]*schema.MetricData, 0, len(metricData.Metrics))
	toPublish, resp := prepareIngest(metricData.Metrics, toPublish)

	select {
	case <-r.Context().Done():
		w.WriteHeader(499)
		fmt.Fprintf(w, "request canceled")
		return
	default:
	}

	err = publish.Publish(toPublish)
	if err != nil {
		log.Errorf("failed to publish metrics. %s", err)
		w.WriteHeader(500)
		fmt.Fprintf(w,"failed to publish metrics. %s", err)
		return
	}

	// track published in the response (it already has discards)
	resp.Published = len(toPublish)
	w.WriteHeader(200)
	json.NewEncoder(w).Encode(resp)
}
<|endoftext|>


<commit_before>package riemann

import (
	"fmt"
	"testing"
	"time"

	"github.com/amir/raidman"
	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/testutil"
	"github.com/stretchr/testify/require"
)

func TestAttributes(t *testing.T) {
	tags := map[string]string{"tag1": "value1", "tag2": "value2"}

	r := &Riemann{}
	require.Equal(t,
		map[string]string{"tag1": "value1", "tag2": "value2"},
		r.attributes("test", tags))

	// enable measurement as attribute, should now be included
	r.MeasurementAsAttribute = true
	require.Equal(t,
		map[string]string{"tag1": "value1", "tag2": "value2", "measurement": "test"},
		r.attributes("test", tags))
}

func TestService(t *testing.T) {
	r := &Riemann{
		Separator: "/",
	}
	require.Equal(t, "test/value", r.service("test", "value"))

	// enable measurement as attribute, should not be part of service name anymore
	r.MeasurementAsAttribute = true
	require.Equal(t, "value", r.service("test", "value"))
}

func TestTags(t *testing.T) {
	tags := map[string]string{"tag1": "value1", "tag2": "value2"}

	// all tag values plus additional tag should be present
	r := &Riemann{
		Tags: []string{"test"},
	}
	require.Equal(t,
		[]string{"test", "value1", "value2"},
		r.tags(tags))

	// only tag2 value plus additional tag should be present
	r.TagKeys = []string{"tag2"}
	require.Equal(t,
		[]string{"test", "value2"},
		r.tags(tags))

	// only tag1 value should be present
	r.Tags = nil
	r.TagKeys = []string{"tag1"}
	require.Equal(t,
		[]string{"value1"},
		r.tags(tags))
}

func TestMetricEvents(t *testing.T) {
	r := &Riemann{
		TTL:                    20.0,
		Separator:              "/",
		MeasurementAsAttribute: false,
		DescriptionText:        "metrics from telegraf",
		Tags:                   []string{"telegraf"},
	}

	// build a single event
	metric, _ := telegraf.NewMetric(
		"test1",
		map[string]string{"tag1": "value1", "host": "abc123"},
		map[string]interface{}{"value": 5.6},
		time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC),
	)

	events := r.buildRiemannEvents(metric)
	require.Len(t, events, 1)

	// is event as expected?
	expectedEvent := &raidman.Event{
		Ttl:         20.0,
		Time:        1257894000,
		Tags:        []string{"telegraf", "value1"},
		Host:        "abc123",
		State:       "",
		Service:     "test1/value",
		Metric:      5.6,
		Description: "metrics from telegraf",
		Attributes:  map[string]string{"tag1": "value1"},
	}
	require.Equal(t, expectedEvent, events[0])

	// build 2 events
	metric, _ = telegraf.NewMetric(
		"test2",
		map[string]string{"host": "xyz987"},
		map[string]interface{}{"point": 1},
		time.Date(2012, time.November, 2, 3, 0, 0, 0, time.UTC),
	)

	events = append(events, r.buildRiemannEvents(metric)...)
	require.Len(t, events, 2)

	// first event should still be the same
	require.Equal(t, expectedEvent, events[0])

	// second event
	expectedEvent = &raidman.Event{
		Ttl:         20.0,
		Time:        1351825200,
		Tags:        []string{"telegraf"},
		Host:        "xyz987",
		State:       "",
		Service:     "test2/point",
		Metric:      int64(1),
		Description: "metrics from telegraf",
		Attributes:  map[string]string{},
	}
	require.Equal(t, expectedEvent, events[1])
}

func TestStateEvents(t *testing.T) {
	r := &Riemann{
		MeasurementAsAttribute: true,
	}

	// string metrics will be skipped unless explicitly enabled
	metric, _ := telegraf.NewMetric(
		"test",
		map[string]string{"host": "host"},
		map[string]interface{}{"value": "running"},
		time.Date(2015, time.November, 9, 22, 0, 0, 0, time.UTC),
	)

	events := r.buildRiemannEvents(metric)
	// no event should be present
	require.Len(t, events, 0)

	// enable string metrics as event states
	r.StringAsState = true
	events = r.buildRiemannEvents(metric)
	require.Len(t, events, 1)

	// is event as expected?
	expectedEvent := &raidman.Event{
		Ttl:         0,
		Time:        1447106400,
		Tags:        nil,
		Host:        "host",
		State:       "running",
		Service:     "value",
		Metric:      nil,
		Description: "",
		Attributes:  map[string]string{"measurement": "test"},
	}
	require.Equal(t, expectedEvent, events[0])
}

func TestConnectAndWrite(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration test in short mode")
	}

	r := &Riemann{
		URL:                    fmt.Sprintf("tcp://%s:5555", testutil.GetLocalHost()),
		TTL:                    15.0,
		Separator:              "/",
		MeasurementAsAttribute: false,
		StringAsState:          true,
		DescriptionText:        "metrics from telegraf",
		Tags:                   []string{"docker"},
	}

	err := r.Connect()
	require.NoError(t, err)

	err = r.Write(testutil.MockMetrics())
	require.NoError(t, err)

	metrics := make([]telegraf.Metric, 0)
	metrics = append(metrics, testutil.TestMetric(2))
	metrics = append(metrics, testutil.TestMetric(3.456789))
	metrics = append(metrics, testutil.TestMetric(uint(0)))
	metrics = append(metrics, testutil.TestMetric("ok"))
	metrics = append(metrics, testutil.TestMetric("running"))
	err = r.Write(metrics)
	require.NoError(t, err)

	time.Sleep(200 * time.Millisecond)

	// are there any "docker" tagged events in Riemann?
	events, err := r.client.Query(`tagged "docker"`)
	require.NoError(t, err)
	require.NotZero(t, len(events))

	// get Riemann events with state = "running", should be 1 event
	events, err = r.client.Query(`state = "running"`)
	require.NoError(t, err)
	require.Len(t, events, 1)

	// is event as expected?
	require.Equal(t, []string{"docker", "value1"}, events[0].Tags)
	require.Equal(t, "running", events[0].State)
	require.Equal(t, "test1/value", events[0].Service)
	require.Equal(t, "metrics from telegraf", events[0].Description)
	require.Equal(t, map[string]string{"tag1": "value1"}, events[0].Attributes)
}
<commit_msg>Fix riemann output unit tests<commit_after>package riemann

import (
	"fmt"
	"testing"
	"time"

	"github.com/amir/raidman"
	"github.com/influxdata/telegraf"
	"github.com/influxdata/telegraf/metric"
	"github.com/influxdata/telegraf/testutil"
	"github.com/stretchr/testify/require"
)

func TestAttributes(t *testing.T) {
	tags := map[string]string{"tag1": "value1", "tag2": "value2"}

	r := &Riemann{}
	require.Equal(t,
		map[string]string{"tag1": "value1", "tag2": "value2"},
		r.attributes("test", tags))

	// enable measurement as attribute, should now be included
	r.MeasurementAsAttribute = true
	require.Equal(t,
		map[string]string{"tag1": "value1", "tag2": "value2", "measurement": "test"},
		r.attributes("test", tags))
}

func TestService(t *testing.T) {
	r := &Riemann{
		Separator: "/",
	}
	require.Equal(t, "test/value", r.service("test", "value"))

	// enable measurement as attribute, should not be part of service name anymore
	r.MeasurementAsAttribute = true
	require.Equal(t, "value", r.service("test", "value"))
}

func TestTags(t *testing.T) {
	tags := map[string]string{"tag1": "value1", "tag2": "value2"}

	// all tag values plus additional tag should be present
	r := &Riemann{
		Tags: []string{"test"},
	}
	require.Equal(t,
		[]string{"test", "value1", "value2"},
		r.tags(tags))

	// only tag2 value plus additional tag should be present
	r.TagKeys = []string{"tag2"}
	require.Equal(t,
		[]string{"test", "value2"},
		r.tags(tags))

	// only tag1 value should be present
	r.Tags = nil
	r.TagKeys = []string{"tag1"}
	require.Equal(t,
		[]string{"value1"},
		r.tags(tags))
}

func TestMetricEvents(t *testing.T) {
	r := &Riemann{
		TTL:                    20.0,
		Separator:              "/",
		MeasurementAsAttribute: false,
		DescriptionText:        "metrics from telegraf",
		Tags:                   []string{"telegraf"},
	}

	// build a single event
	m, _ := metric.New(
		"test1",
		map[string]string{"tag1": "value1", "host": "abc123"},
		map[string]interface{}{"value": 5.6},
		time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC),
	)

	events := r.buildRiemannEvents(m)
	require.Len(t, events, 1)

	// is event as expected?
	expectedEvent := &raidman.Event{
		Ttl:         20.0,
		Time:        1257894000,
		Tags:        []string{"telegraf", "value1"},
		Host:        "abc123",
		State:       "",
		Service:     "test1/value",
		Metric:      5.6,
		Description: "metrics from telegraf",
		Attributes:  map[string]string{"tag1": "value1"},
	}
	require.Equal(t, expectedEvent, events[0])

	// build 2 events
	m, _ = metric.New(
		"test2",
		map[string]string{"host": "xyz987"},
		map[string]interface{}{"point": 1},
		time.Date(2012, time.November, 2, 3, 0, 0, 0, time.UTC),
	)

	events = append(events, r.buildRiemannEvents(m)...)
	require.Len(t, events, 2)

	// first event should still be the same
	require.Equal(t, expectedEvent, events[0])

	// second event
	expectedEvent = &raidman.Event{
		Ttl:         20.0,
		Time:        1351825200,
		Tags:        []string{"telegraf"},
		Host:        "xyz987",
		State:       "",
		Service:     "test2/point",
		Metric:      int64(1),
		Description: "metrics from telegraf",
		Attributes:  map[string]string{},
	}
	require.Equal(t, expectedEvent, events[1])
}

func TestStateEvents(t *testing.T) {
	r := &Riemann{
		MeasurementAsAttribute: true,
	}

	// string metrics will be skipped unless explicitly enabled
	m, _ := metric.New(
		"test",
		map[string]string{"host": "host"},
		map[string]interface{}{"value": "running"},
		time.Date(2015, time.November, 9, 22, 0, 0, 0, time.UTC),
	)

	events := r.buildRiemannEvents(m)
	// no event should be present
	require.Len(t, events, 0)

	// enable string metrics as event states
	r.StringAsState = true
	events = r.buildRiemannEvents(m)
	require.Len(t, events, 1)

	// is event as expected?
	expectedEvent := &raidman.Event{
		Ttl:         0,
		Time:        1447106400,
		Tags:        nil,
		Host:        "host",
		State:       "running",
		Service:     "value",
		Metric:      nil,
		Description: "",
		Attributes:  map[string]string{"measurement": "test"},
	}
	require.Equal(t, expectedEvent, events[0])
}

func TestConnectAndWrite(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping integration test in short mode")
	}

	r := &Riemann{
		URL:                    fmt.Sprintf("tcp://%s:5555", testutil.GetLocalHost()),
		TTL:                    15.0,
		Separator:              "/",
		MeasurementAsAttribute: false,
		StringAsState:          true,
		DescriptionText:        "metrics from telegraf",
		Tags:                   []string{"docker"},
	}

	err := r.Connect()
	require.NoError(t, err)

	err = r.Write(testutil.MockMetrics())
	require.NoError(t, err)

	metrics := make([]telegraf.Metric, 0)
	metrics = append(metrics, testutil.TestMetric(2))
	metrics = append(metrics, testutil.TestMetric(3.456789))
	metrics = append(metrics, testutil.TestMetric(uint(0)))
	metrics = append(metrics, testutil.TestMetric("ok"))
	metrics = append(metrics, testutil.TestMetric("running"))
	err = r.Write(metrics)
	require.NoError(t, err)

	time.Sleep(200 * time.Millisecond)

	// are there any "docker" tagged events in Riemann?
	events, err := r.client.Query(`tagged "docker"`)
	require.NoError(t, err)
	require.NotZero(t, len(events))

	// get Riemann events with state = "running", should be 1 event
	events, err = r.client.Query(`state = "running"`)
	require.NoError(t, err)
	require.Len(t, events, 1)

	// is event as expected?
	require.Equal(t, []string{"docker", "value1"}, events[0].Tags)
	require.Equal(t, "running", events[0].State)
	require.Equal(t, "test1/value", events[0].Service)
	require.Equal(t, "metrics from telegraf", events[0].Description)
	require.Equal(t, map[string]string{"tag1": "value1"}, events[0].Attributes)
}
<|endoftext|>


<commit_before>// unify-labels is a CLI which will add, rename, or change the color of labels so they match the Jekyll org's requirements.
package main

import (
	"flag"
	"log"
	"strings"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/auto-reply/freeze"
)

var desiredLabels = []*github.Label{
	{Name: github.String("bug"), Color: github.String("d41313")},
	{Name: github.String("discussion"), Color: github.String("006b75")},
	{Name: github.String("documentation"), Color: github.String("006b75")},
	{Name: github.String("enhancement"), Color: github.String("009800")},
	{Name: github.String("feature"), Color: github.String("009800")},
	{Name: github.String("fix"), Color: github.String("eb6420")},
	{Name: github.String(freeze.LabelName), Color: github.String("0052cc")},
	{Name: github.String("github"), Color: github.String("222222")},
	{Name: github.String("has-pull-request"), Color: github.String("fbca04")},
	{Name: github.String("help-wanted"), Color: github.String("fbca04")},
	{Name: github.String("internal"), Color: github.String("ededed")},
	{Name: github.String("pending-feedback"), Color: github.String("fbca04")},
	{Name: github.String("pending-rebase"), Color: github.String("eb6420")},
	{Name: github.String("pinned"), Color: github.String("f3f4d3")},
	{Name: github.String("release"), Color: github.String("d4c5f9")},
	{Name: github.String("security"), Color: github.String("e11d21")},
	{Name: github.String("stale"), Color: github.String("bfd4f2")},
	{Name: github.String("suggestion"), Color: github.String("0052cc")},
	{Name: github.String("support"), Color: github.String("5319e7")},
	{Name: github.String("tests"), Color: github.String("d4c5f9")},
	{Name: github.String("undetermined"), Color: github.String("fe3868")},
	{Name: github.String("ux"), Color: github.String("006b75")},
	{Name: github.String("windows"), Color: github.String("fbca04")},
	{Name: github.String("wont-fix"), Color: github.String("e11d21")},
}
var listOpts = github.ListOptions{PerPage: 100}

func allPossibleNames(name string) []string {
	return []string{
		name,
		strings.Replace(name, "-", "", -1),
		strings.Replace(name, "-", " ", -1),
		strings.ToLower(name),
		strings.Title(name),
		strings.Title(strings.Replace(name, "-", "", -1)),
		strings.Title(strings.Replace(name, "-", " ", -1)),
	}
}

func findLabel(labels []*github.Label, desiredLabel *github.Label) *github.Label {
	possibleNames := allPossibleNames(*desiredLabel.Name)
	for _, possibleName := range possibleNames {
		for _, label := range labels {
			if *label.Name == possibleName {
				return label
			}
		}
	}

	return nil
}

func processRepo(context *ctx.Context, repo *github.Repository, perform bool) error {
	owner, repoName := *repo.Owner.Login, *repo.Name
	context.Log("Processing %s", *repo.FullName)

	// 1. Find labels on GitHub.
	labels, _, err := context.GitHub.Issues.ListLabels(owner, repoName, &listOpts)
	if err != nil {
		return context.NewError("error fetching labels for %s: %v", *repo.FullName, err)
	}

	for _, desiredLabel := range desiredLabels {
		matchedLabel := findLabel(labels, desiredLabel)

		// It doesn't exist. Create and continue.
		if matchedLabel == nil {
			if perform {
				context.Log("%s: creating %s with color %s", *repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
				_, _, err := context.GitHub.Issues.CreateLabel(owner, repoName, desiredLabel)
				if err != nil {
					return context.NewError("error creating '%s' for %s: %v", *desiredLabel.Name, *repo.FullName, err)
				}
			} else {
				context.Log("%s: would create %s with color %s", *repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
			}
			continue
		}

		// It does exist, but possibly with incorrect info. Update it.
		if *matchedLabel.Name != *desiredLabel.Name || *matchedLabel.Color != *desiredLabel.Color {
			if perform {
				context.Log("%s: updating %s with data: %v",
					*repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
				_, _, err := context.GitHub.Issues.EditLabel(owner, repoName, *matchedLabel.Name, desiredLabel)
				if err != nil {
					return context.NewError("%s: error updating '%s': %v", *repo.FullName, *matchedLabel.Name, err)
				}
			} else {
				context.Log("%s: would update %s with data: %v", *repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
			}
			continue
		}
	}

	return nil
}

func main() {
	var perform bool
	flag.BoolVar(&perform, "f", false, "Whether to modify the labels (if true) or show dry-run output (if false).")
	flag.Parse()

	context := ctx.NewDefaultContext()
	repos, _, err := context.GitHub.Repositories.List("jekyll", &github.RepositoryListOptions{
		Type: "owner", Sort: "full_name", Direction: "asc", ListOptions: listOpts,
	})
	if err != nil {
		log.Fatalln("error fetching repos:", err)
	}

	for _, repo := range repos {
		if err := processRepo(context, repo, perform); err != nil {
			context.Log("%s: failed!", *repo.FullName)
		}
	}
}
<commit_msg>unify-labels: add 'accepted' label<commit_after>// unify-labels is a CLI which will add, rename, or change the color of labels so they match the Jekyll org's requirements.
package main

import (
	"flag"
	"log"
	"strings"

	"github.com/google/go-github/github"
	"github.com/parkr/auto-reply/ctx"
	"github.com/parkr/auto-reply/freeze"
)

var desiredLabels = []*github.Label{
	{Name: github.String("accepted"), Color: github.String("4bc865")},
	{Name: github.String("bug"), Color: github.String("d41313")},
	{Name: github.String("discussion"), Color: github.String("006b75")},
	{Name: github.String("documentation"), Color: github.String("006b75")},
	{Name: github.String("enhancement"), Color: github.String("009800")},
	{Name: github.String("feature"), Color: github.String("009800")},
	{Name: github.String("fix"), Color: github.String("eb6420")},
	{Name: github.String(freeze.LabelName), Color: github.String("0052cc")},
	{Name: github.String("github"), Color: github.String("222222")},
	{Name: github.String("has-pull-request"), Color: github.String("fbca04")},
	{Name: github.String("help-wanted"), Color: github.String("fbca04")},
	{Name: github.String("internal"), Color: github.String("ededed")},
	{Name: github.String("pending-feedback"), Color: github.String("fbca04")},
	{Name: github.String("pending-rebase"), Color: github.String("eb6420")},
	{Name: github.String("pinned"), Color: github.String("f3f4d3")},
	{Name: github.String("release"), Color: github.String("d4c5f9")},
	{Name: github.String("security"), Color: github.String("e11d21")},
	{Name: github.String("stale"), Color: github.String("bfd4f2")},
	{Name: github.String("suggestion"), Color: github.String("0052cc")},
	{Name: github.String("support"), Color: github.String("5319e7")},
	{Name: github.String("tests"), Color: github.String("d4c5f9")},
	{Name: github.String("undetermined"), Color: github.String("fe3868")},
	{Name: github.String("ux"), Color: github.String("006b75")},
	{Name: github.String("windows"), Color: github.String("fbca04")},
	{Name: github.String("wont-fix"), Color: github.String("e11d21")},
}
var listOpts = github.ListOptions{PerPage: 100}

func allPossibleNames(name string) []string {
	return []string{
		name,
		strings.Replace(name, "-", "", -1),
		strings.Replace(name, "-", " ", -1),
		strings.ToLower(name),
		strings.Title(name),
		strings.Title(strings.Replace(name, "-", "", -1)),
		strings.Title(strings.Replace(name, "-", " ", -1)),
	}
}

func findLabel(labels []*github.Label, desiredLabel *github.Label) *github.Label {
	possibleNames := allPossibleNames(*desiredLabel.Name)
	for _, possibleName := range possibleNames {
		for _, label := range labels {
			if *label.Name == possibleName {
				return label
			}
		}
	}

	return nil
}

func processRepo(context *ctx.Context, repo *github.Repository, perform bool) error {
	owner, repoName := *repo.Owner.Login, *repo.Name
	context.Log("Processing %s", *repo.FullName)

	// 1. Find labels on GitHub.
	labels, _, err := context.GitHub.Issues.ListLabels(owner, repoName, &listOpts)
	if err != nil {
		return context.NewError("error fetching labels for %s: %v", *repo.FullName, err)
	}

	for _, desiredLabel := range desiredLabels {
		matchedLabel := findLabel(labels, desiredLabel)

		// It doesn't exist. Create and continue.
		if matchedLabel == nil {
			if perform {
				context.Log("%s: creating %s with color %s", *repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
				_, _, err := context.GitHub.Issues.CreateLabel(owner, repoName, desiredLabel)
				if err != nil {
					return context.NewError("error creating '%s' for %s: %v", *desiredLabel.Name, *repo.FullName, err)
				}
			} else {
				context.Log("%s: would create %s with color %s", *repo.FullName, *desiredLabel.Name, *desiredLabel.Color)
			}
			continue
		}

		// It does exist, but possibly with incorrect info. Update it.
		if *matchedLabel.Name != *desiredLabel.Name || *matchedLabel.Color != *desiredLabel.Color {
			if perform {
				context.Log("%s: updating %s with data: %v",
					*repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
				_, _, err := context.GitHub.Issues.EditLabel(owner, repoName, *matchedLabel.Name, desiredLabel)
				if err != nil {
					return context.NewError("%s: error updating '%s': %v", *repo.FullName, *matchedLabel.Name, err)
				}
			} else {
				context.Log("%s: would update %s with data: %v", *repo.FullName, *matchedLabel.Name, github.Stringify(desiredLabel))
			}
			continue
		}
	}

	return nil
}

func main() {
	var perform bool
	flag.BoolVar(&perform, "f", false, "Whether to modify the labels (if true) or show dry-run output (if false).")
	flag.Parse()

	context := ctx.NewDefaultContext()
	repos, _, err := context.GitHub.Repositories.List("jekyll", &github.RepositoryListOptions{
		Type: "owner", Sort: "full_name", Direction: "asc", ListOptions: listOpts,
	})
	if err != nil {
		log.Fatalln("error fetching repos:", err)
	}

	for _, repo := range repos {
		if err := processRepo(context, repo, perform); err != nil {
			context.Log("%s: failed!", *repo.FullName)
		}
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Knative 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 v1alpha1

import (
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"knative.dev/pkg/apis"
	duckv1beta1 "knative.dev/pkg/apis/duck/v1beta1"
	"knative.dev/pkg/kmeta"
	net "knative.dev/serving/pkg/apis/networking"
	servingv1beta1 "knative.dev/serving/pkg/apis/serving/v1beta1"
)

// +genclient
// +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object

// PodAutoscaler is a Knative abstraction that encapsulates the interface by which Knative
// components instantiate autoscalers.  This definition is an abstraction that may be backed
// by multiple definitions.  For more information, see the Knative Pluggability presentation:
// https://docs.google.com/presentation/d/10KWynvAJYuOEWy69VBa6bHJVCqIsz1TNdEKosNvcpPY/edit
type PodAutoscaler struct {
	metav1.TypeMeta `json:",inline"`
	// +optional
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec holds the desired state of the PodAutoscaler (from the client).
	// +optional
	Spec PodAutoscalerSpec `json:"spec,omitempty"`

	// Status communicates the observed state of the PodAutoscaler (from the controller).
	// +optional
	Status PodAutoscalerStatus `json:"status,omitempty"`
}

// Verify that PodAutoscaler adheres to the appropriate interfaces.
var (
	// Check that PodAutoscaler can be validated and can be defaulted.
	_ apis.Validatable = (*PodAutoscaler)(nil)
	_ apis.Defaultable = (*PodAutoscaler)(nil)

	// Check that we can create OwnerReferences to a PodAutoscaler.
	_ kmeta.OwnerRefable = (*PodAutoscaler)(nil)
)

// PodAutoscalerSpec holds the desired state of the PodAutoscaler (from the client).
type PodAutoscalerSpec struct {
	// DeprecatedGeneration was used prior in Kubernetes versions <1.11
	// when metadata.generation was not being incremented by the api server
	//
	// This property will be dropped in future Knative releases and should
	// not be used - use metadata.generation
	//
	// Tracking issue: https://knative.dev/serving/issues/643
	//
	// +optional
	DeprecatedGeneration int64 `json:"generation,omitempty"`

	// ContainerConcurrency specifies the maximum allowed
	// in-flight (concurrent) requests per container of the Revision.
	// Defaults to `0` which means unlimited concurrency.
	// +optional
	ContainerConcurrency servingv1beta1.RevisionContainerConcurrencyType `json:"containerConcurrency,omitempty"`

	// ScaleTargetRef defines the /scale-able resource that this PodAutoscaler
	// is responsible for quickly right-sizing.
	ScaleTargetRef corev1.ObjectReference `json:"scaleTargetRef"`

	// DeprecatedServiceName holds the name of a core Kubernetes Service resource that
	// load balances over the pods referenced by the ScaleTargetRef.
	DeprecatedServiceName string `json:"serviceName"`

	// The application-layer protocol. Matches `ProtocolType` inferred from the revision spec.
	ProtocolType net.ProtocolType
}

const (
	// PodAutoscalerConditionReady is set when the revision is starting to materialize
	// runtime resources, and becomes true when those resources are ready.
	PodAutoscalerConditionReady = apis.ConditionReady
	// PodAutoscalerConditionActive is set when the PodAutoscaler's ScaleTargetRef is receiving traffic.
	PodAutoscalerConditionActive apis.ConditionType = "Active"
)

// PodAutoscalerStatus communicates the observed state of the PodAutoscaler (from the controller).
type PodAutoscalerStatus struct {
	duckv1beta1.Status

	// ServiceName is the K8s Service name that serves the revision, scaled by this PA.
	// The service is created and owned by the ServerlessService object owned by this PA.
	ServiceName string `json:"serviceName"`

	// MetricsServiceName is the K8s Service name that provides revision metrics.
	// The service is managed by the PA object.
	MetricsServiceName string `json:"metricsServiceName"`
}

// +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object

// PodAutoscalerList is a list of PodAutoscaler resources
type PodAutoscalerList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata"`

	Items []PodAutoscaler `json:"items"`
}
<commit_msg>Add json struct tag to ProtocolType on PA (#5150)<commit_after>/*
Copyright 2018 The Knative 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 v1alpha1

import (
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"knative.dev/pkg/apis"
	duckv1beta1 "knative.dev/pkg/apis/duck/v1beta1"
	"knative.dev/pkg/kmeta"
	net "knative.dev/serving/pkg/apis/networking"
	servingv1beta1 "knative.dev/serving/pkg/apis/serving/v1beta1"
)

// +genclient
// +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object

// PodAutoscaler is a Knative abstraction that encapsulates the interface by which Knative
// components instantiate autoscalers.  This definition is an abstraction that may be backed
// by multiple definitions.  For more information, see the Knative Pluggability presentation:
// https://docs.google.com/presentation/d/10KWynvAJYuOEWy69VBa6bHJVCqIsz1TNdEKosNvcpPY/edit
type PodAutoscaler struct {
	metav1.TypeMeta `json:",inline"`
	// +optional
	metav1.ObjectMeta `json:"metadata,omitempty"`

	// Spec holds the desired state of the PodAutoscaler (from the client).
	// +optional
	Spec PodAutoscalerSpec `json:"spec,omitempty"`

	// Status communicates the observed state of the PodAutoscaler (from the controller).
	// +optional
	Status PodAutoscalerStatus `json:"status,omitempty"`
}

// Verify that PodAutoscaler adheres to the appropriate interfaces.
var (
	// Check that PodAutoscaler can be validated and can be defaulted.
	_ apis.Validatable = (*PodAutoscaler)(nil)
	_ apis.Defaultable = (*PodAutoscaler)(nil)

	// Check that we can create OwnerReferences to a PodAutoscaler.
	_ kmeta.OwnerRefable = (*PodAutoscaler)(nil)
)

// PodAutoscalerSpec holds the desired state of the PodAutoscaler (from the client).
type PodAutoscalerSpec struct {
	// DeprecatedGeneration was used prior in Kubernetes versions <1.11
	// when metadata.generation was not being incremented by the api server
	//
	// This property will be dropped in future Knative releases and should
	// not be used - use metadata.generation
	//
	// Tracking issue: https://knative.dev/serving/issues/643
	//
	// +optional
	DeprecatedGeneration int64 `json:"generation,omitempty"`

	// ContainerConcurrency specifies the maximum allowed
	// in-flight (concurrent) requests per container of the Revision.
	// Defaults to `0` which means unlimited concurrency.
	// +optional
	ContainerConcurrency servingv1beta1.RevisionContainerConcurrencyType `json:"containerConcurrency,omitempty"`

	// ScaleTargetRef defines the /scale-able resource that this PodAutoscaler
	// is responsible for quickly right-sizing.
	ScaleTargetRef corev1.ObjectReference `json:"scaleTargetRef"`

	// DeprecatedServiceName holds the name of a core Kubernetes Service resource that
	// load balances over the pods referenced by the ScaleTargetRef.
	DeprecatedServiceName string `json:"serviceName"`

	// The application-layer protocol. Matches `ProtocolType` inferred from the revision spec.
	ProtocolType net.ProtocolType `json:"protocolType"`
}

const (
	// PodAutoscalerConditionReady is set when the revision is starting to materialize
	// runtime resources, and becomes true when those resources are ready.
	PodAutoscalerConditionReady = apis.ConditionReady
	// PodAutoscalerConditionActive is set when the PodAutoscaler's ScaleTargetRef is receiving traffic.
	PodAutoscalerConditionActive apis.ConditionType = "Active"
)

// PodAutoscalerStatus communicates the observed state of the PodAutoscaler (from the controller).
type PodAutoscalerStatus struct {
	duckv1beta1.Status

	// ServiceName is the K8s Service name that serves the revision, scaled by this PA.
	// The service is created and owned by the ServerlessService object owned by this PA.
	ServiceName string `json:"serviceName"`

	// MetricsServiceName is the K8s Service name that provides revision metrics.
	// The service is managed by the PA object.
	MetricsServiceName string `json:"metricsServiceName"`
}

// +k8s:deepcopy-gen:interfaces=k8s.io/apimachinery/pkg/runtime.Object

// PodAutoscalerList is a list of PodAutoscaler resources
type PodAutoscalerList struct {
	metav1.TypeMeta `json:",inline"`
	metav1.ListMeta `json:"metadata"`

	Items []PodAutoscaler `json:"items"`
}
<|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 whitebox

import (
	"fmt"

	"k8s.io/api/networking/v1beta1"
	"k8s.io/ingress-gce/pkg/fuzz"
	"k8s.io/ingress-gce/pkg/utils"
)

// Implements a whitebox test to check that the GCLB has the expected number of BackendService's.
type numBackendServicesTest struct {
	uniqSvcPorts map[utils.ServicePortID]bool
}

// Name implements WhiteboxTest.
func (t *numBackendServicesTest) Name() string {
	return "NumBackendServicesTest"

}

// Test implements WhiteboxTest.
func (t *numBackendServicesTest) Test(ing *v1beta1.Ingress, gclb *fuzz.GCLB) error {
	t.uniqSvcPorts = make(map[utils.ServicePortID]bool)
	expectedBackendServices := 0

	if ing.Spec.Backend == nil {
		expectedBackendServices++
	}

	utils.TraverseIngressBackends(ing, func(id utils.ServicePortID) bool {
		if _, ok := t.uniqSvcPorts[id]; !ok {
			expectedBackendServices++
			t.uniqSvcPorts[id] = true
		}
		return false
	})

	if len(gclb.BackendService) != expectedBackendServices {
		return fmt.Errorf("Expected %d BackendService's but got %d", expectedBackendServices, len(gclb.BackendService))
	}

	// Verify that access logs are enabled for GA version.
	for _, cbe := range gclb.BackendService {
		if cbe.GA != nil {
			continue
		}
		if !cbe.GA.LogConfig.Enable {
			return fmt.Errorf("access logs are disabled, expected to be enabled")
		}
	}

	return nil
}
<commit_msg>Fix Access logs enable whitebox test<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 whitebox

import (
	"fmt"

	"k8s.io/api/networking/v1beta1"
	"k8s.io/ingress-gce/pkg/fuzz"
	"k8s.io/ingress-gce/pkg/utils"
)

// Implements a whitebox test to check that the GCLB has the expected number of BackendService's.
type numBackendServicesTest struct {
	uniqSvcPorts map[utils.ServicePortID]bool
}

// Name implements WhiteboxTest.
func (t *numBackendServicesTest) Name() string {
	return "NumBackendServicesTest"

}

// Test implements WhiteboxTest.
func (t *numBackendServicesTest) Test(ing *v1beta1.Ingress, gclb *fuzz.GCLB) error {
	t.uniqSvcPorts = make(map[utils.ServicePortID]bool)
	expectedBackendServices := 0

	if ing.Spec.Backend == nil {
		expectedBackendServices++
	}

	utils.TraverseIngressBackends(ing, func(id utils.ServicePortID) bool {
		if _, ok := t.uniqSvcPorts[id]; !ok {
			expectedBackendServices++
			t.uniqSvcPorts[id] = true
		}
		return false
	})

	if len(gclb.BackendService) != expectedBackendServices {
		return fmt.Errorf("Expected %d BackendService's but got %d", expectedBackendServices, len(gclb.BackendService))
	}

	// Verify that access logs are enabled for GA version.
	for _, cbe := range gclb.BackendService {
		if cbe.GA == nil {
			continue
		}
		if !cbe.GA.LogConfig.Enable {
			return fmt.Errorf("access logs are disabled, expected to be enabled")
		}
	}

	return nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 The Kubernetes Authors 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 whiteblacklist

import (
	"regexp"
	"testing"
)

func TestNew(t *testing.T) {
	t.Run("fails with two non empty maps", func(t *testing.T) {
		_, err := New(map[string]struct{}{"not-empty": {}}, map[string]struct{}{"not-empty": {}})
		if err == nil {
			t.Fatal("expected New() to fail with two non-empty maps")
		}
	})

	t.Run("defaults to blacklisting", func(t *testing.T) {
		l, err := New(map[string]struct{}{}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		if l.isWhiteList {
			t.Fatal("expected whiteBlackList to default to blacklist")
		}
	})

	t.Run("if whitelist set, should be whitelist", func(t *testing.T) {
		list, err := New(map[string]struct{}{"not-empty": {}}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		if !list.isWhiteList {
			t.Fatal("expected list to be whitelist")
		}
	})

	t.Run("if blacklist set, should be blacklist", func(t *testing.T) {
		list, err := New(map[string]struct{}{}, map[string]struct{}{"not-empty": {}})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		if list.isWhiteList {
			t.Fatal("expected list to be blacklist")
		}
	})
}

func TestInclude(t *testing.T) {
	t.Run("adds when whitelist", func(t *testing.T) {
		whitelist, err := New(map[string]struct{}{"not-empty": {}}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		whitelist.Include([]string{"item1"})

		if !whitelist.IsIncluded("item1") {
			t.Fatal("expected included item to be included")
		}
	})
	t.Run("removes when blacklist", func(t *testing.T) {
		item1 := "item1"
		blacklist, err := New(map[string]struct{}{}, map[string]struct{}{item1: {}})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		blacklist.Include([]string{item1})

		if !blacklist.IsIncluded(item1) {
			t.Fatal("expected included item to be included")
		}
	})
}

func TestExclude(t *testing.T) {
	t.Run("removes when whitelist", func(t *testing.T) {
		item1 := "item1"
		whitelist, err := New(map[string]struct{}{item1: {}}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		whitelist.Exclude([]string{item1})

		if whitelist.IsIncluded(item1) {
			t.Fatal("expected excluded item to be excluded")
		}
	})
	t.Run("removes when blacklist", func(t *testing.T) {
		item1 := "item1"
		blacklist, err := New(map[string]struct{}{}, map[string]struct{}{"not-empty": {}})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		blacklist.Exclude([]string{item1})

		if blacklist.IsIncluded(item1) {
			t.Fatal("expected excluded item to be excluded")
		}
	})
}

func TestStatus(t *testing.T) {
	t.Run("status when whitelist has single item", func(t *testing.T) {
		item1 := "item1"
		whitelist, _ := New(map[string]struct{}{item1: {}}, map[string]struct{}{})
		actualStatusString := whitelist.Status()
		expectedStatusString := "whitelisting the following items: " + item1
		if actualStatusString != expectedStatusString {
			t.Errorf("expected status %q but got %q", expectedStatusString, actualStatusString)
		}
	})
	t.Run("status when whitelist has multiple items", func(t *testing.T) {
		item1 := "item1"
		item2 := "item2"
		whitelist, _ := New(map[string]struct{}{item1: {}, item2: {}}, map[string]struct{}{})
		actualStatusString := whitelist.Status()
		expectedRegexPattern := `^whitelisting the following items: item1|item2, item2|item1$`
		matched, _ := regexp.MatchString(expectedRegexPattern, actualStatusString)
		if !matched {
			t.Errorf("expected status %q but got %q", expectedRegexPattern, actualStatusString)
		}
	})
	t.Run("status when blacklist has single item", func(t *testing.T) {
		item1 := "not-empty"
		blacklist, _ := New(map[string]struct{}{}, map[string]struct{}{item1: {}})
		actualStatusString := blacklist.Status()
		expectedStatusString := "blacklisting the following items: " + item1
		if actualStatusString != expectedStatusString {
			t.Errorf("expected status %q but got %q", expectedStatusString, actualStatusString)
		}
	})
	t.Run("status when blacklist has multiple items", func(t *testing.T) {
		item1 := "item1"
		item2 := "item2"
		blacklist, _ := New(map[string]struct{}{}, map[string]struct{}{item1: {}, item2: {}})
		actualStatusString := blacklist.Status()
		expectedRegexPattern := `^blacklisting the following items: item1|item2, item2|item1$`
		matched, _ := regexp.MatchString(expectedRegexPattern, actualStatusString)
		if !matched {
			t.Errorf("expected status %q but got %q", expectedRegexPattern, actualStatusString)
		}
	})
}
<commit_msg>Adding parenthesis in the regex<commit_after>/*
Copyright 2018 The Kubernetes Authors 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 whiteblacklist

import (
	"regexp"
	"testing"
)

func TestNew(t *testing.T) {
	t.Run("fails with two non empty maps", func(t *testing.T) {
		_, err := New(map[string]struct{}{"not-empty": {}}, map[string]struct{}{"not-empty": {}})
		if err == nil {
			t.Fatal("expected New() to fail with two non-empty maps")
		}
	})

	t.Run("defaults to blacklisting", func(t *testing.T) {
		l, err := New(map[string]struct{}{}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		if l.isWhiteList {
			t.Fatal("expected whiteBlackList to default to blacklist")
		}
	})

	t.Run("if whitelist set, should be whitelist", func(t *testing.T) {
		list, err := New(map[string]struct{}{"not-empty": {}}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		if !list.isWhiteList {
			t.Fatal("expected list to be whitelist")
		}
	})

	t.Run("if blacklist set, should be blacklist", func(t *testing.T) {
		list, err := New(map[string]struct{}{}, map[string]struct{}{"not-empty": {}})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		if list.isWhiteList {
			t.Fatal("expected list to be blacklist")
		}
	})
}

func TestInclude(t *testing.T) {
	t.Run("adds when whitelist", func(t *testing.T) {
		whitelist, err := New(map[string]struct{}{"not-empty": {}}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		whitelist.Include([]string{"item1"})

		if !whitelist.IsIncluded("item1") {
			t.Fatal("expected included item to be included")
		}
	})
	t.Run("removes when blacklist", func(t *testing.T) {
		item1 := "item1"
		blacklist, err := New(map[string]struct{}{}, map[string]struct{}{item1: {}})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		blacklist.Include([]string{item1})

		if !blacklist.IsIncluded(item1) {
			t.Fatal("expected included item to be included")
		}
	})
}

func TestExclude(t *testing.T) {
	t.Run("removes when whitelist", func(t *testing.T) {
		item1 := "item1"
		whitelist, err := New(map[string]struct{}{item1: {}}, map[string]struct{}{})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		whitelist.Exclude([]string{item1})

		if whitelist.IsIncluded(item1) {
			t.Fatal("expected excluded item to be excluded")
		}
	})
	t.Run("removes when blacklist", func(t *testing.T) {
		item1 := "item1"
		blacklist, err := New(map[string]struct{}{}, map[string]struct{}{"not-empty": {}})
		if err != nil {
			t.Fatal("expected New() to not fail")
		}

		blacklist.Exclude([]string{item1})

		if blacklist.IsIncluded(item1) {
			t.Fatal("expected excluded item to be excluded")
		}
	})
}

func TestStatus(t *testing.T) {
	t.Run("status when whitelist has single item", func(t *testing.T) {
		item1 := "item1"
		whitelist, _ := New(map[string]struct{}{item1: {}}, map[string]struct{}{})
		actualStatusString := whitelist.Status()
		expectedStatusString := "whitelisting the following items: " + item1
		if actualStatusString != expectedStatusString {
			t.Errorf("expected status %q but got %q", expectedStatusString, actualStatusString)
		}
	})
	t.Run("status when whitelist has multiple items", func(t *testing.T) {
		item1 := "item1"
		item2 := "item2"
		whitelist, _ := New(map[string]struct{}{item1: {}, item2: {}}, map[string]struct{}{})
		actualStatusString := whitelist.Status()
		expectedRegexPattern := `^whitelisting the following items: (item1|item2), (item2|item1)$`
		matched, _ := regexp.MatchString(expectedRegexPattern, actualStatusString)
		if !matched {
			t.Errorf("expected status %q but got %q", expectedRegexPattern, actualStatusString)
		}
	})
	t.Run("status when blacklist has single item", func(t *testing.T) {
		item1 := "not-empty"
		blacklist, _ := New(map[string]struct{}{}, map[string]struct{}{item1: {}})
		actualStatusString := blacklist.Status()
		expectedStatusString := "blacklisting the following items: " + item1
		if actualStatusString != expectedStatusString {
			t.Errorf("expected status %q but got %q", expectedStatusString, actualStatusString)
		}
	})
	t.Run("status when blacklist has multiple items", func(t *testing.T) {
		item1 := "item1"
		item2 := "item2"
		blacklist, _ := New(map[string]struct{}{}, map[string]struct{}{item1: {}, item2: {}})
		actualStatusString := blacklist.Status()
		expectedRegexPattern := `^blacklisting the following items: (item1|item2), (item2|item1)$`
		matched, _ := regexp.MatchString(expectedRegexPattern, actualStatusString)
		if !matched {
			t.Errorf("expected status %q but got %q", expectedRegexPattern, actualStatusString)
		}
	})
}
<|endoftext|>