<commit_before>package kafka

import (
        "github.com/Shopify/sarama"
        "log"
)

const (
  bufferSize     = 256
  initialOffset  = sarama.OffsetOldest // always start listening for the latest event.
)

type Consumer struct {
  consumer           sarama.Consumer
  partitionConsumers []sarama.PartitionConsumer
  messages           chan *sarama.ConsumerMessage
}

func NewConsumer(brokers []string, topic string) *Consumer {
  config := sarama.NewConfig()
  config.Consumer.Return.Errors = true

  consumer, err := sarama.NewConsumer(brokers, config)
  if err != nil {
    log.Fatalln(err)
  }

  partitions, err := consumer.Partitions(topic)
  if err != nil {
    log.Println("Failed to get the list of partitions: %s", err)
  }

  partitionConsumers := make([]sarama.PartitionConsumer, len(partitions))
  messages := make(chan *sarama.ConsumerMessage, bufferSize)

  for _, partition := range partitions {

    partitionConsumer, err := consumer.ConsumePartition(topic, partition, initialOffset)

    if err != nil {
      log.Fatalln("Failed to start consumer for partition %d: %s", partition, err)
    }

    go func(partitionConsumer sarama.PartitionConsumer) {
      for message := range partitionConsumer.Messages() {
        messages <- message
      }
    }(partitionConsumer)

  }

  return &Consumer{
    consumer           : consumer,
    partitionConsumers : partitionConsumers,
    messages           : messages,
  }
}

// Consume messages and process them through the method pass in parameter
func (this *Consumer) Consume(processMessage func(key, value []byte)) {
go func() {
    for message := range this.messages {
      processMessage(message.Key, message.Value)
    }
  }()
}

// Close stops processing messages and releases the corresponding resources
func (this *Consumer) Close() {

  log.Println("Done consuming messages")

  for _, partitionConsumer := range this.partitionConsumers {
    if err := partitionConsumer.Close(); err != nil {
      log.Println("Failed to close partition consumer: ", err)
    }
  }

  if err := this.consumer.Close(); err != nil {
    log.Println("Failed to shutdown kafka consumer cleanly: %s", err)
  }

  close(this.messages)

}<commit_msg>Add some logs<commit_after>package kafka

import (
        "github.com/Shopify/sarama"
        "log"
)

const (
  bufferSize     = 256
  initialOffset  = sarama.OffsetOldest // always start listening for the latest event.
)

type Consumer struct {
  consumer           sarama.Consumer
  partitionConsumers []sarama.PartitionConsumer
  messages           chan *sarama.ConsumerMessage
}

func NewConsumer(brokers []string, topic string) *Consumer {
  config := sarama.NewConfig()
  config.Consumer.Return.Errors = true

  consumer, err := sarama.NewConsumer(brokers, config)
  if err != nil {
    log.Fatalln(err)
  }

  partitions, err := consumer.Partitions(topic)
  if err != nil {
    log.Println("Failed to get the list of partitions: %s", err)
  }

  log.Println("%s partitions found", len(partitions))

  partitionConsumers := make([]sarama.PartitionConsumer, len(partitions))
  messages := make(chan *sarama.ConsumerMessage, bufferSize)

  for _, partition := range partitions {

    partitionConsumer, err := consumer.ConsumePartition(topic, partition, initialOffset)

    if err != nil {
      log.Fatalln("Failed to start consumer for partition %d: %s", partition, err)
    }

    go func(partitionConsumer sarama.PartitionConsumer) {
      for message := range partitionConsumer.Messages() {
        messages <- message
      }
    }(partitionConsumer)

  }

  return &Consumer{
    consumer           : consumer,
    partitionConsumers : partitionConsumers,
    messages           : messages,
  }
}

// Consume messages and process them through the method pass in parameter
func (this *Consumer) Consume(processMessage func(key, value []byte)) {
  go func() {
      log.Println("Start consuming messages ...")

      for message := range this.messages {
        log.Println("Process message with offset %s", message.Offset)
        processMessage(message.Key, message.Value)
      }
    }()
}

// Close stops processing messages and releases the corresponding resources
func (this *Consumer) Close() {

  log.Println("Done consuming messages")

  for _, partitionConsumer := range this.partitionConsumers {
    if err := partitionConsumer.Close(); err != nil {
      log.Println("Failed to close partition consumer: ", err)
    }
  }

  if err := this.consumer.Close(); err != nil {
    log.Println("Failed to shutdown kafka consumer cleanly: %s", err)
  }

  close(this.messages)

}<|endoftext|>


<commit_before>package ethutil

import (
	"bytes"
	"math/big"
	"reflect"
	"testing"
)

func TestRlpValueEncoding(t *testing.T) {
	val := EmptyValue()
	val.AppendList().Append(1).Append(2).Append(3)
	val.Append("4").AppendList().Append(5)

	res := val.Encode()
	exp := Encode([]interface{}{[]interface{}{1, 2, 3}, "4", []interface{}{5}})
	if bytes.Compare(res, exp) != 0 {
		t.Errorf("expected %q, got %q", res, exp)
	}
}

func TestValueSlice(t *testing.T) {
	val := []interface{}{
		"value1",
		"valeu2",
		"value3",
	}

	value := NewValue(val)
	splitVal := value.SliceFrom(1)

	if splitVal.Len() != 2 {
		t.Error("SliceFrom: Expected len", 2, "got", splitVal.Len())
	}

	splitVal = value.SliceTo(2)
	if splitVal.Len() != 2 {
		t.Error("SliceTo: Expected len", 2, "got", splitVal.Len())
	}

	splitVal = value.SliceFromTo(1, 3)
	if splitVal.Len() != 2 {
		t.Error("SliceFromTo: Expected len", 2, "got", splitVal.Len())
	}
}

func TestValue(t *testing.T) {
	value := NewValueFromBytes([]byte("\xcd\x83dog\x83god\x83cat\x01"))
	if value.Get(0).Str() != "dog" {
		t.Errorf("expected '%v', got '%v'", value.Get(0).Str(), "dog")
	}

	if value.Get(3).Uint() != 1 {
		t.Errorf("expected '%v', got '%v'", value.Get(3).Uint(), 1)
	}
}

func TestEncode(t *testing.T) {
	strRes := "\x83dog"
	bytes := Encode("dog")

	str := string(bytes)
	if str != strRes {
		t.Errorf("Expected %q, got %q", strRes, str)
	}

	sliceRes := "\xcc\x83dog\x83god\x83cat"
	strs := []interface{}{"dog", "god", "cat"}
	bytes = Encode(strs)
	slice := string(bytes)
	if slice != sliceRes {
		t.Error("Expected %q, got %q", sliceRes, slice)
	}

	intRes := "\x82\x04\x00"
	bytes = Encode(1024)
	if string(bytes) != intRes {
		t.Errorf("Expected %q, got %q", intRes, bytes)
	}
}

func TestDecode(t *testing.T) {
	single := []byte("\x01")
	b, _ := Decode(single, 0)

	if b.(uint8) != 1 {
		t.Errorf("Expected 1, got %q", b)
	}

	str := []byte("\x83dog")
	b, _ = Decode(str, 0)
	if bytes.Compare(b.([]byte), []byte("dog")) != 0 {
		t.Errorf("Expected dog, got %q", b)
	}

	slice := []byte("\xcc\x83dog\x83god\x83cat")
	res := []interface{}{"dog", "god", "cat"}
	b, _ = Decode(slice, 0)
	if reflect.DeepEqual(b, res) {
		t.Errorf("Expected %q, got %q", res, b)
	}
}

func TestEncodeDecodeBigInt(t *testing.T) {
	bigInt := big.NewInt(1391787038)
	encoded := Encode(bigInt)

	value := NewValueFromBytes(encoded)
	if value.BigInt().Cmp(bigInt) != 0 {
		t.Errorf("Expected %v, got %v", bigInt, value.BigInt())
	}
}

func TestEncodeDecodeBytes(t *testing.T) {
	b := NewValue([]interface{}{[]byte{1, 2, 3, 4, 5}, byte(6)})
	val := NewValueFromBytes(b.Encode())
	if !b.Cmp(val) {
		t.Errorf("Expected %v, got %v", val, b)
	}
}

/*
var ZeroHash256 = make([]byte, 32)
var ZeroHash160 = make([]byte, 20)
var EmptyShaList = Sha3Bin(Encode([]interface{}{}))

var GenisisHeader = []interface{}{
	// Previous hash (none)
	//"",
	ZeroHash256,
	// Sha of uncles
	Sha3Bin(Encode([]interface{}{})),
	// Coinbase
	ZeroHash160,
	// Root state
	"",
	// Sha of transactions
	//EmptyShaList,
	Sha3Bin(Encode([]interface{}{})),
	// Difficulty
	BigPow(2, 22),
	// Time
	//big.NewInt(0),
	int64(0),
	// extra
	"",
	// Nonce
	big.NewInt(42),
}

func TestEnc(t *testing.T) {
	//enc := Encode(GenisisHeader)
	//fmt.Printf("%x (%d)\n", enc, len(enc))
	h, _ := hex.DecodeString("f8a0a00000000000000000000000000000000000000000000000000000000000000000a01dcc4de8dec75d7aab85b567b6ccd41ad312451b948a7413f0a142fd40d49347940000000000000000000000000000000000000000a06d076baa9c4074fb2df222dd16a96b0155a1e6686b3e5748b4e9ca0a208a425ca01dcc4de8dec75d7aab85b567b6ccd41ad312451b948a7413f0a142fd40d493478340000080802a")
	fmt.Printf("%x\n", Sha3Bin(h))
}
*/

func BenchmarkEncodeDecode(b *testing.B) {
	for i := 0; i < b.N; i++ {
		bytes := Encode([]interface{}{"dog", "god", "cat"})
		Decode(bytes, 0)
	}
}
<commit_msg>Added a few tests<commit_after>package ethutil

import (
	"bytes"
	"math/big"
	"reflect"
	"testing"
)

func TestRlpValueEncoding(t *testing.T) {
	val := EmptyValue()
	val.AppendList().Append(1).Append(2).Append(3)
	val.Append("4").AppendList().Append(5)

	res := val.Encode()
	exp := Encode([]interface{}{[]interface{}{1, 2, 3}, "4", []interface{}{5}})
	if bytes.Compare(res, exp) != 0 {
		t.Errorf("expected %q, got %q", res, exp)
	}
}

func TestValueSlice(t *testing.T) {
	val := []interface{}{
		"value1",
		"valeu2",
		"value3",
	}

	value := NewValue(val)
	splitVal := value.SliceFrom(1)

	if splitVal.Len() != 2 {
		t.Error("SliceFrom: Expected len", 2, "got", splitVal.Len())
	}

	splitVal = value.SliceTo(2)
	if splitVal.Len() != 2 {
		t.Error("SliceTo: Expected len", 2, "got", splitVal.Len())
	}

	splitVal = value.SliceFromTo(1, 3)
	if splitVal.Len() != 2 {
		t.Error("SliceFromTo: Expected len", 2, "got", splitVal.Len())
	}
}

func TestValue(t *testing.T) {
	value := NewValueFromBytes([]byte("\xcd\x83dog\x83god\x83cat\x01"))
	if value.Get(0).Str() != "dog" {
		t.Errorf("expected '%v', got '%v'", value.Get(0).Str(), "dog")
	}

	if value.Get(3).Uint() != 1 {
		t.Errorf("expected '%v', got '%v'", value.Get(3).Uint(), 1)
	}
}

func TestEncode(t *testing.T) {
	strRes := "\x83dog"
	bytes := Encode("dog")

	str := string(bytes)
	if str != strRes {
		t.Errorf("Expected %q, got %q", strRes, str)
	}

	sliceRes := "\xcc\x83dog\x83god\x83cat"
	strs := []interface{}{"dog", "god", "cat"}
	bytes = Encode(strs)
	slice := string(bytes)
	if slice != sliceRes {
		t.Error("Expected %q, got %q", sliceRes, slice)
	}

	intRes := "\x82\x04\x00"
	bytes = Encode(1024)
	if string(bytes) != intRes {
		t.Errorf("Expected %q, got %q", intRes, bytes)
	}
}

func TestDecode(t *testing.T) {
	single := []byte("\x01")
	b, _ := Decode(single, 0)

	if b.(uint8) != 1 {
		t.Errorf("Expected 1, got %q", b)
	}

	str := []byte("\x83dog")
	b, _ = Decode(str, 0)
	if bytes.Compare(b.([]byte), []byte("dog")) != 0 {
		t.Errorf("Expected dog, got %q", b)
	}

	slice := []byte("\xcc\x83dog\x83god\x83cat")
	res := []interface{}{"dog", "god", "cat"}
	b, _ = Decode(slice, 0)
	if reflect.DeepEqual(b, res) {
		t.Errorf("Expected %q, got %q", res, b)
	}
}

func TestEncodeDecodeBigInt(t *testing.T) {
	bigInt := big.NewInt(1391787038)
	encoded := Encode(bigInt)

	value := NewValueFromBytes(encoded)
	if value.BigInt().Cmp(bigInt) != 0 {
		t.Errorf("Expected %v, got %v", bigInt, value.BigInt())
	}
}

func TestEncodeDecodeBytes(t *testing.T) {
	b := NewValue([]interface{}{[]byte{1, 2, 3, 4, 5}, byte(6)})
	val := NewValueFromBytes(b.Encode())
	if !b.Cmp(val) {
		t.Errorf("Expected %v, got %v", val, b)
	}
}

func BenchmarkEncodeDecode(b *testing.B) {
	for i := 0; i < b.N; i++ {
		bytes := Encode([]interface{}{"dog", "god", "cat"})
		Decode(bytes, 0)
	}
}
<|endoftext|>


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

package event_test

import (
	"context"
	"io"
	"testing"

	"golang.org/x/exp/event"
	"golang.org/x/exp/event/adapter/logfmt"
	"golang.org/x/exp/event/eventtest"
	"golang.org/x/exp/event/keys"
	"golang.org/x/exp/event/severity"
)

var (
	aValue  = keys.Int(eventtest.A.Name)
	bValue  = keys.String(eventtest.B.Name)
	aCount  = keys.Int64("aCount")
	aStat   = keys.Int("aValue")
	bCount  = keys.Int64("B")
	bLength = keys.Int("BLen")

	eventLog = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			event.To(ctx).With(severity.Info).With(aValue.Of(a)).Log(eventtest.A.Msg)
			return ctx
		},
		AEnd: func(ctx context.Context) {},
		BStart: func(ctx context.Context, b string) context.Context {
			event.To(ctx).With(severity.Info).With(bValue.Of(b)).Log(eventtest.B.Msg)
			return ctx
		},
		BEnd: func(ctx context.Context) {},
	}

	eventLogf = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			event.To(ctx).With(severity.Info).Logf(eventtest.A.Msgf, a)
			return ctx
		},
		AEnd: func(ctx context.Context) {},
		BStart: func(ctx context.Context, b string) context.Context {
			event.To(ctx).With(severity.Info).Logf(eventtest.B.Msgf, b)
			return ctx
		},
		BEnd: func(ctx context.Context) {},
	}

	eventTrace = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			ctx, _ = event.To(ctx).Start(eventtest.A.Msg)
			event.To(ctx).With(aValue.Of(a)).Annotate()
			return ctx
		},
		AEnd: func(ctx context.Context) {
			event.To(ctx).End()
		},
		BStart: func(ctx context.Context, b string) context.Context {
			ctx, _ = event.To(ctx).Start(eventtest.B.Msg)
			event.To(ctx).With(bValue.Of(b)).Annotate()
			return ctx
		},
		BEnd: func(ctx context.Context) {
			event.To(ctx).End()
		},
	}

	eventMetric = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			event.To(ctx).With(aStat.Of(a)).Metric(gauge.Record(1))
			event.To(ctx).With(aCount.Of(1)).Metric(gauge.Record(1))
			return ctx
		},
		AEnd: func(ctx context.Context) {},
		BStart: func(ctx context.Context, b string) context.Context {
			event.To(ctx).With(bLength.Of(len(b))).Metric(gauge.Record(1))
			event.To(ctx).With(bCount.Of(1)).Metric(gauge.Record(1))
			return ctx
		},
		BEnd: func(ctx context.Context) {},
	}
)

func eventNoExporter() context.Context {
	return event.WithExporter(context.Background(), nil)
}

func eventNoop() context.Context {
	return event.WithExporter(context.Background(), event.NewExporter(event.NopHandler{}, eventtest.ExporterOptions()))
}

func eventPrint(w io.Writer) context.Context {
	return event.WithExporter(context.Background(), event.NewExporter(logfmt.NewHandler(w), eventtest.ExporterOptions()))
}

func BenchmarkEventLogNoExporter(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoExporter(), eventLog)
}

func BenchmarkEventLogNoop(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoop(), eventLog)
}

func BenchmarkEventLogDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventLog)
}

func BenchmarkEventLogfDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventLogf)
}

func BenchmarkEventTraceNoop(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoop(), eventTrace)
}

func BenchmarkEventTraceDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventTrace)
}

func BenchmarkEventMetricNoop(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoop(), eventMetric)
}

func BenchmarkEventMetricDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventMetric)
}
<commit_msg>event: add a benchmark with namespace capturing turned on<commit_after>// Copyright 2020 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package event_test

import (
	"context"
	"io"
	"testing"

	"golang.org/x/exp/event"
	"golang.org/x/exp/event/adapter/logfmt"
	"golang.org/x/exp/event/eventtest"
	"golang.org/x/exp/event/keys"
	"golang.org/x/exp/event/severity"
)

var (
	aValue  = keys.Int(eventtest.A.Name)
	bValue  = keys.String(eventtest.B.Name)
	aCount  = keys.Int64("aCount")
	aStat   = keys.Int("aValue")
	bCount  = keys.Int64("B")
	bLength = keys.Int("BLen")

	eventLog = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			event.To(ctx).With(severity.Info).With(aValue.Of(a)).Log(eventtest.A.Msg)
			return ctx
		},
		AEnd: func(ctx context.Context) {},
		BStart: func(ctx context.Context, b string) context.Context {
			event.To(ctx).With(severity.Info).With(bValue.Of(b)).Log(eventtest.B.Msg)
			return ctx
		},
		BEnd: func(ctx context.Context) {},
	}

	eventLogf = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			event.To(ctx).With(severity.Info).Logf(eventtest.A.Msgf, a)
			return ctx
		},
		AEnd: func(ctx context.Context) {},
		BStart: func(ctx context.Context, b string) context.Context {
			event.To(ctx).With(severity.Info).Logf(eventtest.B.Msgf, b)
			return ctx
		},
		BEnd: func(ctx context.Context) {},
	}

	eventTrace = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			ctx, _ = event.To(ctx).Start(eventtest.A.Msg)
			event.To(ctx).With(aValue.Of(a)).Annotate()
			return ctx
		},
		AEnd: func(ctx context.Context) {
			event.To(ctx).End()
		},
		BStart: func(ctx context.Context, b string) context.Context {
			ctx, _ = event.To(ctx).Start(eventtest.B.Msg)
			event.To(ctx).With(bValue.Of(b)).Annotate()
			return ctx
		},
		BEnd: func(ctx context.Context) {
			event.To(ctx).End()
		},
	}

	eventMetric = eventtest.Hooks{
		AStart: func(ctx context.Context, a int) context.Context {
			event.To(ctx).With(aStat.Of(a)).Metric(gauge.Record(1))
			event.To(ctx).With(aCount.Of(1)).Metric(gauge.Record(1))
			return ctx
		},
		AEnd: func(ctx context.Context) {},
		BStart: func(ctx context.Context, b string) context.Context {
			event.To(ctx).With(bLength.Of(len(b))).Metric(gauge.Record(1))
			event.To(ctx).With(bCount.Of(1)).Metric(gauge.Record(1))
			return ctx
		},
		BEnd: func(ctx context.Context) {},
	}
)

func eventNoExporter() context.Context {
	return event.WithExporter(context.Background(), nil)
}

func eventNoop() context.Context {
	return event.WithExporter(context.Background(), event.NewExporter(event.NopHandler{}, eventtest.ExporterOptions()))
}

func eventPrint(w io.Writer) context.Context {
	return event.WithExporter(context.Background(), event.NewExporter(logfmt.NewHandler(w), eventtest.ExporterOptions()))
}

func eventPrintSource(w io.Writer) context.Context {
	opts := eventtest.ExporterOptions()
	opts.EnableNamespaces = true
	return event.WithExporter(context.Background(), event.NewExporter(logfmt.NewHandler(w), opts))
}

func BenchmarkEventLogNoExporter(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoExporter(), eventLog)
}

func BenchmarkEventLogNoop(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoop(), eventLog)
}

func BenchmarkEventLogDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventLog)
}

func BenchmarkEventLogSourceDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrintSource(io.Discard), eventLog)
}

func BenchmarkEventLogfDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventLogf)
}

func BenchmarkEventTraceNoop(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoop(), eventTrace)
}

func BenchmarkEventTraceDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventTrace)
}

func BenchmarkEventMetricNoop(b *testing.B) {
	eventtest.RunBenchmark(b, eventNoop(), eventMetric)
}

func BenchmarkEventMetricDiscard(b *testing.B) {
	eventtest.RunBenchmark(b, eventPrint(io.Discard), eventMetric)
}
<|endoftext|>


<commit_before>package conveyor

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/fsouza/go-dockerclient"
	"github.com/google/go-github/github"
	"github.com/remind101/pkg/reporter"

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

const (
	// Context is used for the commit status context.
	Context = "container/docker"

	// DefaultBuilderImage is the docker image used to build docker images.
	DefaultBuilderImage = "remind101/conveyor-builder"

	// DefaultDataVolume is the default name of a container serving as a
	// data volume for ssh keys and docker credentials. In general, you
	// shouldn't need to change this.
	DefaultDataVolume = "data"
)

type BuildOptions struct {
	// Repository is the repo to build.
	Repository string
	// Sha is the git commit to build.
	Sha string
	// Branch is the name of the branch that this build relates to.
	Branch string
	// Set to true to disable the layer cache. The zero value is to enable
	// caching.
	NoCache bool
}

// Builder represents something that can build a Docker image.
type Builder interface {
	// Build should build the docker image, tag it and push it to the docker
	// registry. This should return the sha256 digest of the image.
	Build(context.Context, Logger, BuildOptions) (string, error)
}

// BuilderFunc is a function that implements the Builder interface.
type BuilderFunc func(context.Context, Logger, BuildOptions) (string, error)

func (fn BuilderFunc) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	return fn(ctx, w, opts)
}

// Conveyor serves as a builder.
type Conveyor struct {
	Builder

	// A Reporter to use to report errors.
	Reporter reporter.Reporter
}

// New returns a new Conveyor instance.
func New(b Builder) *Conveyor {
	return &Conveyor{
		Builder: b,
	}
}

func (c *Conveyor) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	// Embed the reporter in the context.Context.
	ctx = reporter.WithReporter(ctx, c.reporter())
	reporter.AddContext(ctx, "options", opts)
	defer reporter.Monitor(ctx)

	defer func() {
		if err != nil {
			reporter.Report(ctx, err)
		}
	}()

	id, err = c.build(ctx, w, opts)
	return
}

// Build performs the build and ensures that the output stream is closed.
func (c *Conveyor) build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	defer func() {
		var closeErr error
		if w != nil {
			closeErr = w.Close()
		}
		if err == nil {
			// If there was no error from the builder, let the
			// downstream know that there was an error closing the
			// output stream.
			err = closeErr
		}
	}()

	id, err = c.Builder.Build(ctx, w, opts)
	return
}

func (c *Conveyor) reporter() reporter.Reporter {
	if c.Reporter == nil {
		return reporter.ReporterFunc(func(ctx context.Context, err error) error {
			fmt.Fprintf(os.Stderr, "reporting err: %v", err)
			return nil
		})
	}

	return c.Reporter
}

// DockerBuilder is a Builder implementation that runs the build in a docker
// container.
type DockerBuilder struct {
	// dataVolume is the name of the volume that contains ssh keys and
	// configuration data.
	DataVolume string
	// Name of the image to use to build the docker image. Defaults to
	// DefaultBuilderImage.
	Image string
	// Set to true to enable dry runs. This sets the `DRY` environment
	// variable within the builder container to `true`. The behavior of this
	// flag depends on how the builder image handles the `DRY` environment
	// variable.
	DryRun bool

	client *docker.Client
}

// NewDockerBuilder returns a new DockerBuilder backed by the docker client.
func NewDockerBuilder(c *docker.Client) *DockerBuilder {
	return &DockerBuilder{client: c}
}

// NewDockerBuilderFromEnv returns a new DockerBuilder with a docker client
// configured from the standard Docker environment variables.
func NewDockerBuilderFromEnv() (*DockerBuilder, error) {
	c, err := docker.NewClientFromEnv()
	if err != nil {
		return nil, err
	}

	return NewDockerBuilder(c), nil
}

// Build executes the docker image.
func (b *DockerBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	env := []string{
		fmt.Sprintf("REPOSITORY=%s", opts.Repository),
		fmt.Sprintf("BRANCH=%s", opts.Branch),
		fmt.Sprintf("SHA=%s", opts.Sha),
		fmt.Sprintf("DRY=%s", b.dryRun()),
		fmt.Sprintf("CACHE=%s", b.cache(opts)),
	}

	c, err := b.client.CreateContainer(docker.CreateContainerOptions{
		Name: uuid.New(),
		Config: &docker.Config{
			Tty:          true,
			AttachStdout: true,
			AttachStderr: true,
			OpenStdin:    true,
			Image:        b.image(),
			Env:          env,
		},
	})
	if err != nil {
		return "", fmt.Errorf("create container: %v", err)
	}
	defer b.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.ID,
		RemoveVolumes: true,
		Force:         true,
	})

	if err := b.client.StartContainer(c.ID, &docker.HostConfig{
		Privileged:  true,
		VolumesFrom: []string{b.dataVolume()},
	}); err != nil {
		return "", fmt.Errorf("start container: %v", err)
	}

	if err := b.client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    c.ID,
		OutputStream: w,
		ErrorStream:  w,
		Logs:         true,
		Stream:       true,
		Stdout:       true,
		Stderr:       true,
		RawTerminal:  true,
	}); err != nil {
		return "", fmt.Errorf("attach: %v", err)
	}

	exit, err := b.client.WaitContainer(c.ID)
	if err != nil {
		return "", fmt.Errorf("wait container: %v", err)
	}

	// A non-zero exit status means the build failed.
	if exit != 0 {
		return "", fmt.Errorf("container returned a non-zero exit code: %d", exit)
	}

	// TODO: Return sha256
	return "", nil
}

func (b *DockerBuilder) dryRun() string {
	if b.DryRun {
		return "true"
	}
	return ""
}

func (b *DockerBuilder) image() string {
	if b.Image == "" {
		return DefaultBuilderImage
	}
	return b.Image
}

func (b *DockerBuilder) dataVolume() string {
	if b.DataVolume == "" {
		return DefaultDataVolume
	}
	return b.DataVolume
}

func (b *DockerBuilder) cache(opts BuildOptions) string {
	if opts.NoCache {
		return "off"
	}

	return "on"
}

// statusUpdaterBuilder is a Builder implementation that updates the commit
// status in github.
type statusUpdaterBuilder struct {
	Builder
	github GitHubClient
	since  func(time.Time) time.Duration
}

// UpdateGitHubCommitStatus wraps b to update the GitHub commit status when a
// build starts, and stops.
func UpdateGitHubCommitStatus(b Builder, g GitHubClient) *statusUpdaterBuilder {
	return &statusUpdaterBuilder{
		Builder: b,
		github:  g,
		since:   time.Since,
	}
}

func (b *statusUpdaterBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	t := time.Now()

	defer func() {
		duration := b.since(t)
		description := fmt.Sprintf("Image built in %v.", duration)
		status := "success"
		if err != nil {
			status = "failure"
			description = err.Error()
		}
		b.updateStatus(w, opts, status, description)
	}()

	if err = b.updateStatus(w, opts, "pending", "Image building."); err != nil {
		err = fmt.Errorf("status: %v", err)
		return
	}

	id, err = b.Builder.Build(ctx, w, opts)
	return
}

// updateStatus updates the given commit with a new status.
func (b *statusUpdaterBuilder) updateStatus(w Logger, opts BuildOptions, status string, description string) error {
	context := Context
	parts := strings.SplitN(opts.Repository, "/", 2)

	var desc *string
	if description != "" {
		desc = &description
	}
	var url *string
	if status == "success" || status == "failure" || status == "error" {
		url = github.String(w.URL())
	}

	_, _, err := b.github.CreateStatus(parts[0], parts[1], opts.Sha, &github.RepoStatus{
		State:       &status,
		Context:     &context,
		Description: desc,
		TargetURL:   url,
	})
	return err
}

// BuildAsync wraps a Builder to run the build in a goroutine.
func BuildAsync(b Builder) Builder {
	build := func(ctx context.Context, w Logger, opts BuildOptions) {
		if _, err := b.Build(ctx, w, opts); err != nil {
			log.Printf("build err: %v", err)
		}
	}

	return BuilderFunc(func(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
		go build(ctx, w, opts)
		return "", nil
	})
}
<commit_msg>Add container id to reported errors.<commit_after>package conveyor

import (
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"code.google.com/p/go-uuid/uuid"

	"github.com/fsouza/go-dockerclient"
	"github.com/google/go-github/github"
	"github.com/remind101/pkg/reporter"

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

const (
	// Context is used for the commit status context.
	Context = "container/docker"

	// DefaultBuilderImage is the docker image used to build docker images.
	DefaultBuilderImage = "remind101/conveyor-builder"

	// DefaultDataVolume is the default name of a container serving as a
	// data volume for ssh keys and docker credentials. In general, you
	// shouldn't need to change this.
	DefaultDataVolume = "data"
)

type BuildOptions struct {
	// Repository is the repo to build.
	Repository string
	// Sha is the git commit to build.
	Sha string
	// Branch is the name of the branch that this build relates to.
	Branch string
	// Set to true to disable the layer cache. The zero value is to enable
	// caching.
	NoCache bool
}

// Builder represents something that can build a Docker image.
type Builder interface {
	// Build should build the docker image, tag it and push it to the docker
	// registry. This should return the sha256 digest of the image.
	Build(context.Context, Logger, BuildOptions) (string, error)
}

// BuilderFunc is a function that implements the Builder interface.
type BuilderFunc func(context.Context, Logger, BuildOptions) (string, error)

func (fn BuilderFunc) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	return fn(ctx, w, opts)
}

// Conveyor serves as a builder.
type Conveyor struct {
	Builder

	// A Reporter to use to report errors.
	Reporter reporter.Reporter
}

// New returns a new Conveyor instance.
func New(b Builder) *Conveyor {
	return &Conveyor{
		Builder: b,
	}
}

func (c *Conveyor) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	// Embed the reporter in the context.Context.
	ctx = reporter.WithReporter(ctx, c.reporter())
	reporter.AddContext(ctx, "options", opts)
	defer reporter.Monitor(ctx)

	defer func() {
		if err != nil {
			reporter.Report(ctx, err)
		}
	}()

	id, err = c.build(ctx, w, opts)
	return
}

// Build performs the build and ensures that the output stream is closed.
func (c *Conveyor) build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	defer func() {
		var closeErr error
		if w != nil {
			closeErr = w.Close()
		}
		if err == nil {
			// If there was no error from the builder, let the
			// downstream know that there was an error closing the
			// output stream.
			err = closeErr
		}
	}()

	id, err = c.Builder.Build(ctx, w, opts)
	return
}

func (c *Conveyor) reporter() reporter.Reporter {
	if c.Reporter == nil {
		return reporter.ReporterFunc(func(ctx context.Context, err error) error {
			fmt.Fprintf(os.Stderr, "reporting err: %v", err)
			return nil
		})
	}

	return c.Reporter
}

// DockerBuilder is a Builder implementation that runs the build in a docker
// container.
type DockerBuilder struct {
	// dataVolume is the name of the volume that contains ssh keys and
	// configuration data.
	DataVolume string
	// Name of the image to use to build the docker image. Defaults to
	// DefaultBuilderImage.
	Image string
	// Set to true to enable dry runs. This sets the `DRY` environment
	// variable within the builder container to `true`. The behavior of this
	// flag depends on how the builder image handles the `DRY` environment
	// variable.
	DryRun bool

	client *docker.Client
}

// NewDockerBuilder returns a new DockerBuilder backed by the docker client.
func NewDockerBuilder(c *docker.Client) *DockerBuilder {
	return &DockerBuilder{client: c}
}

// NewDockerBuilderFromEnv returns a new DockerBuilder with a docker client
// configured from the standard Docker environment variables.
func NewDockerBuilderFromEnv() (*DockerBuilder, error) {
	c, err := docker.NewClientFromEnv()
	if err != nil {
		return nil, err
	}

	return NewDockerBuilder(c), nil
}

// Build executes the docker image.
func (b *DockerBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
	env := []string{
		fmt.Sprintf("REPOSITORY=%s", opts.Repository),
		fmt.Sprintf("BRANCH=%s", opts.Branch),
		fmt.Sprintf("SHA=%s", opts.Sha),
		fmt.Sprintf("DRY=%s", b.dryRun()),
		fmt.Sprintf("CACHE=%s", b.cache(opts)),
	}

	c, err := b.client.CreateContainer(docker.CreateContainerOptions{
		Name: uuid.New(),
		Config: &docker.Config{
			Tty:          true,
			AttachStdout: true,
			AttachStderr: true,
			OpenStdin:    true,
			Image:        b.image(),
			Env:          env,
		},
	})
	if err != nil {
		return "", fmt.Errorf("create container: %v", err)
	}
	defer b.client.RemoveContainer(docker.RemoveContainerOptions{
		ID:            c.ID,
		RemoveVolumes: true,
		Force:         true,
	})

	reporter.AddContext(ctx, "container_id", c.ID)

	if err := b.client.StartContainer(c.ID, &docker.HostConfig{
		Privileged:  true,
		VolumesFrom: []string{b.dataVolume()},
	}); err != nil {
		return "", fmt.Errorf("start container: %v", err)
	}

	if err := b.client.AttachToContainer(docker.AttachToContainerOptions{
		Container:    c.ID,
		OutputStream: w,
		ErrorStream:  w,
		Logs:         true,
		Stream:       true,
		Stdout:       true,
		Stderr:       true,
		RawTerminal:  true,
	}); err != nil {
		return "", fmt.Errorf("attach: %v", err)
	}

	exit, err := b.client.WaitContainer(c.ID)
	if err != nil {
		return "", fmt.Errorf("wait container: %v", err)
	}

	// A non-zero exit status means the build failed.
	if exit != 0 {
		return "", fmt.Errorf("container returned a non-zero exit code: %d", exit)
	}

	// TODO: Return sha256
	return "", nil
}

func (b *DockerBuilder) dryRun() string {
	if b.DryRun {
		return "true"
	}
	return ""
}

func (b *DockerBuilder) image() string {
	if b.Image == "" {
		return DefaultBuilderImage
	}
	return b.Image
}

func (b *DockerBuilder) dataVolume() string {
	if b.DataVolume == "" {
		return DefaultDataVolume
	}
	return b.DataVolume
}

func (b *DockerBuilder) cache(opts BuildOptions) string {
	if opts.NoCache {
		return "off"
	}

	return "on"
}

// statusUpdaterBuilder is a Builder implementation that updates the commit
// status in github.
type statusUpdaterBuilder struct {
	Builder
	github GitHubClient
	since  func(time.Time) time.Duration
}

// UpdateGitHubCommitStatus wraps b to update the GitHub commit status when a
// build starts, and stops.
func UpdateGitHubCommitStatus(b Builder, g GitHubClient) *statusUpdaterBuilder {
	return &statusUpdaterBuilder{
		Builder: b,
		github:  g,
		since:   time.Since,
	}
}

func (b *statusUpdaterBuilder) Build(ctx context.Context, w Logger, opts BuildOptions) (id string, err error) {
	t := time.Now()

	defer func() {
		duration := b.since(t)
		description := fmt.Sprintf("Image built in %v.", duration)
		status := "success"
		if err != nil {
			status = "failure"
			description = err.Error()
		}
		b.updateStatus(w, opts, status, description)
	}()

	if err = b.updateStatus(w, opts, "pending", "Image building."); err != nil {
		err = fmt.Errorf("status: %v", err)
		return
	}

	id, err = b.Builder.Build(ctx, w, opts)
	return
}

// updateStatus updates the given commit with a new status.
func (b *statusUpdaterBuilder) updateStatus(w Logger, opts BuildOptions, status string, description string) error {
	context := Context
	parts := strings.SplitN(opts.Repository, "/", 2)

	var desc *string
	if description != "" {
		desc = &description
	}
	var url *string
	if status == "success" || status == "failure" || status == "error" {
		url = github.String(w.URL())
	}

	_, _, err := b.github.CreateStatus(parts[0], parts[1], opts.Sha, &github.RepoStatus{
		State:       &status,
		Context:     &context,
		Description: desc,
		TargetURL:   url,
	})
	return err
}

// BuildAsync wraps a Builder to run the build in a goroutine.
func BuildAsync(b Builder) Builder {
	build := func(ctx context.Context, w Logger, opts BuildOptions) {
		if _, err := b.Build(ctx, w, opts); err != nil {
			log.Printf("build err: %v", err)
		}
	}

	return BuilderFunc(func(ctx context.Context, w Logger, opts BuildOptions) (string, error) {
		go build(ctx, w, opts)
		return "", nil
	})
}
<|endoftext|>


<commit_before>package network

import (
	"fmt"
	"net"
	"strings"

	"github.com/docker/cli/cli"
	"github.com/docker/cli/cli/command"
	"github.com/docker/cli/opts"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/network"
	runconfigopts "github.com/docker/docker/runconfig/opts"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
)

type createOptions struct {
	name       string
	driver     string
	driverOpts opts.MapOpts
	labels     opts.ListOpts
	internal   bool
	ipv6       bool
	attachable bool
	ingress    bool

	ipamDriver  string
	ipamSubnet  []string
	ipamIPRange []string
	ipamGateway []string
	ipamAux     opts.MapOpts
	ipamOpt     opts.MapOpts
}

func newCreateCommand(dockerCli *command.DockerCli) *cobra.Command {
	options := createOptions{
		driverOpts: *opts.NewMapOpts(nil, nil),
		labels:     opts.NewListOpts(opts.ValidateEnv),
		ipamAux:    *opts.NewMapOpts(nil, nil),
		ipamOpt:    *opts.NewMapOpts(nil, nil),
	}

	cmd := &cobra.Command{
		Use:   "create [OPTIONS] NETWORK",
		Short: "Create a network",
		Args:  cli.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			options.name = args[0]
			return runCreate(dockerCli, options)
		},
	}

	flags := cmd.Flags()
	flags.StringVarP(&options.driver, "driver", "d", "bridge", "Driver to manage the Network")
	flags.VarP(&options.driverOpts, "opt", "o", "Set driver specific options")
	flags.Var(&options.labels, "label", "Set metadata on a network")
	flags.BoolVar(&options.internal, "internal", false, "Restrict external access to the network")
	flags.BoolVar(&options.ipv6, "ipv6", false, "Enable IPv6 networking")
	flags.BoolVar(&options.attachable, "attachable", false, "Enable manual container attachment")
	flags.SetAnnotation("attachable", "version", []string{"1.25"})
	flags.BoolVar(&options.ingress, "ingress", false, "Create swarm routing-mesh network")
	flags.SetAnnotation("ingress", "version", []string{"1.29"})

	flags.StringVar(&options.ipamDriver, "ipam-driver", "default", "IP Address Management Driver")
	flags.StringSliceVar(&options.ipamSubnet, "subnet", []string{}, "Subnet in CIDR format that represents a network segment")
	flags.StringSliceVar(&options.ipamIPRange, "ip-range", []string{}, "Allocate container ip from a sub-range")
	flags.StringSliceVar(&options.ipamGateway, "gateway", []string{}, "IPv4 or IPv6 Gateway for the master subnet")

	flags.Var(&options.ipamAux, "aux-address", "Auxiliary IPv4 or IPv6 addresses used by Network driver")
	flags.Var(&options.ipamOpt, "ipam-opt", "Set IPAM driver specific options")

	return cmd
}

func runCreate(dockerCli *command.DockerCli, options createOptions) error {
	client := dockerCli.Client()

	ipamCfg, err := consolidateIpam(options.ipamSubnet, options.ipamIPRange, options.ipamGateway, options.ipamAux.GetAll())
	if err != nil {
		return err
	}

	// Construct network create request body
	nc := types.NetworkCreate{
		Driver:  options.driver,
		Options: options.driverOpts.GetAll(),
		IPAM: &network.IPAM{
			Driver:  options.ipamDriver,
			Config:  ipamCfg,
			Options: options.ipamOpt.GetAll(),
		},
		CheckDuplicate: true,
		Internal:       options.internal,
		EnableIPv6:     options.ipv6,
		Attachable:     options.attachable,
		Ingress:        options.ingress,
		Labels:         runconfigopts.ConvertKVStringsToMap(options.labels.GetAll()),
	}

	resp, err := client.NetworkCreate(context.Background(), options.name, nc)
	if err != nil {
		return err
	}
	fmt.Fprintf(dockerCli.Out(), "%s\n", resp.ID)
	return nil
}

// Consolidates the ipam configuration as a group from different related configurations
// user can configure network with multiple non-overlapping subnets and hence it is
// possible to correlate the various related parameters and consolidate them.
// consolidateIpam consolidates subnets, ip-ranges, gateways and auxiliary addresses into
// structured ipam data.
// nolint: gocyclo
func consolidateIpam(subnets, ranges, gateways []string, auxaddrs map[string]string) ([]network.IPAMConfig, error) {
	if len(subnets) < len(ranges) || len(subnets) < len(gateways) {
		return nil, errors.Errorf("every ip-range or gateway must have a corresponding subnet")
	}
	iData := map[string]*network.IPAMConfig{}

	// Populate non-overlapping subnets into consolidation map
	for _, s := range subnets {
		for k := range iData {
			ok1, err := subnetMatches(s, k)
			if err != nil {
				return nil, err
			}
			ok2, err := subnetMatches(k, s)
			if err != nil {
				return nil, err
			}
			if ok1 || ok2 {
				return nil, errors.Errorf("multiple overlapping subnet configuration is not supported")
			}
		}
		iData[s] = &network.IPAMConfig{Subnet: s, AuxAddress: map[string]string{}}
	}

	// Validate and add valid ip ranges
	for _, r := range ranges {
		match := false
		for _, s := range subnets {
			ok, err := subnetMatches(s, r)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			if iData[s].IPRange != "" {
				return nil, errors.Errorf("cannot configure multiple ranges (%s, %s) on the same subnet (%s)", r, iData[s].IPRange, s)
			}
			d := iData[s]
			d.IPRange = r
			match = true
		}
		if !match {
			return nil, errors.Errorf("no matching subnet for range %s", r)
		}
	}

	// Validate and add valid gateways
	for _, g := range gateways {
		match := false
		for _, s := range subnets {
			ok, err := subnetMatches(s, g)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			if iData[s].Gateway != "" {
				return nil, errors.Errorf("cannot configure multiple gateways (%s, %s) for the same subnet (%s)", g, iData[s].Gateway, s)
			}
			d := iData[s]
			d.Gateway = g
			match = true
		}
		if !match {
			return nil, errors.Errorf("no matching subnet for gateway %s", g)
		}
	}

	// Validate and add aux-addresses
	for key, aa := range auxaddrs {
		match := false
		for _, s := range subnets {
			ok, err := subnetMatches(s, aa)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			iData[s].AuxAddress[key] = aa
			match = true
		}
		if !match {
			return nil, errors.Errorf("no matching subnet for aux-address %s", aa)
		}
	}

	idl := []network.IPAMConfig{}
	for _, v := range iData {
		idl = append(idl, *v)
	}
	return idl, nil
}

func subnetMatches(subnet, data string) (bool, error) {
	var (
		ip net.IP
	)

	_, s, err := net.ParseCIDR(subnet)
	if err != nil {
		return false, errors.Errorf("Invalid subnet %s : %v", s, err)
	}

	if strings.Contains(data, "/") {
		ip, _, err = net.ParseCIDR(data)
		if err != nil {
			return false, errors.Errorf("Invalid cidr %s : %v", data, err)
		}
	} else {
		ip = net.ParseIP(data)
	}

	return s.Contains(ip), nil
}
<commit_msg>Add support for config-only, config-from and scope options<commit_after>package network

import (
	"fmt"
	"net"
	"strings"

	"github.com/docker/cli/cli"
	"github.com/docker/cli/cli/command"
	"github.com/docker/cli/opts"
	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/network"
	runconfigopts "github.com/docker/docker/runconfig/opts"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
)

type createOptions struct {
	name       string
	scope      string
	driver     string
	driverOpts opts.MapOpts
	labels     opts.ListOpts
	internal   bool
	ipv6       bool
	attachable bool
	ingress    bool
	configOnly bool
	configFrom string

	ipamDriver  string
	ipamSubnet  []string
	ipamIPRange []string
	ipamGateway []string
	ipamAux     opts.MapOpts
	ipamOpt     opts.MapOpts
}

func newCreateCommand(dockerCli *command.DockerCli) *cobra.Command {
	options := createOptions{
		driverOpts: *opts.NewMapOpts(nil, nil),
		labels:     opts.NewListOpts(opts.ValidateEnv),
		ipamAux:    *opts.NewMapOpts(nil, nil),
		ipamOpt:    *opts.NewMapOpts(nil, nil),
	}

	cmd := &cobra.Command{
		Use:   "create [OPTIONS] NETWORK",
		Short: "Create a network",
		Args:  cli.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			options.name = args[0]
			return runCreate(dockerCli, options)
		},
	}

	flags := cmd.Flags()
	flags.StringVarP(&options.driver, "driver", "d", "bridge", "Driver to manage the Network")
	flags.VarP(&options.driverOpts, "opt", "o", "Set driver specific options")
	flags.Var(&options.labels, "label", "Set metadata on a network")
	flags.BoolVar(&options.internal, "internal", false, "Restrict external access to the network")
	flags.BoolVar(&options.ipv6, "ipv6", false, "Enable IPv6 networking")
	flags.BoolVar(&options.attachable, "attachable", false, "Enable manual container attachment")
	flags.SetAnnotation("attachable", "version", []string{"1.25"})
	flags.BoolVar(&options.ingress, "ingress", false, "Create swarm routing-mesh network")
	flags.SetAnnotation("ingress", "version", []string{"1.29"})
	flags.StringVar(&options.scope, "scope", "", "Control the network's scope")
	flags.SetAnnotation("scope", "version", []string{"1.30"})
	flags.BoolVar(&options.configOnly, "config-only", false, "Create a configuration only network")
	flags.SetAnnotation("config-only", "version", []string{"1.30"})
	flags.StringVar(&options.configFrom, "config-from", "", "The network from which copying the configuration")
	flags.SetAnnotation("config-from", "version", []string{"1.30"})

	flags.StringVar(&options.ipamDriver, "ipam-driver", "default", "IP Address Management Driver")
	flags.StringSliceVar(&options.ipamSubnet, "subnet", []string{}, "Subnet in CIDR format that represents a network segment")
	flags.StringSliceVar(&options.ipamIPRange, "ip-range", []string{}, "Allocate container ip from a sub-range")
	flags.StringSliceVar(&options.ipamGateway, "gateway", []string{}, "IPv4 or IPv6 Gateway for the master subnet")

	flags.Var(&options.ipamAux, "aux-address", "Auxiliary IPv4 or IPv6 addresses used by Network driver")
	flags.Var(&options.ipamOpt, "ipam-opt", "Set IPAM driver specific options")

	return cmd
}

func runCreate(dockerCli *command.DockerCli, options createOptions) error {
	client := dockerCli.Client()

	ipamCfg, err := consolidateIpam(options.ipamSubnet, options.ipamIPRange, options.ipamGateway, options.ipamAux.GetAll())
	if err != nil {
		return err
	}

	// Construct network create request body
	nc := types.NetworkCreate{
		Driver:  options.driver,
		Options: options.driverOpts.GetAll(),
		IPAM: &network.IPAM{
			Driver:  options.ipamDriver,
			Config:  ipamCfg,
			Options: options.ipamOpt.GetAll(),
		},
		CheckDuplicate: true,
		Internal:       options.internal,
		EnableIPv6:     options.ipv6,
		Attachable:     options.attachable,
		Ingress:        options.ingress,
		Scope:          options.scope,
		ConfigOnly:     options.configOnly,
		Labels:         runconfigopts.ConvertKVStringsToMap(options.labels.GetAll()),
	}

	if from := options.configFrom; from != "" {
		nc.ConfigFrom = &network.ConfigReference{
			Network: from,
		}
	}

	resp, err := client.NetworkCreate(context.Background(), options.name, nc)
	if err != nil {
		return err
	}
	fmt.Fprintf(dockerCli.Out(), "%s\n", resp.ID)
	return nil
}

// Consolidates the ipam configuration as a group from different related configurations
// user can configure network with multiple non-overlapping subnets and hence it is
// possible to correlate the various related parameters and consolidate them.
// consolidateIpam consolidates subnets, ip-ranges, gateways and auxiliary addresses into
// structured ipam data.
// nolint: gocyclo
func consolidateIpam(subnets, ranges, gateways []string, auxaddrs map[string]string) ([]network.IPAMConfig, error) {
	if len(subnets) < len(ranges) || len(subnets) < len(gateways) {
		return nil, errors.Errorf("every ip-range or gateway must have a corresponding subnet")
	}
	iData := map[string]*network.IPAMConfig{}

	// Populate non-overlapping subnets into consolidation map
	for _, s := range subnets {
		for k := range iData {
			ok1, err := subnetMatches(s, k)
			if err != nil {
				return nil, err
			}
			ok2, err := subnetMatches(k, s)
			if err != nil {
				return nil, err
			}
			if ok1 || ok2 {
				return nil, errors.Errorf("multiple overlapping subnet configuration is not supported")
			}
		}
		iData[s] = &network.IPAMConfig{Subnet: s, AuxAddress: map[string]string{}}
	}

	// Validate and add valid ip ranges
	for _, r := range ranges {
		match := false
		for _, s := range subnets {
			ok, err := subnetMatches(s, r)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			if iData[s].IPRange != "" {
				return nil, errors.Errorf("cannot configure multiple ranges (%s, %s) on the same subnet (%s)", r, iData[s].IPRange, s)
			}
			d := iData[s]
			d.IPRange = r
			match = true
		}
		if !match {
			return nil, errors.Errorf("no matching subnet for range %s", r)
		}
	}

	// Validate and add valid gateways
	for _, g := range gateways {
		match := false
		for _, s := range subnets {
			ok, err := subnetMatches(s, g)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			if iData[s].Gateway != "" {
				return nil, errors.Errorf("cannot configure multiple gateways (%s, %s) for the same subnet (%s)", g, iData[s].Gateway, s)
			}
			d := iData[s]
			d.Gateway = g
			match = true
		}
		if !match {
			return nil, errors.Errorf("no matching subnet for gateway %s", g)
		}
	}

	// Validate and add aux-addresses
	for key, aa := range auxaddrs {
		match := false
		for _, s := range subnets {
			ok, err := subnetMatches(s, aa)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			iData[s].AuxAddress[key] = aa
			match = true
		}
		if !match {
			return nil, errors.Errorf("no matching subnet for aux-address %s", aa)
		}
	}

	idl := []network.IPAMConfig{}
	for _, v := range iData {
		idl = append(idl, *v)
	}
	return idl, nil
}

func subnetMatches(subnet, data string) (bool, error) {
	var (
		ip net.IP
	)

	_, s, err := net.ParseCIDR(subnet)
	if err != nil {
		return false, errors.Errorf("Invalid subnet %s : %v", s, err)
	}

	if strings.Contains(data, "/") {
		ip, _, err = net.ParseCIDR(data)
		if err != nil {
			return false, errors.Errorf("Invalid cidr %s : %v", data, err)
		}
	} else {
		ip = net.ParseIP(data)
	}

	return s.Contains(ip), nil
}
<|endoftext|>


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

package lib

import (
	"context"
	"fmt"
	"io/ioutil"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/tree"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/client/isolated"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/errors"
	"go.chromium.org/luci/common/system/signals"
)

// CmdArchive returns an object for the `archive` subcommand.
func CmdArchive() *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "archive <options>...",
		ShortDesc: "archive dirs/files to CAS",
		LongDesc: `Given a list of files and directories, creates a CAS tree and digest file for
that and uploads the tree to the CAS server.

When specifying directories and files, you must also specify a current working
directory for each file and directory. The current working directory will not be
included in the archived path. For example, to archive './usr/foo/bar' and have
it appear as 'foo/bar' in the CAS tree, specify '-paths ./usr:foo/bar' or
'-paths usr:foo/bar'. And all workind directories should be the same. When the
CAS tree is then downloaded, it will then appear under 'foo/bar' in the desired
directory.

Note that '.' may be omitted in general, so to upload 'foo' from the current
working directory, '-paths :foo' is sufficient.`,
		CommandRun: func() subcommands.CommandRun {
			c := archiveRun{}
			c.Init()
			c.Flags.Var(&c.paths, "paths", "File(s)/Directory(ies) to archive. Specify as <working directory>:<relative path to file/dir>")
			c.Flags.StringVar(&c.dumpDigest, "dump-digest", "", "Dump uploaded CAS root digest in the format of `<Hash>/<Size>`")
			return &c
		},
	}
}

type archiveRun struct {
	commonFlags
	paths      isolated.ScatterGather
	dumpDigest string
}

func (c *archiveRun) parse(a subcommands.Application, args []string) error {
	if err := c.commonFlags.Parse(); err != nil {
		return err
	}
	if len(args) != 0 {
		return errors.Reason("position arguments not expected").Err()
	}
	return nil
}

// getRoot returns root directory if there is only one working directory.
func getRoot(paths isolated.ScatterGather) (string, error) {
	var rel0, wd0 string
	pickedOne := false
	for rel, wd := range paths {
		if !pickedOne {
			rel0 = rel
			wd0 = wd
			pickedOne = true
			continue
		}

		if wd0 != wd {
			return "", errors.Reason("different root (working) directory is not supported: %s:%s vs %s:%s", wd0, rel0, wd, rel).Err()
		}
	}

	if !pickedOne {
		return "", errors.Reason("-paths should be specified at least once").Err()
	}

	return wd0, nil
}

// Does the archive by uploading to isolate-server.
func (c *archiveRun) doArchive(ctx context.Context) ([]digest.Digest, error) {
	ctx, cancel := context.WithCancel(ctx)
	signals.HandleInterrupt(cancel)

	root, err := getRoot(c.paths)
	if err != nil {
		return nil, err
	}

	is := command.InputSpec{}
	for path := range c.paths {
		is.Inputs = append(is.Inputs, path)
	}

	rootDg, chunkers, _, err := tree.ComputeMerkleTree(root, &is, chunker.DefaultChunkSize, filemetadata.NewNoopCache())
	if err != nil {
		return nil, errors.Annotate(err, "failed to call ComputeMerkleTree").Err()
	}

	client, err := newCasClient(ctx, c.casFlags.Instance, c.casFlags.TokenServerHost, false)
	if err != nil {
		return nil, errors.Annotate(err, "failed to create cas client").Err()
	}
	defer client.Close()

	uploadedDigests, err := client.UploadIfMissing(ctx, chunkers...)
	if err != nil {
		return nil, errors.Annotate(err, "failed to call UploadIfMissing").Err()
	}

	if dd := c.dumpDigest; dd != "" {
		if err := ioutil.WriteFile(dd, []byte(rootDg.String()), 0600); err != nil {
			return nil, errors.Annotate(err, "failed to dump digest").Err()
		}
	}

	return uploadedDigests, nil
}

func (c *archiveRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	ctx := cli.GetContext(a, c, env)

	if err := c.parse(a, args); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	defer c.profiler.Stop()

	// TODO: handle the stats
	if _, err := c.doArchive(ctx); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	return 0
}
<commit_msg>[cas] add dump-stats-json flag to `archive`<commit_after>// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package lib

import (
	"context"
	"fmt"
	"io/ioutil"

	"github.com/bazelbuild/remote-apis-sdks/go/pkg/chunker"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/command"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/digest"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/filemetadata"
	"github.com/bazelbuild/remote-apis-sdks/go/pkg/tree"
	"github.com/maruel/subcommands"

	"go.chromium.org/luci/client/isolated"
	"go.chromium.org/luci/common/cli"
	"go.chromium.org/luci/common/errors"
	isol "go.chromium.org/luci/common/isolated"
	"go.chromium.org/luci/common/system/signals"
)

// CmdArchive returns an object for the `archive` subcommand.
func CmdArchive() *subcommands.Command {
	return &subcommands.Command{
		UsageLine: "archive <options>...",
		ShortDesc: "archive dirs/files to CAS",
		LongDesc: `Given a list of files and directories, creates a CAS tree and digest file for
that and uploads the tree to the CAS server.

When specifying directories and files, you must also specify a current working
directory for each file and directory. The current working directory will not be
included in the archived path. For example, to archive './usr/foo/bar' and have
it appear as 'foo/bar' in the CAS tree, specify '-paths ./usr:foo/bar' or
'-paths usr:foo/bar'. And all workind directories should be the same. When the
CAS tree is then downloaded, it will then appear under 'foo/bar' in the desired
directory.

Note that '.' may be omitted in general, so to upload 'foo' from the current
working directory, '-paths :foo' is sufficient.`,
		CommandRun: func() subcommands.CommandRun {
			c := archiveRun{}
			c.Init()
			c.Flags.Var(&c.paths, "paths", "File(s)/Directory(ies) to archive. Specify as <working directory>:<relative path to file/dir>")
			c.Flags.StringVar(&c.dumpDigest, "dump-digest", "", "Dump uploaded CAS root digest in the format of `<Hash>/<Size>`")
			c.Flags.StringVar(&c.dumpStatsJSON, "dump-stats-json", "", "Dump upload stats to json file.")
			return &c
		},
	}
}

type archiveRun struct {
	commonFlags
	paths         isolated.ScatterGather
	dumpDigest    string
	dumpStatsJSON string
}

func (c *archiveRun) parse(a subcommands.Application, args []string) error {
	if err := c.commonFlags.Parse(); err != nil {
		return err
	}
	if len(args) != 0 {
		return errors.Reason("position arguments not expected").Err()
	}
	return nil
}

// getRoot returns root directory if there is only one working directory.
func getRoot(paths isolated.ScatterGather) (string, error) {
	var rel0, wd0 string
	pickedOne := false
	for rel, wd := range paths {
		if !pickedOne {
			rel0 = rel
			wd0 = wd
			pickedOne = true
			continue
		}

		if wd0 != wd {
			return "", errors.Reason("different root (working) directory is not supported: %s:%s vs %s:%s", wd0, rel0, wd, rel).Err()
		}
	}

	if !pickedOne {
		return "", errors.Reason("-paths should be specified at least once").Err()
	}

	return wd0, nil
}

// Does the archive by uploading to isolate-server.
func (c *archiveRun) doArchive(ctx context.Context) ([]digest.Digest, error) {
	ctx, cancel := context.WithCancel(ctx)
	signals.HandleInterrupt(cancel)

	root, err := getRoot(c.paths)
	if err != nil {
		return nil, err
	}

	is := command.InputSpec{}
	for path := range c.paths {
		is.Inputs = append(is.Inputs, path)
	}

	rootDg, chunkers, _, err := tree.ComputeMerkleTree(root, &is, chunker.DefaultChunkSize, filemetadata.NewNoopCache())
	if err != nil {
		return nil, errors.Annotate(err, "failed to call ComputeMerkleTree").Err()
	}

	client, err := newCasClient(ctx, c.casFlags.Instance, c.casFlags.TokenServerHost, false)
	if err != nil {
		return nil, errors.Annotate(err, "failed to create cas client").Err()
	}
	defer client.Close()

	uploadedDigests, err := client.UploadIfMissing(ctx, chunkers...)
	if err != nil {
		return nil, errors.Annotate(err, "failed to call UploadIfMissing").Err()
	}

	if dd := c.dumpDigest; dd != "" {
		if err := ioutil.WriteFile(dd, []byte(rootDg.String()), 0600); err != nil {
			return nil, errors.Annotate(err, "failed to dump digest").Err()
		}
	}

	if dsj := c.dumpStatsJSON; dsj != "" {
		uploaded := make([]int64, 0, len(uploadedDigests))
		uploadedSet := make(map[digest.Digest]struct{})
		for _, d := range uploadedDigests {
			uploaded = append(uploaded, d.Size)
			uploadedSet[d] = struct{}{}
		}

		notUploaded := make([]int64, 0, len(chunkers)-len(uploadedDigests))
		for _, d := range chunkers {
			if _, ok := uploadedSet[d.Digest()]; !ok {
				notUploaded = append(notUploaded, d.Digest().Size)
			}
		}

		if err := isol.WriteStats(dsj, notUploaded, uploaded); err != nil {
			return nil, errors.Annotate(err, "failed to write stats json").Err()
		}
	}

	return uploadedDigests, nil
}

func (c *archiveRun) Run(a subcommands.Application, args []string, env subcommands.Env) int {
	ctx := cli.GetContext(a, c, env)

	if err := c.parse(a, args); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}
	defer c.profiler.Stop()

	if _, err := c.doArchive(ctx); err != nil {
		errors.Log(ctx, err)
		fmt.Fprintf(a.GetErr(), "%s: %s\n", a.GetName(), err)
		return 1
	}

	return 0
}
<|endoftext|>


<commit_before>package fingerprint

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	log "github.com/hashicorp/go-hclog"

	cleanhttp "github.com/hashicorp/go-cleanhttp"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// This is where the AWS metadata server normally resides. We hardcode the
	// "instance" path as well since it's the only one we access here.
	DEFAULT_AWS_URL = "http://169.254.169.254/latest/meta-data/"

	// AwsMetadataTimeout is the timeout used when contacting the AWS metadata
	// service
	AwsMetadataTimeout = 2 * time.Second
)

// map of instance type to approximate speed, in Mbits/s
// Estimates from http://stackoverflow.com/a/35806587
// This data is meant for a loose approximation
var ec2InstanceSpeedMap = map[*regexp.Regexp]int{
	regexp.MustCompile("t2.nano"):      30,
	regexp.MustCompile("t2.micro"):     70,
	regexp.MustCompile("t2.small"):     125,
	regexp.MustCompile("t2.medium"):    300,
	regexp.MustCompile("m3.medium"):    400,
	regexp.MustCompile("c4.8xlarge"):   4000,
	regexp.MustCompile("x1.16xlarge"):  5000,
	regexp.MustCompile(`.*\.large`):    500,
	regexp.MustCompile(`.*\.xlarge`):   750,
	regexp.MustCompile(`.*\.2xlarge`):  1000,
	regexp.MustCompile(`.*\.4xlarge`):  2000,
	regexp.MustCompile(`.*\.8xlarge`):  10000,
	regexp.MustCompile(`.*\.10xlarge`): 10000,
	regexp.MustCompile(`.*\.16xlarge`): 10000,
	regexp.MustCompile(`.*\.32xlarge`): 10000,
}

// EnvAWSFingerprint is used to fingerprint AWS metadata
type EnvAWSFingerprint struct {
	StaticFingerprinter
	timeout time.Duration
	logger  log.Logger
}

// NewEnvAWSFingerprint is used to create a fingerprint from AWS metadata
func NewEnvAWSFingerprint(logger log.Logger) Fingerprint {
	f := &EnvAWSFingerprint{
		logger:  logger.Named("env_aws"),
		timeout: AwsMetadataTimeout,
	}
	return f
}

func (f *EnvAWSFingerprint) Fingerprint(request *FingerprintRequest, response *FingerprintResponse) error {
	cfg := request.Config

	// Check if we should tighten the timeout
	if cfg.ReadBoolDefault(TightenNetworkTimeoutsConfig, false) {
		f.timeout = 1 * time.Millisecond
	}

	if !f.isAWS() {
		return nil
	}

	// newNetwork is populated and added to the Nodes resources
	newNetwork := &structs.NetworkResource{
		Device: "eth0",
	}

	metadataURL := os.Getenv("AWS_ENV_URL")
	if metadataURL == "" {
		metadataURL = DEFAULT_AWS_URL
	}

	client := &http.Client{
		Timeout:   f.timeout,
		Transport: cleanhttp.DefaultTransport(),
	}

	// Keys and whether they should be namespaced as unique. Any key whose value
	// uniquely identifies a node, such as ip, should be marked as unique. When
	// marked as unique, the key isn't included in the computed node class.
	keys := map[string]bool{
		"ami-id":                      false,
		"hostname":                    true,
		"instance-id":                 true,
		"instance-type":               false,
		"local-hostname":              true,
		"local-ipv4":                  true,
		"public-hostname":             true,
		"public-ipv4":                 true,
		"placement/availability-zone": false,
	}
	for k, unique := range keys {
		res, err := client.Get(metadataURL + k)
		if res.StatusCode != http.StatusOK {
			f.logger.Debug("could not read attribute value", "attribute", k)
			continue
		}
		if err != nil {
			// if it's a URL error, assume we're not in an AWS environment
			// TODO: better way to detect AWS? Check xen virtualization?
			if _, ok := err.(*url.Error); ok {
				return nil
			}
			// not sure what other errors it would return
			return err
		}
		resp, err := ioutil.ReadAll(res.Body)
		res.Body.Close()
		if err != nil {
			f.logger.Error("error reading response body for AWS attribute", "attribute", k, "error", err)
		}

		// assume we want blank entries
		key := "platform.aws." + strings.Replace(k, "/", ".", -1)
		if unique {
			key = structs.UniqueNamespace(key)
		}

		response.AddAttribute(key, strings.Trim(string(resp), "\n"))
	}

	// copy over network specific information
	if val, ok := response.Attributes["unique.platform.aws.local-ipv4"]; ok && val != "" {
		response.AddAttribute("unique.network.ip-address", val)
		newNetwork.IP = val
		newNetwork.CIDR = newNetwork.IP + "/32"
	}

	// find LinkSpeed from lookup
	throughput := f.linkSpeed()
	if cfg.NetworkSpeed != 0 {
		throughput = cfg.NetworkSpeed
	} else if throughput == 0 {
		// Failed to determine speed. Check if the network fingerprint got it
		found := false
		if request.Node.Resources != nil && len(request.Node.Resources.Networks) > 0 {
			for _, n := range request.Node.Resources.Networks {
				if n.IP == newNetwork.IP {
					throughput = n.MBits
					found = true
					break
				}
			}
		}

		// Nothing detected so default
		if !found {
			throughput = defaultNetworkSpeed
		}
	}

	newNetwork.MBits = throughput
	response.NodeResources = &structs.NodeResources{
		Networks: []*structs.NetworkResource{newNetwork},
	}

	// populate Links
	response.AddLink("aws.ec2", fmt.Sprintf("%s.%s",
		response.Attributes["platform.aws.placement.availability-zone"],
		response.Attributes["unique.platform.aws.instance-id"]))
	response.Detected = true

	return nil
}

func (f *EnvAWSFingerprint) isAWS() bool {
	// Read the internal metadata URL from the environment, allowing test files to
	// provide their own
	metadataURL := os.Getenv("AWS_ENV_URL")
	if metadataURL == "" {
		metadataURL = DEFAULT_AWS_URL
	}

	client := &http.Client{
		Timeout:   f.timeout,
		Transport: cleanhttp.DefaultTransport(),
	}

	// Query the metadata url for the ami-id, to verify we're on AWS
	resp, err := client.Get(metadataURL + "ami-id")
	if err != nil {
		f.logger.Debug("error querying AWS Metadata URL, skipping")
		return false
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		// URL not found, which indicates that this isn't AWS
		return false
	}

	instanceID, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		f.logger.Debug("error reading AWS Instance ID, skipping")
		return false
	}

	match, err := regexp.MatchString("ami-*", string(instanceID))
	if err != nil || !match {
		return false
	}

	return true
}

// EnvAWSFingerprint uses lookup table to approximate network speeds
func (f *EnvAWSFingerprint) linkSpeed() int {

	// Query the API for the instance type, and use the table above to approximate
	// the network speed
	metadataURL := os.Getenv("AWS_ENV_URL")
	if metadataURL == "" {
		metadataURL = DEFAULT_AWS_URL
	}

	client := &http.Client{
		Timeout:   f.timeout,
		Transport: cleanhttp.DefaultTransport(),
	}

	res, err := client.Get(metadataURL + "instance-type")
	if err != nil {
		f.logger.Error("error reading instance-type", "error", err)
		return 0
	}

	body, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		f.logger.Error("error reading response body for instance-type", "error", err)
		return 0
	}

	key := strings.Trim(string(body), "\n")
	netSpeed := 0
	for reg, speed := range ec2InstanceSpeedMap {
		if reg.MatchString(key) {
			netSpeed = speed
			break
		}
	}

	return netSpeed
}
<commit_msg>fix nil pointer in fingerprinting AWS env leading to crash<commit_after>package fingerprint

import (
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"regexp"
	"strings"
	"time"

	log "github.com/hashicorp/go-hclog"

	cleanhttp "github.com/hashicorp/go-cleanhttp"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// This is where the AWS metadata server normally resides. We hardcode the
	// "instance" path as well since it's the only one we access here.
	DEFAULT_AWS_URL = "http://169.254.169.254/latest/meta-data/"

	// AwsMetadataTimeout is the timeout used when contacting the AWS metadata
	// service
	AwsMetadataTimeout = 2 * time.Second
)

// map of instance type to approximate speed, in Mbits/s
// Estimates from http://stackoverflow.com/a/35806587
// This data is meant for a loose approximation
var ec2InstanceSpeedMap = map[*regexp.Regexp]int{
	regexp.MustCompile("t2.nano"):      30,
	regexp.MustCompile("t2.micro"):     70,
	regexp.MustCompile("t2.small"):     125,
	regexp.MustCompile("t2.medium"):    300,
	regexp.MustCompile("m3.medium"):    400,
	regexp.MustCompile("c4.8xlarge"):   4000,
	regexp.MustCompile("x1.16xlarge"):  5000,
	regexp.MustCompile(`.*\.large`):    500,
	regexp.MustCompile(`.*\.xlarge`):   750,
	regexp.MustCompile(`.*\.2xlarge`):  1000,
	regexp.MustCompile(`.*\.4xlarge`):  2000,
	regexp.MustCompile(`.*\.8xlarge`):  10000,
	regexp.MustCompile(`.*\.10xlarge`): 10000,
	regexp.MustCompile(`.*\.16xlarge`): 10000,
	regexp.MustCompile(`.*\.32xlarge`): 10000,
}

// EnvAWSFingerprint is used to fingerprint AWS metadata
type EnvAWSFingerprint struct {
	StaticFingerprinter
	timeout time.Duration
	logger  log.Logger
}

// NewEnvAWSFingerprint is used to create a fingerprint from AWS metadata
func NewEnvAWSFingerprint(logger log.Logger) Fingerprint {
	f := &EnvAWSFingerprint{
		logger:  logger.Named("env_aws"),
		timeout: AwsMetadataTimeout,
	}
	return f
}

func (f *EnvAWSFingerprint) Fingerprint(request *FingerprintRequest, response *FingerprintResponse) error {
	cfg := request.Config

	// Check if we should tighten the timeout
	if cfg.ReadBoolDefault(TightenNetworkTimeoutsConfig, false) {
		f.timeout = 1 * time.Millisecond
	}

	if !f.isAWS() {
		return nil
	}

	// newNetwork is populated and added to the Nodes resources
	newNetwork := &structs.NetworkResource{
		Device: "eth0",
	}

	metadataURL := os.Getenv("AWS_ENV_URL")
	if metadataURL == "" {
		metadataURL = DEFAULT_AWS_URL
	}

	client := &http.Client{
		Timeout:   f.timeout,
		Transport: cleanhttp.DefaultTransport(),
	}

	// Keys and whether they should be namespaced as unique. Any key whose value
	// uniquely identifies a node, such as ip, should be marked as unique. When
	// marked as unique, the key isn't included in the computed node class.
	keys := map[string]bool{
		"ami-id":                      false,
		"hostname":                    true,
		"instance-id":                 true,
		"instance-type":               false,
		"local-hostname":              true,
		"local-ipv4":                  true,
		"public-hostname":             true,
		"public-ipv4":                 true,
		"placement/availability-zone": false,
	}
	for k, unique := range keys {
		res, err := client.Get(metadataURL + k)
		if err != nil {
			// if it's a URL error, assume we're not in an AWS environment
			// TODO: better way to detect AWS? Check xen virtualization?
			if _, ok := err.(*url.Error); ok {
				return nil
			}
			// not sure what other errors it would return
			return err
		} else if res.StatusCode != http.StatusOK {
			f.logger.Debug("could not read attribute value", "attribute", k)
			continue
		}
		resp, err := ioutil.ReadAll(res.Body)
		res.Body.Close()
		if err != nil {
			f.logger.Error("error reading response body for AWS attribute", "attribute", k, "error", err)
		}

		// assume we want blank entries
		key := "platform.aws." + strings.Replace(k, "/", ".", -1)
		if unique {
			key = structs.UniqueNamespace(key)
		}

		response.AddAttribute(key, strings.Trim(string(resp), "\n"))
	}

	// copy over network specific information
	if val, ok := response.Attributes["unique.platform.aws.local-ipv4"]; ok && val != "" {
		response.AddAttribute("unique.network.ip-address", val)
		newNetwork.IP = val
		newNetwork.CIDR = newNetwork.IP + "/32"
	}

	// find LinkSpeed from lookup
	throughput := f.linkSpeed()
	if cfg.NetworkSpeed != 0 {
		throughput = cfg.NetworkSpeed
	} else if throughput == 0 {
		// Failed to determine speed. Check if the network fingerprint got it
		found := false
		if request.Node.Resources != nil && len(request.Node.Resources.Networks) > 0 {
			for _, n := range request.Node.Resources.Networks {
				if n.IP == newNetwork.IP {
					throughput = n.MBits
					found = true
					break
				}
			}
		}

		// Nothing detected so default
		if !found {
			throughput = defaultNetworkSpeed
		}
	}

	newNetwork.MBits = throughput
	response.NodeResources = &structs.NodeResources{
		Networks: []*structs.NetworkResource{newNetwork},
	}

	// populate Links
	response.AddLink("aws.ec2", fmt.Sprintf("%s.%s",
		response.Attributes["platform.aws.placement.availability-zone"],
		response.Attributes["unique.platform.aws.instance-id"]))
	response.Detected = true

	return nil
}

func (f *EnvAWSFingerprint) isAWS() bool {
	// Read the internal metadata URL from the environment, allowing test files to
	// provide their own
	metadataURL := os.Getenv("AWS_ENV_URL")
	if metadataURL == "" {
		metadataURL = DEFAULT_AWS_URL
	}

	client := &http.Client{
		Timeout:   f.timeout,
		Transport: cleanhttp.DefaultTransport(),
	}

	// Query the metadata url for the ami-id, to verify we're on AWS
	resp, err := client.Get(metadataURL + "ami-id")
	if err != nil {
		f.logger.Debug("error querying AWS Metadata URL, skipping")
		return false
	}
	defer resp.Body.Close()

	if resp.StatusCode >= 400 {
		// URL not found, which indicates that this isn't AWS
		return false
	}

	instanceID, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		f.logger.Debug("error reading AWS Instance ID, skipping")
		return false
	}

	match, err := regexp.MatchString("ami-*", string(instanceID))
	if err != nil || !match {
		return false
	}

	return true
}

// EnvAWSFingerprint uses lookup table to approximate network speeds
func (f *EnvAWSFingerprint) linkSpeed() int {

	// Query the API for the instance type, and use the table above to approximate
	// the network speed
	metadataURL := os.Getenv("AWS_ENV_URL")
	if metadataURL == "" {
		metadataURL = DEFAULT_AWS_URL
	}

	client := &http.Client{
		Timeout:   f.timeout,
		Transport: cleanhttp.DefaultTransport(),
	}

	res, err := client.Get(metadataURL + "instance-type")
	if err != nil {
		f.logger.Error("error reading instance-type", "error", err)
		return 0
	}

	body, err := ioutil.ReadAll(res.Body)
	res.Body.Close()
	if err != nil {
		f.logger.Error("error reading response body for instance-type", "error", err)
		return 0
	}

	key := strings.Trim(string(body), "\n")
	netSpeed := 0
	for reg, speed := range ec2InstanceSpeedMap {
		if reg.MatchString(key) {
			netSpeed = speed
			break
		}
	}

	return netSpeed
}
<|endoftext|>


<commit_before>/*
Copyright 2014 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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	kclientcmd "github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd"
	kclientcmdapi "github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/controller/framework"
	kcontrollerFramework "github.com/GoogleCloudPlatform/kubernetes/pkg/controller/framework"
	kSelector "github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client
	// DNS domain name.
	domain string
}

func (ks *kube2consul) removeDNS(record string) error {
	glog.V(2).Infof("Removing %s from DNS", record)
	return ks.consulClient.Agent().ServiceDeregister(record)
}

func (ks *kube2consul) addDNS(record string, service *kapi.Service) error {
	if strings.Contains(record, ".") {
		glog.V(1).Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if PortalIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.V(1).Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
		asr := &consulapi.AgentServiceRegistration{
			ID:			 record,
			Name: 	 record,
			Address: service.Spec.PortalIP,
			Port:    service.Spec.Ports[0].Port,
		}

		glog.V(2).Infof("Setting DNS record: %v -> %v:%d\n", record, service.Spec.PortalIP, service.Spec.Ports[i].Port)
		if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
			return err
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1beta3",
		}
	} else {
		var err error
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(
			&kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile},
			&kclientcmd.ConfigOverrides{ClusterInfo: kclientcmdapi.Cluster{Server: masterUrl, InsecureSkipTLSVerify: true}}).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	return fmt.Sprintf("%s.%s", service, namespace)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *cache.ListWatch {
	return cache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.addDNS(s.Name, s); err != nil {
			glog.V(1).Infof("Failed to add service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.removeDNS(s.Name); err != nil {
			glog.V(1).Infof("Failed to remove service: %v due to: %v", name, err)
		}
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = framework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		framework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: func(oldObj, newObj interface{}) {
				ks.newService(newObj)
			},
		},
	)
	serviceController.Run(util.NeverStop)
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := kube2consul{}

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	watchForServices(kubeClient, &ks)
}
<commit_msg>update API Version tio v1<commit_after>/*
Copyright 2014 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.
*/

// kube2consul is a bridge between Kubernetes and Consul.  It watches the
// Kubernetes master for changes in Services and creates new DNS records on the
// consul agent.

package main

import (
	"flag"
	"fmt"
	"strings"
	"net/url"
	"os"
	"time"

	kapi "github.com/GoogleCloudPlatform/kubernetes/pkg/api"
	kclient "github.com/GoogleCloudPlatform/kubernetes/pkg/client"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/client/cache"
	kclientcmd "github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd"
	kclientcmdapi "github.com/GoogleCloudPlatform/kubernetes/pkg/client/clientcmd/api"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/controller/framework"
	kcontrollerFramework "github.com/GoogleCloudPlatform/kubernetes/pkg/controller/framework"
	kSelector "github.com/GoogleCloudPlatform/kubernetes/pkg/fields"
	"github.com/GoogleCloudPlatform/kubernetes/pkg/util"
	"github.com/golang/glog"
	consulapi "github.com/hashicorp/consul/api"
)

var (
	argConsulAgent         = flag.String("consul-agent", "http://127.0.0.1:8500", "URL to consul agent")
	argKubecfgFile         = flag.String("kubecfg_file", "", "Location of kubecfg file for access to kubernetes service")
	argKubeMasterUrl       = flag.String("kube_master_url", "https://${KUBERNETES_SERVICE_HOST}:${KUBERNETES_SERVICE_PORT}", "Url to reach kubernetes master. Env variables in this flag will be expanded.")
)

const (
	// Maximum number of attempts to connect to consul server.
	maxConnectAttempts = 12
	// Resync period for the kube controller loop.
	resyncPeriod = 5 * time.Second
)

type kube2consul struct {
	// Consul client.
	consulClient *consulapi.Client
	// DNS domain name.
	domain string
}

func (ks *kube2consul) removeDNS(record string) error {
	glog.V(2).Infof("Removing %s from DNS", record)
	return ks.consulClient.Agent().ServiceDeregister(record)
}

func (ks *kube2consul) addDNS(record string, service *kapi.Service) error {
	if strings.Contains(record, ".") {
		glog.V(1).Infof("Service names containing '.' are not supported: %s\n", service.Name)
		return nil
	}

	// if PortalIP is not set, do not create a DNS records
	if !kapi.IsServiceIPSet(service) {
		glog.V(1).Infof("Skipping dns record for headless service: %s\n", service.Name)
		return nil
	}

	for i := range service.Spec.Ports {
		asr := &consulapi.AgentServiceRegistration{
			ID:			 record,
			Name: 	 record,
			Address: service.Spec.PortalIP,
			Port:    service.Spec.Ports[0].Port,
		}

		glog.V(2).Infof("Setting DNS record: %v -> %v:%d\n", record, service.Spec.PortalIP, service.Spec.Ports[i].Port)
		if err := ks.consulClient.Agent().ServiceRegister(asr); err != nil {
			return err
		}
	}
	return nil
}

func newConsulClient(consulAgent string) (*consulapi.Client, error) {
	var (
		client *consulapi.Client
		err    error
	)

	consulConfig := consulapi.DefaultConfig()
	consulAgentUrl, err := url.Parse(consulAgent)
	if err != nil {
			glog.Infof("Error parsing Consul url")
			return nil, err
	}

	if consulAgentUrl.Host != "" {
	  consulConfig.Address = consulAgentUrl.Host
	}

	if consulAgentUrl.Scheme != "" {
		consulConfig.Scheme = consulAgentUrl.Scheme
	}

	client, err = consulapi.NewClient(consulConfig)
	if err != nil {
			glog.Infof("Error creating Consul client")
			return nil, err
	}

	for attempt := 1; attempt <= maxConnectAttempts; attempt++ {
		if _, err = client.Agent().Self(); err == nil {
			break
		}

		if attempt == maxConnectAttempts {
			break
		}

		glog.Infof("[Attempt: %d] Attempting access to Consul after 5 second sleep", attempt)
		time.Sleep(5 * time.Second)
	}

	if err != nil {
		return nil, fmt.Errorf("failed to connect to Consul agent: %v, error: %v", consulAgent, err)
	}
	glog.Infof("Consul agent found: %v", consulAgent)

	return client, nil
}

func getKubeMasterUrl() (string, error) {
	if *argKubeMasterUrl == "" {
		return "", fmt.Errorf("no --kube_master_url specified")
	}
	parsedUrl, err := url.Parse(os.ExpandEnv(*argKubeMasterUrl))
	if err != nil {
		return "", fmt.Errorf("failed to parse --kube_master_url %s - %v", *argKubeMasterUrl, err)
	}
	if parsedUrl.Scheme == "" || parsedUrl.Host == "" || parsedUrl.Host == ":" {
		return "", fmt.Errorf("invalid --kube_master_url specified %s", *argKubeMasterUrl)
	}
	return parsedUrl.String(), nil
}

// TODO: evaluate using pkg/client/clientcmd
func newKubeClient() (*kclient.Client, error) {
	var config *kclient.Config
	masterUrl, err := getKubeMasterUrl()
	if err != nil {
		return nil, err
	}
	if *argKubecfgFile == "" {
		config = &kclient.Config{
			Host:    masterUrl,
			Version: "v1",
		}
	} else {
		var err error
		if config, err = kclientcmd.NewNonInteractiveDeferredLoadingClientConfig(
			&kclientcmd.ClientConfigLoadingRules{ExplicitPath: *argKubecfgFile},
			&kclientcmd.ConfigOverrides{ClusterInfo: kclientcmdapi.Cluster{Server: masterUrl, InsecureSkipTLSVerify: true}}).ClientConfig(); err != nil {
			return nil, err
		}
	}
	glog.Infof("Using %s for kubernetes master", config.Host)
	glog.Infof("Using kubernetes API %s", config.Version)
	return kclient.New(config)
}

func buildNameString(service, namespace string) string {
	return fmt.Sprintf("%s.%s", service, namespace)
}

// Returns a cache.ListWatch that gets all changes to services.
func createServiceLW(kubeClient *kclient.Client) *cache.ListWatch {
	return cache.NewListWatchFromClient(kubeClient, "services", kapi.NamespaceAll, kSelector.Everything())
}

func (ks *kube2consul) newService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.addDNS(s.Name, s); err != nil {
			glog.V(1).Infof("Failed to add service: %v due to: %v", name, err)
		}
	}
}

func (ks *kube2consul) removeService(obj interface{}) {
	if s, ok := obj.(*kapi.Service); ok {
		name := buildNameString(s.Name, s.Namespace)
		if err := ks.removeDNS(s.Name); err != nil {
			glog.V(1).Infof("Failed to remove service: %v due to: %v", name, err)
		}
	}
}

func watchForServices(kubeClient *kclient.Client, ks *kube2consul) {
	var serviceController *kcontrollerFramework.Controller
	_, serviceController = framework.NewInformer(
		createServiceLW(kubeClient),
		&kapi.Service{},
		resyncPeriod,
		framework.ResourceEventHandlerFuncs{
			AddFunc:    ks.newService,
			DeleteFunc: ks.removeService,
			UpdateFunc: func(oldObj, newObj interface{}) {
				ks.newService(newObj)
			},
		},
	)
	serviceController.Run(util.NeverStop)
}

func main() {
	flag.Parse()
	var err error
	// TODO: Validate input flags.
	ks := kube2consul{}

	if ks.consulClient, err = newConsulClient(*argConsulAgent); err != nil {
		glog.Fatalf("Failed to create Consul client - %v", err)
	}

	kubeClient, err := newKubeClient()
	if err != nil {
		glog.Fatalf("Failed to create a kubernetes client: %v", err)
	}

	watchForServices(kubeClient, &ks)
}
<|endoftext|>