<commit_before>package builder

import (
	"github.com/blablacar/cnt/config"
	"github.com/blablacar/cnt/log"
	"github.com/blablacar/cnt/utils"
)

func (p *Pod) Push() {
	log.Get().Info("Push POD", p.manifest.Name)

	for _, e := range p.manifest.Pod.Apps {
		aci, err := NewAciWithManifest(p.path+"/"+e.Name, p.args, p.toAciManifest(e))
		if err != nil {
			log.Get().Panic(err)
		}
		aci.PodName = &p.manifest.Name
		aci.Push()
	}

	utils.ExecCmd("curl", "-i",
		"-F", "r=releases",
		"-F", "hasPom=false",
		"-F", "e=pod",
		"-F", "g=com.blablacar.aci.linux.amd64",
		"-F", "p=pod",
		"-F", "v="+p.manifest.Name.Version(),
		"-F", "a="+p.manifest.Name.ShortName(),
		"-F", "file=@"+p.target+POD_TARGET_MANIFEST,
		"-u", config.GetConfig().Push.Username+":"+config.GetConfig().Push.Password,
		config.GetConfig().Push.Url+"/service/local/artifact/maven/content")

}
<commit_msg>build before push pod<commit_after>package builder

import (
	"github.com/blablacar/cnt/config"
	"github.com/blablacar/cnt/log"
	"github.com/blablacar/cnt/utils"
)

func (p *Pod) Push() {
	log.Get().Info("Push POD", p.manifest.Name)

	p.Build()

	for _, e := range p.manifest.Pod.Apps {
		aci, err := NewAciWithManifest(p.path+"/"+e.Name, p.args, p.toAciManifest(e))
		if err != nil {
			log.Get().Panic(err)
		}
		aci.PodName = &p.manifest.Name
		aci.Push()
	}

	utils.ExecCmd("curl", "-i",
		"-F", "r=releases",
		"-F", "hasPom=false",
		"-F", "e=pod",
		"-F", "g=com.blablacar.aci.linux.amd64",
		"-F", "p=pod",
		"-F", "v="+p.manifest.Name.Version(),
		"-F", "a="+p.manifest.Name.ShortName(),
		"-F", "file=@"+p.target+POD_TARGET_MANIFEST,
		"-u", config.GetConfig().Push.Username+":"+config.GetConfig().Push.Password,
		config.GetConfig().Push.Url+"/service/local/artifact/maven/content")

}
<|endoftext|>


<commit_before>package nsq

import (
	"context"
	"sync"
	"sync/atomic"

	gonsq "github.com/bitly/go-nsq"
)

func NewConsumer(topic, channel string, opt *Option) (*Consumer, error) {
	if opt == nil {
		opt = &Option{}
	}

	// initialized at 0 to pause reading on startup
	maxInFlight := 0

	// create nsq consumer config
	nsqConf := gonsq.NewConfig()
	nsqConf.MaxInFlight = maxInFlight

	// create context for clean shutdown
	ctx, cncl := context.WithCancel(context.Background())

	// create lazy consumer
	c := &Consumer{
		opt:       opt,
		nsqConf:   nsqConf,
		closeChan: make(chan interface{}),
		msgChan:   make(chan []byte),
		ctx:       ctx,
		cncl:      cncl,
	}

	// connect to nsqd
	err := c.connect(topic, channel)
	if err != nil {
		return nil, err
	}

	return c, nil
}

type Consumer struct {
	opt      *Option
	nsqConf  *gonsq.Config
	consumer *gonsq.Consumer

	// msgRequested and msgReceived is necessary for lazy
	// loading b/c of the delay in communicating with nsqd
	// the new RDY count.
	msgRequested int64 // total messages requested
	msgReceived  int64 // total messages received

	// bytes are assumed serializable as a valid Task
	// msgChan should be created with a consumer and should never be closed.
	msgChan chan []byte

	// close signal to shut down Consumer. Will stop any in-process messages
	// (messages are re-queued). Will also shutdown any open nsq consumers.
	closeChan chan interface{}

	// wait group is to make sure all consumers are closed during
	// a shutdown.
	sync.WaitGroup

	// mutex for updating consumer maxInFlight
	sync.Mutex

	// context for clean shutdown
	ctx  context.Context
	cncl context.CancelFunc
}

// connect will connect nsq. An error is returned if there
// is a problem connecting.
func (c *Consumer) connect(topic, channel string) error {
	// initialize nsq consumer - does not connect
	consumer, err := gonsq.NewConsumer(topic, channel, c.nsqConf)
	if err != nil {
		return err
	}

	// set custom logger
	if c.opt.Logger != nil {
		consumer.SetLogger(c.opt.Logger, c.opt.LogLvl)
	}

	consumer.AddHandler(c)

	// attempt to connect to lookupds (if provided)
	// or attempt to connect to nsqds (if provided)
	// if neither is provided attempt to connect to localhost
	if len(c.opt.LookupdAddrs) > 0 {
		err = consumer.ConnectToNSQLookupds(c.opt.LookupdAddrs)
		if err != nil {
			return err
		}
	} else if len(c.opt.NSQdAddrs) > 0 {
		err = consumer.ConnectToNSQDs(c.opt.NSQdAddrs)
		if err != nil {
			return err
		}
	} else {
		err = consumer.ConnectToNSQD("localhost:4150")
		if err != nil {
			return err
		}
	}

	// update nsq consumer
	c.consumer = consumer

	return nil
}

// reqMsg will un-pause messages in flight (if paused)
// (by setting maxInflight = 1)
func (c *Consumer) reqMsg() {
	c.Lock()
	// increase total message requested count
	atomic.AddInt64(&c.msgRequested, 1)

	// start flow of messages
	// can always  set the maxInFlight to 1
	// b/c nsq consumer will not update maxInFlight
	// if there is no change.
	c.consumer.ChangeMaxInFlight(1)

	c.Unlock()
}

// recMsg will attempt to 'receive' the message. If the
// message is received then 'true' is returned'. If the
// message could not be received (b/c all requested messages
// have been received) then 'false' is returned.
func (c *Consumer) recMsg() bool {
	c.Lock()
	defer c.Unlock()

	// check if all messages have been received
	if atomic.LoadInt64(&c.msgReceived) == atomic.LoadInt64(&c.msgRequested) {
		return false // all requested messages have been received
	}

	// increase total message received
	atomic.AddInt64(&c.msgReceived, 1)

	// pause flow if all requested messages have been received
	if atomic.LoadInt64(&c.msgReceived) == atomic.LoadInt64(&c.msgRequested) {
		// pause flow
		c.consumer.ChangeMaxInFlight(0)
	}

	return true // message was received
}

func (c *Consumer) HandleMessage(msg *gonsq.Message) error {
	msg.DisableAutoResponse()
	body := msg.Body

	// check if the message can be received
	// - we can get messages that were not requested
	// b/c of the delay in turning off the flow of messages.
	if !c.recMsg() {
		// re-queue message so that a ready worker somewhere else
		// can accept it
		msg.Requeue(-1)
		return nil
	}

	// the message should be ready to accept immediately
	// or else the calling application risks that the message
	// reaches the timeout limit and is re-queued.
	select {
	case <-c.ctx.Done():
		// requeue so message is not lost during shutdown
		msg.Requeue(-1)
		return nil
	case c.msgChan <- body:
		msg.Finish() // successful - ack message
		return nil
	}

	return nil
}

// Msg will block until it receives one and only one message.
//
// Safe to call concurrently. Safe to call after Stop() and if this
// is the case will simple return done=true.
func (c *Consumer) Msg() (msg []byte, done bool, err error) {
	if c.ctx.Err() != nil {
		// should not attempt to read if already stopped
		done = true
		return msg, done, nil
	}

	c.reqMsg() // request message

	// wait for a message
	select {
	case msg = <-c.msgChan:
		break
	case <-c.ctx.Done():
		done = true
		break
	}

	return msg, done, err
}

func (c *Consumer) Stop() error {
	if c.ctx.Err() != nil {
		return nil
	}

	c.cncl()

	// stop the consumer
	if c.consumer != nil {
		c.consumer.Stop()
		<-c.consumer.StopChan // wait for consumer to finish closing
	}

	return nil
}
<commit_msg>unlimited nsq requeue attempts (#52)<commit_after>package nsq

import (
	"context"
	"sync"
	"sync/atomic"
	"time"

	gonsq "github.com/bitly/go-nsq"
)

func NewConsumer(topic, channel string, opt *Option) (*Consumer, error) {
	if opt == nil {
		opt = &Option{}
	}

	// initialized at 0 to pause reading on startup
	maxInFlight := 0

	// create nsq consumer config
	nsqConf := gonsq.NewConfig()
	nsqConf.MaxInFlight = maxInFlight
	//nsqConf.MaxAttempts = 65535
	nsqConf.MaxAttempts = 0 // unlimited attempts
	nsqConf.BackoffMultiplier = 0
	nsqConf.MaxBackoffDuration = 0
	nsqConf.DefaultRequeueDelay = time.Second * 2
	nsqConf.MaxRequeueDelay = time.Second * 6

	// create context for clean shutdown
	ctx, cncl := context.WithCancel(context.Background())

	// create lazy consumer
	c := &Consumer{
		opt:       opt,
		nsqConf:   nsqConf,
		closeChan: make(chan interface{}),
		msgChan:   make(chan []byte),
		ctx:       ctx,
		cncl:      cncl,
	}

	// connect to nsqd
	err := c.connect(topic, channel)
	if err != nil {
		return nil, err
	}

	return c, nil
}

type Consumer struct {
	opt      *Option
	nsqConf  *gonsq.Config
	consumer *gonsq.Consumer

	// msgRequested and msgReceived is necessary for lazy
	// loading b/c of the delay in communicating with nsqd
	// the new RDY count.
	msgRequested int64 // total messages requested
	msgReceived  int64 // total messages received

	// bytes are assumed serializable as a valid Task
	// msgChan should be created with a consumer and should never be closed.
	msgChan chan []byte

	// close signal to shut down Consumer. Will stop any in-process messages
	// (messages are re-queued). Will also shutdown any open nsq consumers.
	closeChan chan interface{}

	// wait group is to make sure all consumers are closed during
	// a shutdown.
	sync.WaitGroup

	// mutex for updating consumer maxInFlight
	sync.Mutex

	// context for clean shutdown
	ctx  context.Context
	cncl context.CancelFunc
}

// connect will connect nsq. An error is returned if there
// is a problem connecting.
func (c *Consumer) connect(topic, channel string) error {
	// initialize nsq consumer - does not connect
	consumer, err := gonsq.NewConsumer(topic, channel, c.nsqConf)
	if err != nil {
		return err
	}

	// set custom logger
	if c.opt.Logger != nil {
		consumer.SetLogger(c.opt.Logger, c.opt.LogLvl)
	}

	consumer.AddHandler(c)

	// attempt to connect to lookupds (if provided)
	// or attempt to connect to nsqds (if provided)
	// if neither is provided attempt to connect to localhost
	if len(c.opt.LookupdAddrs) > 0 {
		err = consumer.ConnectToNSQLookupds(c.opt.LookupdAddrs)
		if err != nil {
			return err
		}
	} else if len(c.opt.NSQdAddrs) > 0 {
		err = consumer.ConnectToNSQDs(c.opt.NSQdAddrs)
		if err != nil {
			return err
		}
	} else {
		err = consumer.ConnectToNSQD("localhost:4150")
		if err != nil {
			return err
		}
	}

	// update nsq consumer
	c.consumer = consumer

	return nil
}

// reqMsg will un-pause messages in flight (if paused)
// (by setting maxInflight = 1)
func (c *Consumer) reqMsg() {
	c.Lock()
	// increase total message requested count
	atomic.AddInt64(&c.msgRequested, 1)

	// start flow of messages
	// can always  set the maxInFlight to 1
	// b/c nsq consumer will not update maxInFlight
	// if there is no change.
	c.consumer.ChangeMaxInFlight(1)

	c.Unlock()
}

// recMsg will attempt to 'receive' the message. If the
// message is received then 'true' is returned'. If the
// message could not be received (b/c all requested messages
// have been received) then 'false' is returned.
func (c *Consumer) recMsg() bool {
	c.Lock()
	defer c.Unlock()

	// check if all messages have been received
	if atomic.LoadInt64(&c.msgReceived) == atomic.LoadInt64(&c.msgRequested) {
		return false // all requested messages have been received
	}

	// increase total message received
	atomic.AddInt64(&c.msgReceived, 1)

	// pause flow if all requested messages have been received
	if atomic.LoadInt64(&c.msgReceived) == atomic.LoadInt64(&c.msgRequested) {
		// pause flow
		c.consumer.ChangeMaxInFlight(0)
	}

	return true // message was received
}

func (c *Consumer) HandleMessage(msg *gonsq.Message) error {
	msg.DisableAutoResponse()
	body := msg.Body

	// check if the message can be received
	// - we can get messages that were not requested
	// b/c of the delay in turning off the flow of messages.
	if !c.recMsg() {
		// re-queue message so that a ready worker somewhere else
		// can accept it
		msg.Requeue(-1)
		return nil
	}

	// the message should be ready to accept immediately
	// or else the calling application risks that the message
	// reaches the timeout limit and is re-queued.
	select {
	case <-c.ctx.Done():
		// requeue so message is not lost during shutdown
		msg.Requeue(-1)
		return nil
	case c.msgChan <- body:
		msg.Finish() // successful - ack message
		return nil
	}

	return nil
}

// Msg will block until it receives one and only one message.
//
// Safe to call concurrently. Safe to call after Stop() and if this
// is the case will simple return done=true.
func (c *Consumer) Msg() (msg []byte, done bool, err error) {
	if c.ctx.Err() != nil {
		// should not attempt to read if already stopped
		done = true
		return msg, done, nil
	}

	c.reqMsg() // request message

	// wait for a message
	select {
	case msg = <-c.msgChan:
		break
	case <-c.ctx.Done():
		done = true
		break
	}

	return msg, done, err
}

func (c *Consumer) Stop() error {
	if c.ctx.Err() != nil {
		return nil
	}

	c.cncl()

	// stop the consumer
	if c.consumer != nil {
		c.consumer.Stop()
		<-c.consumer.StopChan // wait for consumer to finish closing
	}

	return nil
}
<|endoftext|>


<commit_before>package treap

import (
	"fmt"
	"testing"
)

func StringLess(p, q interface{}) bool {
	return p.(string) < q.(string)
}

func IntLess(p, q interface{}) bool {
	return p.(int) < q.(int)
}

func TestEmpty(t *testing.T) {
	tree := NewTree(StringLess)
	if tree.Len() != 0 {
		t.Errorf("expected tree len 0")
	}

	x := tree.Get("asdf")
	if x != nil {
		t.Errorf("expected nil for nonexistent key")
	}
}

func TestInsert(t *testing.T) {
	tree := NewTree(StringLess)
	tree.Insert("xyz", "adsf")
	x := tree.Get("xyz")
	if x != "adsf" {
		t.Errorf("expected adsf, got %v", x)
	}
}

func TestFromDoc(t *testing.T) {
	tree := NewTree(IntLess)
	tree.Insert(5, "a")
	tree.Insert(7, "b")
	x := tree.Get(5)
	if x != "a" {
		t.Errorf("expected a, got %v", x)
	}
	x = tree.Get(7)
	if x != "b" {
		t.Errorf("expected b, got %v", x)
	}
	tree.Insert(2, "c")
	x = tree.Get(2)
	if x != "c" {
		t.Errorf("exepcted c, got %v", x)
	}
	tree.Insert(2, "d")
	x = tree.Get(2)
	if x != "d" {
		t.Errorf("exepcted d, got %v", x)
	}
	tree.Delete(5)
	if tree.Exists(5) {
		t.Errorf("expected 5 to be removed from tree")
	}
}

func TestBalance(t *testing.T) {
	tree := NewTree(IntLess)
	for i := 0; i < 1000; i++ {
		tree.Insert(i, false)
	}
	for i := 0; i < 1000; i += 50 {
		fmt.Printf("%d: height = %d\n", i, tree.Height(i))
	}
}

// tests copied from petar's llrb
func TestCases(t *testing.T) {
	tree := NewTree(IntLess)
	tree.Insert(1, true)
	tree.Insert(1, false)
	if tree.Len() != 1 {
		t.Errorf("expecting len 1")
	}
	if !tree.Exists(1) {
		t.Errorf("expecting to find key=1")
	}

	tree.Delete(1)
	if tree.Len() != 0 {
		t.Errorf("expecting len 0")
	}
	if tree.Exists(1) {
		t.Errorf("not expecting to find key=1")
	}

	tree.Delete(1)
	if tree.Len() != 0 {
		t.Errorf("expecting len 0")
	}
	if tree.Exists(1) {
		t.Errorf("not expecting to find key=1")
	}
}

func TestReverseInsertOrder(t *testing.T) {
	tree := NewTree(IntLess)
	n := 100
	for i := 0; i < n; i++ {
		tree.Insert(n-i, true)
	}
	c := tree.IterKeysAscend()
	for j, item := 1, <-c; item != nil; j, item = j+1, <-c {
		if item.(int) != j {
			t.Fatalf("bad order")
		}
	}
}

func TestIterateOverlapNoFunc(t *testing.T) {
	tree := NewTree(IntLess)
	n := 100
	for i := 0; i < n; i++ {
		tree.Insert(n-i, true)
	}
	for v := range tree.IterateOverlap(50) {
		t.Errorf("didn't expect to have any overlap since fn not defined: %v", v)
	}
}

type BucketKey struct {
	Start    int64
	Duration int64
}

type BucketVal struct {
	Start    int64
	Duration int64
	Value    float64
}

func BucketLess(a, b interface{}) bool {
	aa := a.(*BucketKey)
	bb := b.(*BucketKey)
	if aa.Start < bb.Start {
		return true
	}
	if aa.Start == bb.Start {
		return aa.Duration < bb.Duration
	}
	return false
}

func BucketOverlap(a, b interface{}) bool {
	aa := a.(*BucketKey)
	bb := b.(*BucketKey)
	return aa.Start+aa.Duration < bb.Start
}

func TestIterateOverlap(t *testing.T) {
	tree := NewOverlapTree(BucketLess, BucketOverlap)

	tree.Insert(&BucketKey{100, 10}, &BucketVal{100, 10, 5.0})
	tree.Insert(&BucketKey{110, 10}, &BucketVal{110, 10, 6.0})
	tree.Insert(&BucketKey{120, 10}, &BucketVal{120, 10, 7.0})
	tree.Insert(&BucketKey{130, 10}, &BucketVal{130, 10, 8.0})

	for v := range tree.IterateOverlap(&BucketKey{105, 7}) {
		fmt.Printf("val: %v\n", v)
	}
}

/*
before:                            after:

        A                             B
       / \                           / \
      B   C                         D   A
     / \ / \                           / \
    D  E F G                          E   C
                                         / \
                                        F   G

*/
func TestLeftRotate(t *testing.T) {
	// create a tree by hand...
	a := newNode("a", "a", 1)
	b := newNode("b", "b", 2)
	c := newNode("c", "c", 3)
	d := newNode("d", "d", 4)
	e := newNode("e", "e", 5)
	f := newNode("f", "g", 5)
	g := newNode("g", "g", 5)
	a.left = b
	a.right = c
	b.left = d
	b.right = e
	c.left = f
	c.right = g

	x := new(Tree)
	root := x.leftRotate(a)
	if root != b {
		t.Errorf("expected root to be b")
	}
	if root.left != d {
		t.Errorf("expected root.left to be d")
	}
	if root.right != a {
		t.Errorf("expected root.right to be a")
	}
	if a.left != e {
		t.Errorf("expected a.left to be e")
	}
	if a.right != c {
		t.Errorf("expected a.right to be c")
	}
	if c.left != f {
		t.Errorf("expected c.left to be f")
	}
	if c.right != g {
		t.Errorf("expected c.right to be g")
	}
}

/*
before:                            after:

        A                             C
       / \                           / \
      B   C                         A   G
     / \ / \                       / \
    D  E F G                      B   F
                                 / \
                                D  E

*/
func TestRightRotate(t *testing.T) {
	// create a tree by hand...
	a := newNode("a", "a", 1)
	b := newNode("b", "b", 2)
	c := newNode("c", "c", 3)
	d := newNode("d", "d", 4)
	e := newNode("e", "e", 5)
	f := newNode("f", "g", 5)
	g := newNode("g", "g", 5)
	a.left = b
	a.right = c
	b.left = d
	b.right = e
	c.left = f
	c.right = g

	x := new(Tree)
	root := x.rightRotate(a)
	if root != c {
		t.Errorf("expected root to be c")
	}
	if root.left != a {
		t.Errorf("expected root.left to be a")
	}
	if root.right != g {
		t.Errorf("expected root.right to be g")
	}
	if a.left != b {
		t.Errorf("expected a.left to be b")
	}
	if a.right != f {
		t.Errorf("expected a.right to be f")
	}
	if b.left != d {
		t.Errorf("expected b.left to be d")
	}
	if b.right != e {
		t.Errorf("expected b.right to be e")
	}
}
<commit_msg>added benchmark tests from @jeffallen<commit_after>package treap

import (
	"fmt"
        "math/rand"
	"testing"
)

func init() {
	// so that every run is the same seq of rand numbers
	rand.Seed(0)
}

func StringLess(p, q interface{}) bool {
	return p.(string) < q.(string)
}

func IntLess(p, q interface{}) bool {
	return p.(int) < q.(int)
}

func TestEmpty(t *testing.T) {
	tree := NewTree(StringLess)
	if tree.Len() != 0 {
		t.Errorf("expected tree len 0")
	}

	x := tree.Get("asdf")
	if x != nil {
		t.Errorf("expected nil for nonexistent key")
	}
}

func TestInsert(t *testing.T) {
	tree := NewTree(StringLess)
	tree.Insert("xyz", "adsf")
	x := tree.Get("xyz")
	if x != "adsf" {
		t.Errorf("expected adsf, got %v", x)
	}
}

func TestFromDoc(t *testing.T) {
	tree := NewTree(IntLess)
	tree.Insert(5, "a")
	tree.Insert(7, "b")
	x := tree.Get(5)
	if x != "a" {
		t.Errorf("expected a, got %v", x)
	}
	x = tree.Get(7)
	if x != "b" {
		t.Errorf("expected b, got %v", x)
	}
	tree.Insert(2, "c")
	x = tree.Get(2)
	if x != "c" {
		t.Errorf("exepcted c, got %v", x)
	}
	tree.Insert(2, "d")
	x = tree.Get(2)
	if x != "d" {
		t.Errorf("exepcted d, got %v", x)
	}
	tree.Delete(5)
	if tree.Exists(5) {
		t.Errorf("expected 5 to be removed from tree")
	}
}

func TestBalance(t *testing.T) {
	tree := NewTree(IntLess)
	for i := 0; i < 1000; i++ {
		tree.Insert(i, false)
	}
	for i := 0; i < 1000; i += 50 {
		fmt.Printf("%d: height = %d\n", i, tree.Height(i))
	}
}

// tests copied from petar's llrb
func TestCases(t *testing.T) {
	tree := NewTree(IntLess)
	tree.Insert(1, true)
	tree.Insert(1, false)
	if tree.Len() != 1 {
		t.Errorf("expecting len 1")
	}
	if !tree.Exists(1) {
		t.Errorf("expecting to find key=1")
	}

	tree.Delete(1)
	if tree.Len() != 0 {
		t.Errorf("expecting len 0")
	}
	if tree.Exists(1) {
		t.Errorf("not expecting to find key=1")
	}

	tree.Delete(1)
	if tree.Len() != 0 {
		t.Errorf("expecting len 0")
	}
	if tree.Exists(1) {
		t.Errorf("not expecting to find key=1")
	}
}

func TestReverseInsertOrder(t *testing.T) {
	tree := NewTree(IntLess)
	n := 100
	for i := 0; i < n; i++ {
		tree.Insert(n-i, true)
	}
	c := tree.IterKeysAscend()
	for j, item := 1, <-c; item != nil; j, item = j+1, <-c {
		if item.(int) != j {
			t.Fatalf("bad order")
		}
	}
}

func TestIterateOverlapNoFunc(t *testing.T) {
	tree := NewTree(IntLess)
	n := 100
	for i := 0; i < n; i++ {
		tree.Insert(n-i, true)
	}
	for v := range tree.IterateOverlap(50) {
		t.Errorf("didn't expect to have any overlap since fn not defined: %v", v)
	}
}

type BucketKey struct {
	Start    int64
	Duration int64
}

type BucketVal struct {
	Start    int64
	Duration int64
	Value    float64
}

func BucketLess(a, b interface{}) bool {
	aa := a.(*BucketKey)
	bb := b.(*BucketKey)
	if aa.Start < bb.Start {
		return true
	}
	if aa.Start == bb.Start {
		return aa.Duration < bb.Duration
	}
	return false
}

func BucketOverlap(a, b interface{}) bool {
	aa := a.(*BucketKey)
	bb := b.(*BucketKey)
	return aa.Start+aa.Duration < bb.Start
}

func TestIterateOverlap(t *testing.T) {
	tree := NewOverlapTree(BucketLess, BucketOverlap)

	tree.Insert(&BucketKey{100, 10}, &BucketVal{100, 10, 5.0})
	tree.Insert(&BucketKey{110, 10}, &BucketVal{110, 10, 6.0})
	tree.Insert(&BucketKey{120, 10}, &BucketVal{120, 10, 7.0})
	tree.Insert(&BucketKey{130, 10}, &BucketVal{130, 10, 8.0})

	for v := range tree.IterateOverlap(&BucketKey{105, 7}) {
		fmt.Printf("val: %v\n", v)
	}
}

/*
before:                            after:

        A                             B
       / \                           / \
      B   C                         D   A
     / \ / \                           / \
    D  E F G                          E   C
                                         / \
                                        F   G

*/
func TestLeftRotate(t *testing.T) {
	// create a tree by hand...
	a := newNode("a", "a", 1)
	b := newNode("b", "b", 2)
	c := newNode("c", "c", 3)
	d := newNode("d", "d", 4)
	e := newNode("e", "e", 5)
	f := newNode("f", "g", 5)
	g := newNode("g", "g", 5)
	a.left = b
	a.right = c
	b.left = d
	b.right = e
	c.left = f
	c.right = g

	x := new(Tree)
	root := x.leftRotate(a)
	if root != b {
		t.Errorf("expected root to be b")
	}
	if root.left != d {
		t.Errorf("expected root.left to be d")
	}
	if root.right != a {
		t.Errorf("expected root.right to be a")
	}
	if a.left != e {
		t.Errorf("expected a.left to be e")
	}
	if a.right != c {
		t.Errorf("expected a.right to be c")
	}
	if c.left != f {
		t.Errorf("expected c.left to be f")
	}
	if c.right != g {
		t.Errorf("expected c.right to be g")
	}
}

/*
before:                            after:

        A                             C
       / \                           / \
      B   C                         A   G
     / \ / \                       / \
    D  E F G                      B   F
                                 / \
                                D  E

*/
func TestRightRotate(t *testing.T) {
	// create a tree by hand...
	a := newNode("a", "a", 1)
	b := newNode("b", "b", 2)
	c := newNode("c", "c", 3)
	d := newNode("d", "d", 4)
	e := newNode("e", "e", 5)
	f := newNode("f", "g", 5)
	g := newNode("g", "g", 5)
	a.left = b
	a.right = c
	b.left = d
	b.right = e
	c.left = f
	c.right = g

	x := new(Tree)
	root := x.rightRotate(a)
	if root != c {
		t.Errorf("expected root to be c")
	}
	if root.left != a {
		t.Errorf("expected root.left to be a")
	}
	if root.right != g {
		t.Errorf("expected root.right to be g")
	}
	if a.left != b {
		t.Errorf("expected a.left to be b")
	}
	if a.right != f {
		t.Errorf("expected a.right to be f")
	}
	if b.left != d {
		t.Errorf("expected b.left to be d")
	}
	if b.right != e {
		t.Errorf("expected b.right to be e")
	}
}

func treeOfInts(ints []int) (tree *Tree) {
	tree = NewTree(IntLess)
	for _, i := range ints {
		tree.Insert(i, i)
	}
	return
}

func BenchmarkInsert(b *testing.B) {
	b.StopTimer()
	ints := rand.Perm(b.N)
	b.StartTimer()
	_ = treeOfInts(ints)
}

func BenchmarkDelete(b *testing.B) {
	b.StopTimer()
	ints := rand.Perm(b.N)
	tree := treeOfInts(ints)

	b.StartTimer()
	for i := 0; i < b.N; i++ {
		tree.Delete(i)
	}
}

func BenchmarkLookup(b *testing.B) {
	b.StopTimer()
	ints := rand.Perm(b.N)
	tree := treeOfInts(ints)
	b.StartTimer()
	for j := 0; j < 10; j++ {
		for i := 0; i < len(ints)/10; i++ {
			_ = tree.Exists(ints[i])
		}
	}
}
<|endoftext|>


<commit_before>package cfg

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/scrapli/scrapligo/driver/base"

	"github.com/scrapli/scrapligo/util"

	"github.com/scrapli/scrapligo/logging"

	"github.com/scrapli/scrapligo/driver/network"
)

const (
	runningConfig = "running"
)

var ErrNoConfigSourcesProvided = errors.New("no configuration sources provided, cannot continue")
var ErrInvalidConfigTarget = errors.New("provided config source is not valid")

var ErrConnectionNotOpen = errors.New(
	"underlying scrapli connection is not open and `DedicatedConnection` is false, cannot continue")
var ErrVersionError = errors.New("failed getting or parsing device version")

var ErrPrepareNotCalled = errors.New("the Prepare method has not been called, cannot continue")

var ErrConfigSessionAlreadyExists = errors.New(
	"configuration session name already exists, cannot create it")

var ErrInvalidSource = errors.New("invalid config source/target provided")

var ErrFailedToDetermineDeviceState = errors.New(
	"failed to determine device state, ex: failed to determine available space")
var ErrInsufficientSpaceAvailable = errors.New("insufficient space available on device")

var ErrGetConfigFailed = errors.New("get config operation failed")
var ErrLoadConfigFailed = errors.New("load config operation failed")
var ErrAbortConfigFailed = errors.New("abort config operation failed")
var ErrCommitConfigFailed = errors.New("commit config operation failed")
var ErrDiffConfigFailed = errors.New("diff config operation failed")

// Platform -- interface describing the methods the vendor specific platforms must implement, note
// that this is also similar (but not the same!) to the same api surface of the Cfg object that
// users see.
type Platform interface {
	GetVersion() (string, []*base.Response, error)
	GetConfig(source string) (string, []*base.Response, error)
	LoadConfig(config string, replace bool, options *OperationOptions) ([]*base.Response, error)
	AbortConfig() ([]*base.Response, error)
	CommitConfig(source string) ([]*base.Response, error)
	DiffConfig(source, candidateConfig string) ([]*base.Response, string, string, string, error)
	ClearConfigSession()
}

func FormatLogMessage(conn *network.Driver, level, msg string) string {
	return logging.FormatLogMessage(level, conn.Host, conn.Transport.BaseTransportArgs.Port, msg)
}

func setPlatformOptions(p Platform, options ...Option) error {
	for _, option := range options {
		err := option(p)

		if err != nil {
			if errors.Is(err, ErrIgnoredOption) {
				continue
			} else {
				return err
			}
		}
	}

	return nil
}

func parseOperationOptions(o []OperationOption) *OperationOptions {
	opts := &OperationOptions{Source: "running", DiffColorize: true, AutoClean: true}

	if len(o) > 0 && o[0] != nil {
		for _, option := range o {
			option(opts)
		}
	}

	return opts
}

func determineCandidateConfigFilename(c string) string {
	candidateConfigFilename := c

	if c != "" {
		candidateConfigFilename = fmt.Sprintf("scrapli_cfg_%d", time.Now().Unix())
	}

	return candidateConfigFilename
}

// Cfg primary/base cfg platform struct.
type Cfg struct {
	ConfigSources       []string
	OnPrepare           func(*network.Driver) error
	DedicatedConnection bool
	IgnoreVersion       bool

	CandidateConfig string
	VersionString   string
	prepared        bool

	Platform Platform
	conn     *network.Driver
}

// newCfg returns a new instance of Cfg; private because users should be calling the platform
// specific new functions (or using the factory).
func newCfg(
	conn *network.Driver,
	options ...Option,
) (*Cfg, error) {
	c := &Cfg{
		OnPrepare:           nil,
		DedicatedConnection: false,
		IgnoreVersion:       false,
		prepared:            false,
		conn:                conn,
	}

	for _, option := range options {
		err := option(c)

		if err != nil {
			if errors.Is(err, ErrIgnoredOption) {
				continue
			} else {
				return nil, err
			}
		}
	}

	if len(c.ConfigSources) == 0 {
		// if for some reason we dont have config sources we cant really do anything... this should
		// be set by the specific platform so this *shouldn't* happen but... who knows!
		return nil, ErrNoConfigSourcesProvided
	}

	return c, nil
}

func (d *Cfg) invalidConfigSource(err error) error {
	logging.LogError(
		FormatLogMessage(
			d.conn,
			"error",
			"invalid configuration source",
		),
	)

	return err
}

func (d *Cfg) prepareOk() error {
	if d.OnPrepare != nil && !d.prepared {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"OnPrepare provided, but prepare method not called. call prepare method prior "+
					"to using the Cfg object",
			),
		)

		return ErrPrepareNotCalled
	}

	return nil
}

func (d *Cfg) versionOk() error {
	if !d.IgnoreVersion && d.VersionString == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"IgnoreVersion is false, but version has not yet been fetched. call prepare method prior "+
					"to using the Cfg object to ensure version is properly gathered",
			),
		)

		return ErrPrepareNotCalled
	}

	return nil
}

func (d *Cfg) operationOk() error {
	prepareErr := d.prepareOk()
	if prepareErr != nil {
		return prepareErr
	}

	versionErr := d.versionOk()

	if versionErr != nil {
		return versionErr
	}

	return nil
}

func (d *Cfg) validateAndSetVersion(versionResponse *Response) error {
	if versionResponse.Failed != nil {
		logging.LogError(FormatLogMessage(d.conn, "error", "failed getting version from device"))
		return ErrVersionError
	}

	if versionResponse.Result == "" {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed parsing version string from device output"),
		)

		return ErrVersionError
	}

	d.VersionString = versionResponse.Result

	return nil
}

func (d *Cfg) open() error {
	if d.conn.Transport.IsAlive() {
		// nothing to do, connection is already open!
		return nil
	}

	if d.DedicatedConnection {
		err := d.conn.Open()
		return err
	}

	return ErrConnectionNotOpen
}

// Prepare the connection.
func (d *Cfg) Prepare() error {
	logging.LogDebug(FormatLogMessage(d.conn, "info", "preparing cfg connection"))

	err := d.open()
	if err != nil {
		return err
	}

	if !d.IgnoreVersion {
		logging.LogDebug(
			FormatLogMessage(d.conn, "debug", "IgnoreVersion is false, fetching device version"),
		)

		versionResponse, getVersionErr := d.GetVersion()

		if getVersionErr != nil {
			return getVersionErr
		}

		validateVersionErr := d.validateAndSetVersion(versionResponse)
		if validateVersionErr != nil {
			return validateVersionErr
		}
	}

	if d.OnPrepare != nil {
		logging.LogDebug(FormatLogMessage(d.conn, "debug", "OnPrepare provided, executing now"))

		prepareErr := d.OnPrepare(d.conn)
		if prepareErr != nil {
			return prepareErr
		}
	}

	d.prepared = true

	return nil
}

func (d *Cfg) clearConfigSession() {
	logging.LogDebug(
		FormatLogMessage(
			d.conn,
			"debug",
			"resetting config session data",
		),
	)

	d.CandidateConfig = ""
	d.Platform.ClearConfigSession()
}

func (d *Cfg) close() error {
	if d.DedicatedConnection && d.conn.Transport.IsAlive() {
		logging.LogDebug(
			FormatLogMessage(
				d.conn,
				"info",
				"DedicatedConnection is true, closing scrapli connection",
			),
		)

		err := d.conn.Close()

		return err
	}

	return nil
}

// Cleanup cleans up the cfg session.
func (d *Cfg) Cleanup() error {
	err := d.close()
	if err != nil {
		return err
	}

	d.VersionString = ""
	d.prepared = false

	d.clearConfigSession()

	return nil
}

// RenderSubstitutedConfig renders a config with provided substitutions.
func (d *Cfg) RenderSubstitutedConfig() (string, error) {
	return "", nil
}

func (d *Cfg) configSourceValid(source string) bool {
	for _, configSource := range d.ConfigSources {
		if configSource == source {
			return true
		}
	}

	return false
}

// GetVersion gets the version from the device.
func (d *Cfg) GetVersion() (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "get version requested"),
	)

	r := NewResponse(d.conn.Host, "GetVersion", nil)

	versionString, scrapliResponses, err := d.Platform.GetVersion()

	r.Record(scrapliResponses, versionString)

	if r.Failed != nil {
		logging.LogDebug(FormatLogMessage(d.conn, "warning", "failed to fetch device version"))
	}

	if r.Result == "" {
		logging.LogDebug(FormatLogMessage(d.conn, "warning", "failed to parse device version"))
	}

	return r, err
}

// GetConfig gets the configuration of a source datastore from the device.
func (d *Cfg) GetConfig(source string) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(
			d.conn,
			"info",
			fmt.Sprintf("get config requested for config source '%s'", source),
		),
	)

	r := NewResponse(d.conn.Host, "GetConfig", ErrGetConfigFailed)

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	if !d.configSourceValid(source) {
		return r, d.invalidConfigSource(ErrInvalidSource)
	}

	cfgString, scrapliResponses, err := d.Platform.GetConfig(source)

	r.Record(scrapliResponses, cfgString)

	if r.Failed != nil {
		logging.LogError(FormatLogMessage(d.conn, "debug", "failed to fetch config from device"))
	}

	return r, err
}

// LoadConfig loads a candidate configuration.
func (d *Cfg) LoadConfig(
	config string,
	replace bool,
	options ...OperationOption,
) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "load config requested"),
	)

	opts := parseOperationOptions(options)

	d.CandidateConfig = config
	r := NewResponse(d.conn.Host, "LoadConfig", ErrLoadConfigFailed)

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	scrapliResponses, err := d.Platform.LoadConfig(config, replace, opts)

	r.Record(scrapliResponses, "")

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to load candidate configuration"),
		)
	}

	return r, err
}

// LoadConfigFromFile loads a candidate configuration from a provided file.
func (d *Cfg) LoadConfigFromFile(
	f string,
	replace bool,
	options ...OperationOption,
) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "load config from file requested"),
	)

	c, err := util.LoadFileLines(f)
	if err != nil {
		return nil, err
	}

	return d.LoadConfig(strings.Join(c, "\n"), replace, options...)
}

// AbortConfig aborts the loaded candidate configuration.
func (d *Cfg) AbortConfig() (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "abort config requested"),
	)

	r := NewResponse(d.conn.Host, "AbortConfig", ErrAbortConfigFailed)

	if d.CandidateConfig == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"no candidate configuration exists, you must load a config in order to abort it!",
			),
		)

		return r, ErrAbortConfigFailed
	}

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	scrapliResponses, err := d.Platform.AbortConfig()

	r.Record(scrapliResponses, "")

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to abort candidate configuration"),
		)
	}

	d.clearConfigSession()

	return r, err
}

// CommitConfig commits the loaded candidate configuration.
func (d *Cfg) CommitConfig(options ...OperationOption) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "commit config requested"),
	)

	opts := parseOperationOptions(options)

	r := NewResponse(d.conn.Host, "CommitConfig", ErrCommitConfigFailed)

	if d.CandidateConfig == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"no candidate configuration exists, you must load a config in order to commit it!",
			),
		)

		return r, ErrCommitConfigFailed
	}

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	if !d.configSourceValid(opts.Source) {
		return r, d.invalidConfigSource(ErrInvalidSource)
	}

	scrapliResponses, err := d.Platform.CommitConfig(opts.Source)

	r.Record(scrapliResponses, "")

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to commit candidate configuration"),
		)
	}

	d.clearConfigSession()

	return r, err
}

// DiffConfig diffs the candidate configuration against a source config.
func (d *Cfg) DiffConfig(options ...OperationOption) (*DiffResponse, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "diff config requested"),
	)

	opts := parseOperationOptions(options)

	r := NewDiffResponse(d.conn.Host, opts.Source, opts.DiffColorize, opts.DiffSideBySideWidth)

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	if d.CandidateConfig == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"no candidate configuration exists, you must load a config in order to diff it!",
			),
		)

		return r, ErrDiffConfigFailed
	}

	if !d.configSourceValid(opts.Source) {
		return r, d.invalidConfigSource(ErrDiffConfigFailed)
	}

	scrapliResponses, sourceConfig, candidateConfig, deviceDiff, err := d.Platform.DiffConfig(
		opts.Source,
		d.CandidateConfig,
	)

	r.Record(scrapliResponses, "")
	r.RecordDiff(sourceConfig, candidateConfig, deviceDiff)

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to diff configuration"),
		)
	}

	return r, err
}
<commit_msg>fix exported candidate filename check<commit_after>package cfg

import (
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/scrapli/scrapligo/driver/base"

	"github.com/scrapli/scrapligo/util"

	"github.com/scrapli/scrapligo/logging"

	"github.com/scrapli/scrapligo/driver/network"
)

const (
	runningConfig = "running"
)

var ErrNoConfigSourcesProvided = errors.New("no configuration sources provided, cannot continue")
var ErrInvalidConfigTarget = errors.New("provided config source is not valid")

var ErrConnectionNotOpen = errors.New(
	"underlying scrapli connection is not open and `DedicatedConnection` is false, cannot continue")
var ErrVersionError = errors.New("failed getting or parsing device version")

var ErrPrepareNotCalled = errors.New("the Prepare method has not been called, cannot continue")

var ErrConfigSessionAlreadyExists = errors.New(
	"configuration session name already exists, cannot create it")

var ErrInvalidSource = errors.New("invalid config source/target provided")

var ErrFailedToDetermineDeviceState = errors.New(
	"failed to determine device state, ex: failed to determine available space")
var ErrInsufficientSpaceAvailable = errors.New("insufficient space available on device")

var ErrGetConfigFailed = errors.New("get config operation failed")
var ErrLoadConfigFailed = errors.New("load config operation failed")
var ErrAbortConfigFailed = errors.New("abort config operation failed")
var ErrCommitConfigFailed = errors.New("commit config operation failed")
var ErrDiffConfigFailed = errors.New("diff config operation failed")

// Platform -- interface describing the methods the vendor specific platforms must implement, note
// that this is also similar (but not the same!) to the same api surface of the Cfg object that
// users see.
type Platform interface {
	GetVersion() (string, []*base.Response, error)
	GetConfig(source string) (string, []*base.Response, error)
	LoadConfig(config string, replace bool, options *OperationOptions) ([]*base.Response, error)
	AbortConfig() ([]*base.Response, error)
	CommitConfig(source string) ([]*base.Response, error)
	DiffConfig(source, candidateConfig string) ([]*base.Response, string, string, string, error)
	ClearConfigSession()
}

func FormatLogMessage(conn *network.Driver, level, msg string) string {
	return logging.FormatLogMessage(level, conn.Host, conn.Transport.BaseTransportArgs.Port, msg)
}

func setPlatformOptions(p Platform, options ...Option) error {
	for _, option := range options {
		err := option(p)

		if err != nil {
			if errors.Is(err, ErrIgnoredOption) {
				continue
			} else {
				return err
			}
		}
	}

	return nil
}

func parseOperationOptions(o []OperationOption) *OperationOptions {
	opts := &OperationOptions{Source: "running", DiffColorize: true, AutoClean: true}

	if len(o) > 0 && o[0] != nil {
		for _, option := range o {
			option(opts)
		}
	}

	return opts
}

func determineCandidateConfigFilename(c string) string {
	candidateConfigFilename := c

	if c == "" {
		candidateConfigFilename = fmt.Sprintf("scrapli_cfg_%d", time.Now().Unix())
	}

	return candidateConfigFilename
}

// Cfg primary/base cfg platform struct.
type Cfg struct {
	ConfigSources       []string
	OnPrepare           func(*network.Driver) error
	DedicatedConnection bool
	IgnoreVersion       bool

	CandidateConfig string
	VersionString   string
	prepared        bool

	Platform Platform
	conn     *network.Driver
}

// newCfg returns a new instance of Cfg; private because users should be calling the platform
// specific new functions (or using the factory).
func newCfg(
	conn *network.Driver,
	options ...Option,
) (*Cfg, error) {
	c := &Cfg{
		OnPrepare:           nil,
		DedicatedConnection: false,
		IgnoreVersion:       false,
		prepared:            false,
		conn:                conn,
	}

	for _, option := range options {
		err := option(c)

		if err != nil {
			if errors.Is(err, ErrIgnoredOption) {
				continue
			} else {
				return nil, err
			}
		}
	}

	if len(c.ConfigSources) == 0 {
		// if for some reason we dont have config sources we cant really do anything... this should
		// be set by the specific platform so this *shouldn't* happen but... who knows!
		return nil, ErrNoConfigSourcesProvided
	}

	return c, nil
}

func (d *Cfg) invalidConfigSource(err error) error {
	logging.LogError(
		FormatLogMessage(
			d.conn,
			"error",
			"invalid configuration source",
		),
	)

	return err
}

func (d *Cfg) prepareOk() error {
	if d.OnPrepare != nil && !d.prepared {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"OnPrepare provided, but prepare method not called. call prepare method prior "+
					"to using the Cfg object",
			),
		)

		return ErrPrepareNotCalled
	}

	return nil
}

func (d *Cfg) versionOk() error {
	if !d.IgnoreVersion && d.VersionString == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"IgnoreVersion is false, but version has not yet been fetched. call prepare method prior "+
					"to using the Cfg object to ensure version is properly gathered",
			),
		)

		return ErrPrepareNotCalled
	}

	return nil
}

func (d *Cfg) operationOk() error {
	prepareErr := d.prepareOk()
	if prepareErr != nil {
		return prepareErr
	}

	versionErr := d.versionOk()

	if versionErr != nil {
		return versionErr
	}

	return nil
}

func (d *Cfg) validateAndSetVersion(versionResponse *Response) error {
	if versionResponse.Failed != nil {
		logging.LogError(FormatLogMessage(d.conn, "error", "failed getting version from device"))
		return ErrVersionError
	}

	if versionResponse.Result == "" {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed parsing version string from device output"),
		)

		return ErrVersionError
	}

	d.VersionString = versionResponse.Result

	return nil
}

func (d *Cfg) open() error {
	if d.conn.Transport.IsAlive() {
		// nothing to do, connection is already open!
		return nil
	}

	if d.DedicatedConnection {
		err := d.conn.Open()
		return err
	}

	return ErrConnectionNotOpen
}

// Prepare the connection.
func (d *Cfg) Prepare() error {
	logging.LogDebug(FormatLogMessage(d.conn, "info", "preparing cfg connection"))

	err := d.open()
	if err != nil {
		return err
	}

	if !d.IgnoreVersion {
		logging.LogDebug(
			FormatLogMessage(d.conn, "debug", "IgnoreVersion is false, fetching device version"),
		)

		versionResponse, getVersionErr := d.GetVersion()

		if getVersionErr != nil {
			return getVersionErr
		}

		validateVersionErr := d.validateAndSetVersion(versionResponse)
		if validateVersionErr != nil {
			return validateVersionErr
		}
	}

	if d.OnPrepare != nil {
		logging.LogDebug(FormatLogMessage(d.conn, "debug", "OnPrepare provided, executing now"))

		prepareErr := d.OnPrepare(d.conn)
		if prepareErr != nil {
			return prepareErr
		}
	}

	d.prepared = true

	return nil
}

func (d *Cfg) clearConfigSession() {
	logging.LogDebug(
		FormatLogMessage(
			d.conn,
			"debug",
			"resetting config session data",
		),
	)

	d.CandidateConfig = ""
	d.Platform.ClearConfigSession()
}

func (d *Cfg) close() error {
	if d.DedicatedConnection && d.conn.Transport.IsAlive() {
		logging.LogDebug(
			FormatLogMessage(
				d.conn,
				"info",
				"DedicatedConnection is true, closing scrapli connection",
			),
		)

		err := d.conn.Close()

		return err
	}

	return nil
}

// Cleanup cleans up the cfg session.
func (d *Cfg) Cleanup() error {
	err := d.close()
	if err != nil {
		return err
	}

	d.VersionString = ""
	d.prepared = false

	d.clearConfigSession()

	return nil
}

// RenderSubstitutedConfig renders a config with provided substitutions.
func (d *Cfg) RenderSubstitutedConfig() (string, error) {
	return "", nil
}

func (d *Cfg) configSourceValid(source string) bool {
	for _, configSource := range d.ConfigSources {
		if configSource == source {
			return true
		}
	}

	return false
}

// GetVersion gets the version from the device.
func (d *Cfg) GetVersion() (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "get version requested"),
	)

	r := NewResponse(d.conn.Host, "GetVersion", nil)

	versionString, scrapliResponses, err := d.Platform.GetVersion()

	r.Record(scrapliResponses, versionString)

	if r.Failed != nil {
		logging.LogDebug(FormatLogMessage(d.conn, "warning", "failed to fetch device version"))
	}

	if r.Result == "" {
		logging.LogDebug(FormatLogMessage(d.conn, "warning", "failed to parse device version"))
	}

	return r, err
}

// GetConfig gets the configuration of a source datastore from the device.
func (d *Cfg) GetConfig(source string) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(
			d.conn,
			"info",
			fmt.Sprintf("get config requested for config source '%s'", source),
		),
	)

	r := NewResponse(d.conn.Host, "GetConfig", ErrGetConfigFailed)

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	if !d.configSourceValid(source) {
		return r, d.invalidConfigSource(ErrInvalidSource)
	}

	cfgString, scrapliResponses, err := d.Platform.GetConfig(source)

	r.Record(scrapliResponses, cfgString)

	if r.Failed != nil {
		logging.LogError(FormatLogMessage(d.conn, "debug", "failed to fetch config from device"))
	}

	return r, err
}

// LoadConfig loads a candidate configuration.
func (d *Cfg) LoadConfig(
	config string,
	replace bool,
	options ...OperationOption,
) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "load config requested"),
	)

	opts := parseOperationOptions(options)

	d.CandidateConfig = config
	r := NewResponse(d.conn.Host, "LoadConfig", ErrLoadConfigFailed)

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	scrapliResponses, err := d.Platform.LoadConfig(config, replace, opts)

	r.Record(scrapliResponses, "")

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to load candidate configuration"),
		)
	}

	return r, err
}

// LoadConfigFromFile loads a candidate configuration from a provided file.
func (d *Cfg) LoadConfigFromFile(
	f string,
	replace bool,
	options ...OperationOption,
) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "load config from file requested"),
	)

	c, err := util.LoadFileLines(f)
	if err != nil {
		return nil, err
	}

	return d.LoadConfig(strings.Join(c, "\n"), replace, options...)
}

// AbortConfig aborts the loaded candidate configuration.
func (d *Cfg) AbortConfig() (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "abort config requested"),
	)

	r := NewResponse(d.conn.Host, "AbortConfig", ErrAbortConfigFailed)

	if d.CandidateConfig == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"no candidate configuration exists, you must load a config in order to abort it!",
			),
		)

		return r, ErrAbortConfigFailed
	}

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	scrapliResponses, err := d.Platform.AbortConfig()

	r.Record(scrapliResponses, "")

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to abort candidate configuration"),
		)
	}

	d.clearConfigSession()

	return r, err
}

// CommitConfig commits the loaded candidate configuration.
func (d *Cfg) CommitConfig(options ...OperationOption) (*Response, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "commit config requested"),
	)

	opts := parseOperationOptions(options)

	r := NewResponse(d.conn.Host, "CommitConfig", ErrCommitConfigFailed)

	if d.CandidateConfig == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"no candidate configuration exists, you must load a config in order to commit it!",
			),
		)

		return r, ErrCommitConfigFailed
	}

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	if !d.configSourceValid(opts.Source) {
		return r, d.invalidConfigSource(ErrInvalidSource)
	}

	scrapliResponses, err := d.Platform.CommitConfig(opts.Source)

	r.Record(scrapliResponses, "")

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to commit candidate configuration"),
		)
	}

	d.clearConfigSession()

	return r, err
}

// DiffConfig diffs the candidate configuration against a source config.
func (d *Cfg) DiffConfig(options ...OperationOption) (*DiffResponse, error) {
	logging.LogDebug(
		FormatLogMessage(d.conn, "info", "diff config requested"),
	)

	opts := parseOperationOptions(options)

	r := NewDiffResponse(d.conn.Host, opts.Source, opts.DiffColorize, opts.DiffSideBySideWidth)

	operationOkErr := d.operationOk()
	if operationOkErr != nil {
		return r, operationOkErr
	}

	if d.CandidateConfig == "" {
		logging.LogError(
			FormatLogMessage(
				d.conn,
				"error",
				"no candidate configuration exists, you must load a config in order to diff it!",
			),
		)

		return r, ErrDiffConfigFailed
	}

	if !d.configSourceValid(opts.Source) {
		return r, d.invalidConfigSource(ErrDiffConfigFailed)
	}

	scrapliResponses, sourceConfig, candidateConfig, deviceDiff, err := d.Platform.DiffConfig(
		opts.Source,
		d.CandidateConfig,
	)

	r.Record(scrapliResponses, "")
	r.RecordDiff(sourceConfig, candidateConfig, deviceDiff)

	if r.Failed != nil {
		logging.LogError(
			FormatLogMessage(d.conn, "error", "failed to diff configuration"),
		)
	}

	return r, err
}
<|endoftext|>


<commit_before>package memory

import (
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"strings"
)

// the collision avoidance window defines how many times we try to find a higher
// slot that's free if two record hashes collide
var collisionAvoidanceWindow = uint32(1024)

// the function we use to get the hash for hashing the meta records
// it can be replaced for mocking in tests
var queryHash func() hash.Hash32

func init() {
	queryHash = fnv.New32a
}

// list of meta records keyed by a unique identifier used as ID
type metaTagRecords map[recordId]metaTagRecord

type metaTagRecord struct {
	metaTags []kv
	queries  []expression
}

type recordId uint32

// index structure keyed by tag -> value -> list of meta record IDs
type metaTagValue map[string][]recordId
type metaTagIndex map[string]metaTagValue

func (m metaTagIndex) deleteRecord(keyValue kv, id recordId) {
	if values, ok := m[keyValue.key]; ok {
		if ids, ok := values[keyValue.value]; ok {
			for i := 0; i < len(ids); i++ {
				if ids[i] == id {
					// no need to keep the order
					ids[i] = ids[len(ids)-1]
					values[keyValue.value] = ids[:len(ids)-1]

					// no id should ever be present more than once
					return
				}
			}
		}
	}
}

func (m metaTagIndex) insertRecord(keyValue kv, id recordId) {
	var values metaTagValue
	var ok bool

	if values, ok = m[keyValue.key]; !ok {
		values = make(metaTagValue)
		m[keyValue.key] = values
	}

	values[keyValue.value] = append(values[keyValue.value], id)
}

// upsert inserts or updates a meta tag record according to the given specifications
// it uses the set of tag query expressions as the identity of the record, if a record with the
// same identity is already present then its meta tags get updated to the specified ones.
// If the new record contains no meta tags, then the update is equivalent to a delete.
// Those are the return values:
// 1) The id at which the new record got inserted
// 2) Pointer to the inserted metaTagRecord
// 3) The id of the record that has been replaced if an update was performed
// 4) Pointer to the metaTagRecord that has been replaced if an update was performed, otherwise nil
// 5) Error if an error occurred, otherwise it's nil
func (m metaTagRecords) upsert(metaTags []string, tagQueryExpressions []string) (recordId, *metaTagRecord, recordId, *metaTagRecord, error) {
	record, err := metaTagRecordFromStrings(metaTags, tagQueryExpressions)
	if err != nil {
		return 0, nil, 0, nil, err
	}

	record.sortQueries()
	id := record.hashQueries()
	var oldRecord *metaTagRecord
	var oldId recordId

	// loop over existing records, starting from id, trying to find one that has
	// the exact same queries as the one we're upserting
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		if existingRecord, ok := m[id+recordId(i)]; ok {
			if record.matchesQueries(existingRecord) {
				oldRecord = &existingRecord
				oldId = id + recordId(i)
				delete(m, oldId)
				break
			}
		}
	}

	if !record.hasMetaTags() {
		return 0, &record, oldId, oldRecord, nil
	}

	// now find the best position to insert the new/updated record, starting from id
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		// if we find a free slot, then insert the new record there
		if _, ok := m[id]; !ok {
			m[id] = record

			return id, &record, oldId, oldRecord, nil
		}

		id++
	}

	return 0, nil, 0, nil, fmt.Errorf("MetaTagRecordUpsert: Unable to find a slot to insert record")
}

// metaTagRecordFromStrings takes two slices of strings, parses them and returns a metaTagRecord
// The first slice of strings has the meta tags & values
// The second slice has the tag query expressions which the meta tags & values refer to
// On parsing error the second returned value is an error, otherwise it is nil
func metaTagRecordFromStrings(metaTags []string, tagQueryExpressions []string) (metaTagRecord, error) {
	record := metaTagRecord{
		metaTags: make([]kv, 0, len(metaTags)),
		queries:  make([]expression, 0, len(tagQueryExpressions)),
	}

	if len(tagQueryExpressions) == 0 {
		return record, fmt.Errorf("Requiring at least one tag query expression, 0 given")
	}

	for _, tag := range metaTags {
		tagSplits := strings.SplitN(tag, "=", 2)
		if len(tagSplits) < 2 {
			return record, fmt.Errorf("Missing \"=\" sign in tag %s", tag)
		}
		if len(tagSplits[0]) == 0 || len(tagSplits[1]) == 0 {
			return record, fmt.Errorf("Tag/Value cannot be empty in %s", tag)
		}

		record.metaTags = append(record.metaTags, kv{key: tagSplits[0], value: tagSplits[1]})
	}

	for _, query := range tagQueryExpressions {
		parsed, err := parseExpression(query)
		if err != nil {
			return record, err
		}
		record.queries = append(record.queries, parsed)
	}

	return record, nil
}

func (m *metaTagRecord) metaTagStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.metaTags))

	for i, tag := range m.metaTags {
		tag.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

func (m *metaTagRecord) queryStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.queries))

	for i, query := range m.queries {
		query.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

// hashQueries generates a hash of all the queries in the record
func (m *metaTagRecord) hashQueries() recordId {
	builder := strings.Builder{}
	for _, query := range m.queries {
		query.stringIntoBuilder(&builder)

		// trailing ";" doesn't matter, this is only hash input
		builder.WriteString(";")
	}

	h := queryHash()
	h.Write([]byte(builder.String()))
	return recordId(h.Sum32())
}

// sortQueries sorts all the queries first by key, then by value, then by
// operator. The order itself is not relevant, we only need to guarantee
// that the order is consistent for a given set of queries after every call
// to this function, because the queries will then be used as hash input
func (m *metaTagRecord) sortQueries() {
	sort.Slice(m.queries, func(i, j int) bool {
		if m.queries[i].key == m.queries[j].key {
			if m.queries[i].value == m.queries[j].value {
				return m.queries[i].operator < m.queries[j].operator
			}
			return m.queries[i].value < m.queries[j].value
		}
		return m.queries[i].key < m.queries[j].key
	})
}

// matchesQueries compares another tag record's queries to this
// one's queries. Returns true if they are equal, otherwise false.
// It is assumed that all the queries are already sorted
func (m *metaTagRecord) matchesQueries(other metaTagRecord) bool {
	if len(m.queries) != len(other.queries) {
		return false
	}

	for i, query := range m.queries {
		if query.key != other.queries[i].key {
			return false
		}

		if query.operator != other.queries[i].operator {
			return false
		}

		if query.value != other.queries[i].value {
			return false
		}
	}

	return true
}

// hasMetaTags returns true if the meta tag record has one or more
// meta tags, otherwise it returns false
func (m *metaTagRecord) hasMetaTags() bool {
	return len(m.metaTags) > 0
}
<commit_msg>rearrange order of types/methods<commit_after>package memory

import (
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"strings"
)

// the collision avoidance window defines how many times we try to find a higher
// slot that's free if two record hashes collide
var collisionAvoidanceWindow = uint32(1024)

// the function we use to get the hash for hashing the meta records
// it can be replaced for mocking in tests
var queryHash func() hash.Hash32

func init() {
	queryHash = fnv.New32a
}

// list of meta records keyed by a unique identifier used as ID
type metaTagRecords map[recordId]metaTagRecord

type recordId uint32

// upsert inserts or updates a meta tag record according to the given specifications
// it uses the set of tag query expressions as the identity of the record, if a record with the
// same identity is already present then its meta tags get updated to the specified ones.
// If the new record contains no meta tags, then the update is equivalent to a delete.
// Those are the return values:
// 1) The id at which the new record got inserted
// 2) Pointer to the inserted metaTagRecord
// 3) The id of the record that has been replaced if an update was performed
// 4) Pointer to the metaTagRecord that has been replaced if an update was performed, otherwise nil
// 5) Error if an error occurred, otherwise it's nil
func (m metaTagRecords) upsert(metaTags []string, tagQueryExpressions []string) (recordId, *metaTagRecord, recordId, *metaTagRecord, error) {
	record, err := metaTagRecordFromStrings(metaTags, tagQueryExpressions)
	if err != nil {
		return 0, nil, 0, nil, err
	}

	record.sortQueries()
	id := record.hashQueries()
	var oldRecord *metaTagRecord
	var oldId recordId

	// loop over existing records, starting from id, trying to find one that has
	// the exact same queries as the one we're upserting
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		if existingRecord, ok := m[id+recordId(i)]; ok {
			if record.matchesQueries(existingRecord) {
				oldRecord = &existingRecord
				oldId = id + recordId(i)
				delete(m, oldId)
				break
			}
		}
	}

	if !record.hasMetaTags() {
		return 0, &record, oldId, oldRecord, nil
	}

	// now find the best position to insert the new/updated record, starting from id
	for i := uint32(0); i < collisionAvoidanceWindow; i++ {
		// if we find a free slot, then insert the new record there
		if _, ok := m[id]; !ok {
			m[id] = record

			return id, &record, oldId, oldRecord, nil
		}

		id++
	}

	return 0, nil, 0, nil, fmt.Errorf("MetaTagRecordUpsert: Unable to find a slot to insert record")
}

type metaTagRecord struct {
	metaTags []kv
	queries  []expression
}

// metaTagRecordFromStrings takes two slices of strings, parses them and returns a metaTagRecord
// The first slice of strings has the meta tags & values
// The second slice has the tag query expressions which the meta tags & values refer to
// On parsing error the second returned value is an error, otherwise it is nil
func metaTagRecordFromStrings(metaTags []string, tagQueryExpressions []string) (metaTagRecord, error) {
	record := metaTagRecord{
		metaTags: make([]kv, 0, len(metaTags)),
		queries:  make([]expression, 0, len(tagQueryExpressions)),
	}

	if len(tagQueryExpressions) == 0 {
		return record, fmt.Errorf("Requiring at least one tag query expression, 0 given")
	}

	for _, tag := range metaTags {
		tagSplits := strings.SplitN(tag, "=", 2)
		if len(tagSplits) < 2 {
			return record, fmt.Errorf("Missing \"=\" sign in tag %s", tag)
		}
		if len(tagSplits[0]) == 0 || len(tagSplits[1]) == 0 {
			return record, fmt.Errorf("Tag/Value cannot be empty in %s", tag)
		}

		record.metaTags = append(record.metaTags, kv{key: tagSplits[0], value: tagSplits[1]})
	}

	for _, query := range tagQueryExpressions {
		parsed, err := parseExpression(query)
		if err != nil {
			return record, err
		}
		record.queries = append(record.queries, parsed)
	}

	return record, nil
}

func (m *metaTagRecord) metaTagStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.metaTags))

	for i, tag := range m.metaTags {
		tag.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

func (m *metaTagRecord) queryStrings(builder *strings.Builder) []string {
	res := make([]string, len(m.queries))

	for i, query := range m.queries {
		query.stringIntoBuilder(builder)
		res[i] = builder.String()
		builder.Reset()
	}

	return res
}

// hashQueries generates a hash of all the queries in the record
func (m *metaTagRecord) hashQueries() recordId {
	builder := strings.Builder{}
	for _, query := range m.queries {
		query.stringIntoBuilder(&builder)

		// trailing ";" doesn't matter, this is only hash input
		builder.WriteString(";")
	}

	h := queryHash()
	h.Write([]byte(builder.String()))
	return recordId(h.Sum32())
}

// sortQueries sorts all the queries first by key, then by value, then by
// operator. The order itself is not relevant, we only need to guarantee
// that the order is consistent for a given set of queries after every call
// to this function, because the queries will then be used as hash input
func (m *metaTagRecord) sortQueries() {
	sort.Slice(m.queries, func(i, j int) bool {
		if m.queries[i].key == m.queries[j].key {
			if m.queries[i].value == m.queries[j].value {
				return m.queries[i].operator < m.queries[j].operator
			}
			return m.queries[i].value < m.queries[j].value
		}
		return m.queries[i].key < m.queries[j].key
	})
}

// matchesQueries compares another tag record's queries to this
// one's queries. Returns true if they are equal, otherwise false.
// It is assumed that all the queries are already sorted
func (m *metaTagRecord) matchesQueries(other metaTagRecord) bool {
	if len(m.queries) != len(other.queries) {
		return false
	}

	for i, query := range m.queries {
		if query.key != other.queries[i].key {
			return false
		}

		if query.operator != other.queries[i].operator {
			return false
		}

		if query.value != other.queries[i].value {
			return false
		}
	}

	return true
}

// hasMetaTags returns true if the meta tag record has one or more
// meta tags, otherwise it returns false
func (m *metaTagRecord) hasMetaTags() bool {
	return len(m.metaTags) > 0
}

// index structure keyed by tag -> value -> list of meta record IDs
type metaTagValue map[string][]recordId
type metaTagIndex map[string]metaTagValue

func (m metaTagIndex) deleteRecord(keyValue kv, id recordId) {
	if values, ok := m[keyValue.key]; ok {
		if ids, ok := values[keyValue.value]; ok {
			for i := 0; i < len(ids); i++ {
				if ids[i] == id {
					// no need to keep the order
					ids[i] = ids[len(ids)-1]
					values[keyValue.value] = ids[:len(ids)-1]

					// no id should ever be present more than once
					return
				}
			}
		}
	}
}

func (m metaTagIndex) insertRecord(keyValue kv, id recordId) {
	var values metaTagValue
	var ok bool

	if values, ok = m[keyValue.key]; !ok {
		values = make(metaTagValue)
		m[keyValue.key] = values
	}

	values[keyValue.value] = append(values[keyValue.value], id)
}
<|endoftext|>


<commit_before>package kafkamdm

import (
	"flag"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"

	"github.com/bsm/sarama-cluster"
	"github.com/raintank/met"
	"github.com/raintank/metrictank/idx"
	"github.com/raintank/metrictank/in"
	"github.com/raintank/metrictank/mdata"
	"github.com/raintank/metrictank/usage"
)

type KafkaMdm struct {
	in.In
	consumer *cluster.Consumer
	stats    met.Backend

	wg sync.WaitGroup
	// read from this channel to block until consumer is cleanly stopped
	StopChan chan int
}

var LogLevel int
var Enabled bool
var brokerStr string
var brokers []string
var topicStr string
var topics []string
var group string
var config *cluster.Config
var channelBufferSize int
var consumerFetchMin int
var consumerFetchDefault int
var consumerMaxWaitTime string
var consumerMaxProcessingTime string
var netMaxOpenRequests int

func ConfigSetup() {
	inKafkaMdm := flag.NewFlagSet("kafka-mdm-in", flag.ExitOnError)
	inKafkaMdm.BoolVar(&Enabled, "enabled", false, "")
	inKafkaMdm.StringVar(&brokerStr, "brokers", "kafka:9092", "tcp address for kafka (may be be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&topicStr, "topics", "mdm", "kafka topic (may be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&group, "group", "group1", "kafka consumer group")
	inKafkaMdm.IntVar(&channelBufferSize, "channel-buffer-size", 1000000, "The number of metrics to buffer in internal and external channels")
	inKafkaMdm.IntVar(&consumerFetchMin, "consumer-fetch-min", 1024000, "The minimum number of message bytes to fetch in a request")
	inKafkaMdm.IntVar(&consumerFetchDefault, "consumer-fetch-default", 4096000, "The default number of message bytes to fetch in a request")
	inKafkaMdm.StringVar(&consumerMaxWaitTime, "consumer-max-wait-time", "1s", "The maximum amount of time the broker will wait for Consumer.Fetch.Min bytes to become available before it returns fewer than that anyway")
	inKafkaMdm.StringVar(&consumerMaxProcessingTime, "consumer-max-processing-time", "1s", "The maximum amount of time the consumer expects a message takes to process")
	inKafkaMdm.IntVar(&netMaxOpenRequests, "net-max-open-requests", 100, "How many outstanding requests a connection is allowed to have before sending on it blocks")
	globalconf.Register("kafka-mdm-in", inKafkaMdm)
}

func ConfigProcess(instance string) {
	if !Enabled {
		return
	}

	waitTime, err := time.ParseDuration(consumerMaxWaitTime)
	if err != nil {
		log.Fatal(4, "kafka-mdm invalid config, could not parse consumer-max-wait-time: %s", err)
	}
	processingTime, err := time.ParseDuration(consumerMaxProcessingTime)
	if err != nil {
		log.Fatal(4, "kafka-mdm invalid config, could not parse consumer-max-processing-time: %s", err)
	}

	brokers = strings.Split(brokerStr, ",")
	topics = strings.Split(topicStr, ",")

	config = cluster.NewConfig()
	// see https://github.com/raintank/metrictank/issues/236
	config.Consumer.Offsets.Initial = sarama.OffsetNewest
	config.ClientID = instance + "-mdm"
	config.Group.Return.Notifications = true
	config.ChannelBufferSize = channelBufferSize
	config.Consumer.Fetch.Min = int32(consumerFetchMin)
	config.Consumer.Fetch.Default = int32(consumerFetchDefault)
	config.Consumer.MaxWaitTime = waitTime
	config.Consumer.MaxProcessingTime = processingTime
	config.Net.MaxOpenRequests = netMaxOpenRequests
	config.Config.Version = sarama.V0_10_0_0
	err = config.Validate()
	if err != nil {
		log.Fatal(2, "kafka-mdm invalid config: %s", err)
	}
}

func New(stats met.Backend) *KafkaMdm {
	consumer, err := cluster.NewConsumer(brokers, group, topics, config)
	if err != nil {
		log.Fatal(2, "kafka-mdm failed to start consumer: %s", err)
	}
	log.Info("kafka-mdm consumer started without error")
	k := KafkaMdm{
		consumer: consumer,
		stats:    stats,
		StopChan: make(chan int),
	}

	return &k
}

func (k *KafkaMdm) Start(metrics mdata.Metrics, metricIndex idx.MetricIndex, usg *usage.Usage) {
	k.In = in.New(metrics, metricIndex, usg, "kafka-mdm", k.stats)
	go k.notifications()
	go k.consume()
}

func (k *KafkaMdm) consume() {
	k.wg.Add(1)
	messageChan := k.consumer.Messages()
	for msg := range messageChan {
		if LogLevel < 2 {
			log.Debug("kafka-mdm received message: Topic %s, Partition: %d, Offset: %d, Key: %x", msg.Topic, msg.Partition, msg.Offset, msg.Key)
		}
		k.In.Handle(msg.Value)
		k.consumer.MarkOffset(msg, "")
	}
	log.Info("kafka-mdm consumer ended.")
	k.wg.Done()
}

func (k *KafkaMdm) notifications() {
	k.wg.Add(1)
	for msg := range k.consumer.Notifications() {
		if len(msg.Claimed) > 0 {
			for topic, partitions := range msg.Claimed {
				log.Info("kafka-mdm consumer claimed %d partitions on topic: %s", len(partitions), topic)
			}
		}
		if len(msg.Released) > 0 {
			for topic, partitions := range msg.Released {
				log.Info("kafka-mdm consumer released %d partitions on topic: %s", len(partitions), topic)
			}
		}

		if len(msg.Current) == 0 {
			log.Info("kafka-mdm consumer is no longer consuming from any partitions.")
		} else {
			log.Info("kafka-mdm Current partitions:")
			for topic, partitions := range msg.Current {
				log.Info("kafka-mdm Current partitions: %s: %v", topic, partitions)
			}
		}
	}
	log.Info("kafka-mdm notification processing stopped")
	k.wg.Done()
}

// Stop will initiate a graceful stop of the Consumer (permanent)
//
// NOTE: receive on StopChan to block until this process completes
func (k *KafkaMdm) Stop() {
	// closes notifications and messages channels, amongst others
	k.consumer.Close()

	go func() {
		k.wg.Wait()
		close(k.StopChan)
	}()
}
<commit_msg>refactor kafkaMdm to manage its own offsets.<commit_after>package kafkamdm

import (
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"

	"github.com/raintank/met"
	"github.com/raintank/metrictank/idx"
	"github.com/raintank/metrictank/in"
	"github.com/raintank/metrictank/mdata"
	"github.com/raintank/metrictank/usage"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/storage"
)

type KafkaMdm struct {
	in.In
	consumer sarama.Consumer
	client   sarama.Client
	stats    met.Backend

	wg sync.WaitGroup
	// read from this channel to block until consumer is cleanly stopped
	StopChan chan int

	// signal to PartitionConsumers to shutdown
	stopConsuming chan struct{}
}

var LogLevel int
var Enabled bool
var brokerStr string
var brokers []string
var topicStr string
var topics []string
var offset string
var dataDir string
var config *sarama.Config
var channelBufferSize int
var consumerFetchMin int
var consumerFetchDefault int
var consumerMaxWaitTime time.Duration
var consumerMaxProcessingTime time.Duration
var netMaxOpenRequests int
var offsetMgr *OffsetMgr
var offsetCommitInterval time.Duration

func ConfigSetup() {
	inKafkaMdm := flag.NewFlagSet("kafka-mdm-in", flag.ExitOnError)
	inKafkaMdm.BoolVar(&Enabled, "enabled", false, "")
	inKafkaMdm.StringVar(&brokerStr, "brokers", "kafka:9092", "tcp address for kafka (may be be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&topicStr, "topics", "mdm", "kafka topic (may be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&offset, "offset", "last", "Set the offset to start consuming from. Can be one of newest, oldest,last or a time duration")
	inKafkaMdm.DurationVar(&offsetCommitInterval, "offset-commit-interval", time.Second*5, "Interval at which offsets should be saved.")
	inKafkaMdm.StringVar(&dataDir, "data-dir", "", "Directory to store partition offsets index")
	inKafkaMdm.IntVar(&channelBufferSize, "channel-buffer-size", 1000000, "The number of metrics to buffer in internal and external channels")
	inKafkaMdm.IntVar(&consumerFetchMin, "consumer-fetch-min", 1024000, "The minimum number of message bytes to fetch in a request")
	inKafkaMdm.IntVar(&consumerFetchDefault, "consumer-fetch-default", 4096000, "The default number of message bytes to fetch in a request")
	inKafkaMdm.DurationVar(&consumerMaxWaitTime, "consumer-max-wait-time", time.Second, "The maximum amount of time the broker will wait for Consumer.Fetch.Min bytes to become available before it returns fewer than that anyway")
	inKafkaMdm.DurationVar(&consumerMaxProcessingTime, "consumer-max-processing-time", time.Second, "The maximum amount of time the consumer expects a message takes to process")
	inKafkaMdm.IntVar(&netMaxOpenRequests, "net-max-open-requests", 100, "How many outstanding requests a connection is allowed to have before sending on it blocks")
	globalconf.Register("kafka-mdm-in", inKafkaMdm)
}

func ConfigProcess(instance string) {
	if !Enabled {
		return
	}
	switch offset {
	case "last":
	case "oldest":
	case "newest":
	default:
		_, err := time.ParseDuration(offest)
		if err != nil {
			log.Fatal(4, "kafkamdm: invalid offest format. %s", err)
		}
	}

	var err error
	offsetMgr, err = NewOffsetMgr(dataDir)
	if err != nil {
		log.Fatal(4, "kafka-mdm couldnt create offsetMgr. %s", err)
	}
	brokers = strings.Split(brokerStr, ",")
	topics = strings.Split(topicStr, ",")

	config = sarama.NewConfig()

	config.ClientID = instance + "-mdm"
	config.ChannelBufferSize = channelBufferSize
	config.Consumer.Fetch.Min = int32(consumerFetchMin)
	config.Consumer.Fetch.Default = int32(consumerFetchDefault)
	config.Consumer.MaxWaitTime = consumerMaxWaitTime
	config.Consumer.MaxProcessingTime = consumerMaxProcessingTime
	config.Net.MaxOpenRequests = netMaxOpenRequests
	config.Version = sarama.V0_10_0_0
	err = config.Validate()
	if err != nil {
		log.Fatal(2, "kafka-mdm invalid config: %s", err)
	}
}

func New(stats met.Backend) *KafkaMdm {
	client, err := sarama.NewClient(brokers, config)
	if err != nil {
		log.Fatal(4, "kafka-mdm failed to create client. %s", err)
	}
	consumer, err := sarama.NewConsumerFromClient(client)
	if err != nil {
		log.Fatal(2, "kafka-mdm failed to create consumer: %s", err)
	}
	log.Info("kafka-mdm consumer created without error")
	k := KafkaMdm{
		consumer:      consumer,
		client:        client,
		stats:         stats,
		StopChan:      make(chan int),
		stopConsuming: make(chan struct{}),
	}

	return &k
}

func (k *KafkaMdm) Start(metrics mdata.Metrics, metricIndex idx.MetricIndex, usg *usage.Usage) {
	k.In = in.New(metrics, metricIndex, usg, "kafka-mdm", k.stats)
	for _, topic := range topics {
		// get partitions.
		partitions, err := k.consumer.Partitions(topic)
		if err != nil {
			log.Fatal(4, "kafka-mdm: Faild to get partitions for topic %s. %s", topic, err)
		}
		for _, partition := range partitions {
			switch offset {
			case "oldest":
				go k.consumePartition(topic, partition, -2)
			case "newest":
				go k.consumePartition(topic, partition, -1)
			case "last":
				o, err := offsetMgr.Last(topic, partition)
				if err != nil {
					log.Fatal(4, "kafka-mdm: Failed to get offset for %s:%d. %s", topic, partition, err)
				}
				go k.consumePartition(topic, partition, o)
			default:
				// treat offset as a time duration.
				dur, err := time.ParseDuration(offset)
				if err != nil {
					log.Fatal(4, "kafka-mdm: Failed to parse offset to time.Duration. %s", err)
				}
				o, err := k.client.GetOffset(topic, partition, time.Now().Add(-1*dur).UnixNano()/int64(time.Millisecond))
				if err != nil {
					log.Fatal(4, "kafka-mdm: failed to get offset for %s:%d.  %s", topic, partition, err)
				}
				go k.consumePartition(topic, partition, o)
			}

		}
	}
}

// this will continually consume from the topic until k.stopConsuming is triggered.
func (k *KafkaMdm) consumePartition(topic string, partition int32, partitionOffset int64) {
	k.wg.Add(1)
	defer k.wg.Done()

	pc, err := k.consumer.ConsumePartition(topic, partition, partitionOffset)
	if err != nil {
		log.Fatal(4, "kafka-mdm: failed to start partitionConsumer for %s:%d. %s", topic, partition, err)
	}
	log.Info("kafka-mdm: consuming from %s:%d from offset %d", topic, partition, partitionOffset)
	currentOffset := partitionOffset
	messages := pc.Messages()
	ticker := time.NewTicker(time.Second * 5)
	for {
		select {
		case msg := <-messages:
			if LogLevel < 2 {
				log.Debug("kafka-mdm received message: Topic %s, Partition: %d, Offset: %d, Key: %x", msg.Topic, msg.Partition, msg.Offset, msg.Key)
			}
			k.In.Handle(msg.Value)
			currentOffset = msg.Offset
		case <-ticker.C:
			if err := offsetMgr.Commit(topic, partition, currentOffset); err != nil {
				log.Error(3, "kafka-mdm failed to commit offset for %s:%d, %s", topic, partition, err)
			}
		case <-k.stopConsuming:
			pc.Close()
			offsetMgr.Commit(topic, partition, currentOffset)
			log.Info("kafka-mdm consumer for %s:%d ended.", topic, partition)
			return
		}
	}
}

// Stop will initiate a graceful stop of the Consumer (permanent)
//
// NOTE: receive on StopChan to block until this process completes
func (k *KafkaMdm) Stop() {
	// closes notifications and messages channels, amongst others
	close(k.stopConsuming)
	go func() {
		k.wg.Wait()
		offsetMgr.Close()
		close(k.StopChan)
	}()
}

type OffsetMgr struct {
	db *leveldb.DB
}

func NewOffsetMgr(dir string) (*OffsetMgr, error) {
	dbFile := filepath.Join(dir, "partitionOffsets.db")
	db, err := leveldb.OpenFile(dbFile, &opt.Options{})
	if err != nil {
		if _, ok := err.(*storage.ErrCorrupted); ok {
			db, err = leveldb.RecoverFile(dbFile, &opt.Options{})
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	return &OffsetMgr{
		db: db,
	}, nil
}

func (o *OffsetMgr) Close() {
	o.Close()
}

func (o *OffsetMgr) Commit(topic string, partition int32, offset int64) error {
	key := new(bytes.Buffer)
	key.WriteString(fmt.Sprintf("T:%s-P:%d", topic, partition))
	data := new(bytes.Buffer)
	if err := binary.Write(data, binary.LittleEndian, offset); err != nil {
		return err
	}
	return o.db.Put(key.Bytes(), data.Bytes(), &opt.WriteOptions{Sync: true})
}

func (o *OffsetMgr) Last(topic string, partition int32) (int64, error) {
	key := new(bytes.Buffer)
	key.WriteString(fmt.Sprintf("T:%s-P:%d", topic, partition))
	data, err := o.db.Get(key.Bytes(), nil)
	if err != nil {
		if err == leveldb.ErrNotFound {
			return -1, nil
		}
		return 0, err
	}
	var offset int64
	err = binary.Read(bytes.NewBuffer(data), binary.LittleEndian, &offset)
	if err != nil {
		return 0, err
	}
	return offset, nil
}
<|endoftext|>


<commit_before>package kafkamdm

import (
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"

	"github.com/raintank/met"
	"github.com/raintank/metrictank/idx"
	"github.com/raintank/metrictank/in"
	"github.com/raintank/metrictank/mdata"
	"github.com/raintank/metrictank/usage"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/storage"
)

type KafkaMdm struct {
	in.In
	consumer sarama.Consumer
	client   sarama.Client
	stats    met.Backend

	wg sync.WaitGroup
	// read from this channel to block until consumer is cleanly stopped
	StopChan chan int

	// signal to PartitionConsumers to shutdown
	stopConsuming chan struct{}
}

var LogLevel int
var Enabled bool
var brokerStr string
var brokers []string
var topicStr string
var topics []string
var offset string
var dataDir string
var config *sarama.Config
var channelBufferSize int
var consumerFetchMin int
var consumerFetchDefault int
var consumerMaxWaitTime time.Duration
var consumerMaxProcessingTime time.Duration
var netMaxOpenRequests int
var offsetMgr *OffsetMgr
var offsetCommitInterval time.Duration

func ConfigSetup() {
	inKafkaMdm := flag.NewFlagSet("kafka-mdm-in", flag.ExitOnError)
	inKafkaMdm.BoolVar(&Enabled, "enabled", false, "")
	inKafkaMdm.StringVar(&brokerStr, "brokers", "kafka:9092", "tcp address for kafka (may be be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&topicStr, "topics", "mdm", "kafka topic (may be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&offset, "offset", "last", "Set the offset to start consuming from. Can be one of newest, oldest,last or a time duration")
	inKafkaMdm.DurationVar(&offsetCommitInterval, "offset-commit-interval", time.Second*5, "Interval at which offsets should be saved.")
	inKafkaMdm.StringVar(&dataDir, "data-dir", "", "Directory to store partition offsets index")
	inKafkaMdm.IntVar(&channelBufferSize, "channel-buffer-size", 1000000, "The number of metrics to buffer in internal and external channels")
	inKafkaMdm.IntVar(&consumerFetchMin, "consumer-fetch-min", 1024000, "The minimum number of message bytes to fetch in a request")
	inKafkaMdm.IntVar(&consumerFetchDefault, "consumer-fetch-default", 4096000, "The default number of message bytes to fetch in a request")
	inKafkaMdm.DurationVar(&consumerMaxWaitTime, "consumer-max-wait-time", time.Second, "The maximum amount of time the broker will wait for Consumer.Fetch.Min bytes to become available before it returns fewer than that anyway")
	inKafkaMdm.DurationVar(&consumerMaxProcessingTime, "consumer-max-processing-time", time.Second, "The maximum amount of time the consumer expects a message takes to process")
	inKafkaMdm.IntVar(&netMaxOpenRequests, "net-max-open-requests", 100, "How many outstanding requests a connection is allowed to have before sending on it blocks")
	globalconf.Register("kafka-mdm-in", inKafkaMdm)
}

func ConfigProcess(instance string) {
	if !Enabled {
		return
	}

	if offsetCommitInterval == 0 {
		log.Fatal("kafkamdm: offset-commit-interval must be greater then 0")
	}
	if consumerMaxWaitTime == 0 {
		log.Fatal("kafkamdm: consumer-max-wait-time must be greater then 0")
	}
	if consumerMaxProcessingTime == 0 {
		log.Fatal("kafkamdm: consumer-max-processing-time must be greater then 0")
	}

	switch offset {
	case "last":
	case "oldest":
	case "newest":
	default:
		_, err := time.ParseDuration(offset)
		if err != nil {
			log.Fatal(4, "kafkamdm: invalid offest format. %s", err)
		}
	}

	var err error
	offsetMgr, err = NewOffsetMgr(dataDir)
	if err != nil {
		log.Fatal(4, "kafka-mdm couldnt create offsetMgr. %s", err)
	}
	brokers = strings.Split(brokerStr, ",")
	topics = strings.Split(topicStr, ",")

	config = sarama.NewConfig()

	config.ClientID = instance + "-mdm"
	config.ChannelBufferSize = channelBufferSize
	config.Consumer.Fetch.Min = int32(consumerFetchMin)
	config.Consumer.Fetch.Default = int32(consumerFetchDefault)
	config.Consumer.MaxWaitTime = consumerMaxWaitTime
	config.Consumer.MaxProcessingTime = consumerMaxProcessingTime
	config.Net.MaxOpenRequests = netMaxOpenRequests
	config.Version = sarama.V0_10_0_0
	err = config.Validate()
	if err != nil {
		log.Fatal(2, "kafka-mdm invalid config: %s", err)
	}
}

func New(stats met.Backend) *KafkaMdm {
	client, err := sarama.NewClient(brokers, config)
	if err != nil {
		log.Fatal(4, "kafka-mdm failed to create client. %s", err)
	}
	consumer, err := sarama.NewConsumerFromClient(client)
	if err != nil {
		log.Fatal(2, "kafka-mdm failed to create consumer: %s", err)
	}
	log.Info("kafka-mdm consumer created without error")
	k := KafkaMdm{
		consumer:      consumer,
		client:        client,
		stats:         stats,
		StopChan:      make(chan int),
		stopConsuming: make(chan struct{}),
	}

	return &k
}

func (k *KafkaMdm) Start(metrics mdata.Metrics, metricIndex idx.MetricIndex, usg *usage.Usage) {
	k.In = in.New(metrics, metricIndex, usg, "kafka-mdm", k.stats)
	for _, topic := range topics {
		// get partitions.
		partitions, err := k.consumer.Partitions(topic)
		if err != nil {
			log.Fatal(4, "kafka-mdm: Faild to get partitions for topic %s. %s", topic, err)
		}
		for _, partition := range partitions {
			switch offset {
			case "oldest":
				go k.consumePartition(topic, partition, -2)
			case "newest":
				go k.consumePartition(topic, partition, -1)
			case "last":
				o, err := offsetMgr.Last(topic, partition)
				if err != nil {
					log.Fatal(4, "kafka-mdm: Failed to get offset for %s:%d. %s", topic, partition, err)
				}
				go k.consumePartition(topic, partition, o)
			default:
				// treat offset as a time duration.
				dur, err := time.ParseDuration(offset)
				if err != nil {
					log.Fatal(4, "kafka-mdm: Failed to parse offset to time.Duration. %s", err)
				}
				o, err := k.client.GetOffset(topic, partition, time.Now().Add(-1*dur).UnixNano()/int64(time.Millisecond))
				if err != nil {
					log.Fatal(4, "kafka-mdm: failed to get offset for %s:%d.  %s", topic, partition, err)
				}
				go k.consumePartition(topic, partition, o)
			}

		}
	}
}

// this will continually consume from the topic until k.stopConsuming is triggered.
func (k *KafkaMdm) consumePartition(topic string, partition int32, partitionOffset int64) {
	k.wg.Add(1)
	defer k.wg.Done()

	pc, err := k.consumer.ConsumePartition(topic, partition, partitionOffset)
	if err != nil {
		log.Fatal(4, "kafka-mdm: failed to start partitionConsumer for %s:%d. %s", topic, partition, err)
	}
	log.Info("kafka-mdm: consuming from %s:%d from offset %d", topic, partition, partitionOffset)
	currentOffset := partitionOffset
	messages := pc.Messages()
	ticker := time.NewTicker(time.Second * 5)
	for {
		select {
		case msg := <-messages:
			if LogLevel < 2 {
				log.Debug("kafka-mdm received message: Topic %s, Partition: %d, Offset: %d, Key: %x", msg.Topic, msg.Partition, msg.Offset, msg.Key)
			}
			k.In.Handle(msg.Value)
			currentOffset = msg.Offset
		case <-ticker.C:
			if err := offsetMgr.Commit(topic, partition, currentOffset); err != nil {
				log.Error(3, "kafka-mdm failed to commit offset for %s:%d, %s", topic, partition, err)
			}
		case <-k.stopConsuming:
			pc.Close()
			offsetMgr.Commit(topic, partition, currentOffset)
			log.Info("kafka-mdm consumer for %s:%d ended.", topic, partition)
			return
		}
	}
}

// Stop will initiate a graceful stop of the Consumer (permanent)
//
// NOTE: receive on StopChan to block until this process completes
func (k *KafkaMdm) Stop() {
	// closes notifications and messages channels, amongst others
	close(k.stopConsuming)
	go func() {
		k.wg.Wait()
		offsetMgr.Close()
		close(k.StopChan)
	}()
}

type OffsetMgr struct {
	db *leveldb.DB
}

func NewOffsetMgr(dir string) (*OffsetMgr, error) {
	dbFile := filepath.Join(dir, "partitionOffsets.db")
	db, err := leveldb.OpenFile(dbFile, &opt.Options{})
	if err != nil {
		if _, ok := err.(*storage.ErrCorrupted); ok {
			log.Warn("partitionOffsets.db is corrupt. Recovering.")
			db, err = leveldb.RecoverFile(dbFile, &opt.Options{})
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	log.Info("Opened %s", dbFile)
	return &OffsetMgr{
		db: db,
	}, nil
}

func (o *OffsetMgr) Close() {
	log.Info("Closing partitionsOffset DB.")
	o.db.Close()
}

func (o *OffsetMgr) Commit(topic string, partition int32, offset int64) error {
	key := new(bytes.Buffer)
	key.WriteString(fmt.Sprintf("T:%s-P:%d", topic, partition))
	data := new(bytes.Buffer)
	if err := binary.Write(data, binary.LittleEndian, offset); err != nil {
		return err
	}
	log.Debug("commiting offset %d for %s:%d to partitionsOffset.db", offset, topic, partition)
	return o.db.Put(key.Bytes(), data.Bytes(), &opt.WriteOptions{Sync: true})
}

func (o *OffsetMgr) Last(topic string, partition int32) (int64, error) {
	key := new(bytes.Buffer)
	key.WriteString(fmt.Sprintf("T:%s-P:%d", topic, partition))
	data, err := o.db.Get(key.Bytes(), nil)
	if err != nil {
		if err == leveldb.ErrNotFound {
			log.Debug("no offset recorded for %s:%d", topic, partition)
			return -1, nil
		}
		return 0, err
	}
	var offset int64
	err = binary.Read(bytes.NewBuffer(data), binary.LittleEndian, &offset)
	if err != nil {
		return 0, err
	}
	log.Debug("found saved offset %d for %s:%d", offset, topic, partition)
	return offset, nil
}
<commit_msg>only parse offsetDuration once.<commit_after>package kafkamdm

import (
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"github.com/Shopify/sarama"
	"github.com/raintank/worldping-api/pkg/log"
	"github.com/rakyll/globalconf"

	"github.com/raintank/met"
	"github.com/raintank/metrictank/idx"
	"github.com/raintank/metrictank/in"
	"github.com/raintank/metrictank/mdata"
	"github.com/raintank/metrictank/usage"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"github.com/syndtr/goleveldb/leveldb/storage"
)

type KafkaMdm struct {
	in.In
	consumer sarama.Consumer
	client   sarama.Client
	stats    met.Backend

	wg sync.WaitGroup
	// read from this channel to block until consumer is cleanly stopped
	StopChan chan int

	// signal to PartitionConsumers to shutdown
	stopConsuming chan struct{}
}

var LogLevel int
var Enabled bool
var brokerStr string
var brokers []string
var topicStr string
var topics []string
var offset string
var dataDir string
var config *sarama.Config
var channelBufferSize int
var consumerFetchMin int
var consumerFetchDefault int
var consumerMaxWaitTime time.Duration
var consumerMaxProcessingTime time.Duration
var netMaxOpenRequests int
var offsetMgr *OffsetMgr
var offsetDuration time.Duration
var offsetCommitInterval time.Duration

func ConfigSetup() {
	inKafkaMdm := flag.NewFlagSet("kafka-mdm-in", flag.ExitOnError)
	inKafkaMdm.BoolVar(&Enabled, "enabled", false, "")
	inKafkaMdm.StringVar(&brokerStr, "brokers", "kafka:9092", "tcp address for kafka (may be be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&topicStr, "topics", "mdm", "kafka topic (may be given multiple times as a comma-separated list)")
	inKafkaMdm.StringVar(&offset, "offset", "last", "Set the offset to start consuming from. Can be one of newest, oldest,last or a time duration")
	inKafkaMdm.DurationVar(&offsetCommitInterval, "offset-commit-interval", time.Second*5, "Interval at which offsets should be saved.")
	inKafkaMdm.StringVar(&dataDir, "data-dir", "", "Directory to store partition offsets index")
	inKafkaMdm.IntVar(&channelBufferSize, "channel-buffer-size", 1000000, "The number of metrics to buffer in internal and external channels")
	inKafkaMdm.IntVar(&consumerFetchMin, "consumer-fetch-min", 1024000, "The minimum number of message bytes to fetch in a request")
	inKafkaMdm.IntVar(&consumerFetchDefault, "consumer-fetch-default", 4096000, "The default number of message bytes to fetch in a request")
	inKafkaMdm.DurationVar(&consumerMaxWaitTime, "consumer-max-wait-time", time.Second, "The maximum amount of time the broker will wait for Consumer.Fetch.Min bytes to become available before it returns fewer than that anyway")
	inKafkaMdm.DurationVar(&consumerMaxProcessingTime, "consumer-max-processing-time", time.Second, "The maximum amount of time the consumer expects a message takes to process")
	inKafkaMdm.IntVar(&netMaxOpenRequests, "net-max-open-requests", 100, "How many outstanding requests a connection is allowed to have before sending on it blocks")
	globalconf.Register("kafka-mdm-in", inKafkaMdm)
}

func ConfigProcess(instance string) {
	if !Enabled {
		return
	}

	if offsetCommitInterval == 0 {
		log.Fatal("kafkamdm: offset-commit-interval must be greater then 0")
	}
	if consumerMaxWaitTime == 0 {
		log.Fatal("kafkamdm: consumer-max-wait-time must be greater then 0")
	}
	if consumerMaxProcessingTime == 0 {
		log.Fatal("kafkamdm: consumer-max-processing-time must be greater then 0")
	}

	switch offset {
	case "last":
	case "oldest":
	case "newest":
	default:
		offsetDuration, err = time.ParseDuration(offset)
		if err != nil {
			log.Fatal(4, "kafkamdm: invalid offest format. %s", err)
		}
	}

	var err error
	offsetMgr, err = NewOffsetMgr(dataDir)
	if err != nil {
		log.Fatal(4, "kafka-mdm couldnt create offsetMgr. %s", err)
	}
	brokers = strings.Split(brokerStr, ",")
	topics = strings.Split(topicStr, ",")

	config = sarama.NewConfig()

	config.ClientID = instance + "-mdm"
	config.ChannelBufferSize = channelBufferSize
	config.Consumer.Fetch.Min = int32(consumerFetchMin)
	config.Consumer.Fetch.Default = int32(consumerFetchDefault)
	config.Consumer.MaxWaitTime = consumerMaxWaitTime
	config.Consumer.MaxProcessingTime = consumerMaxProcessingTime
	config.Net.MaxOpenRequests = netMaxOpenRequests
	config.Version = sarama.V0_10_0_0
	err = config.Validate()
	if err != nil {
		log.Fatal(2, "kafka-mdm invalid config: %s", err)
	}
}

func New(stats met.Backend) *KafkaMdm {
	client, err := sarama.NewClient(brokers, config)
	if err != nil {
		log.Fatal(4, "kafka-mdm failed to create client. %s", err)
	}
	consumer, err := sarama.NewConsumerFromClient(client)
	if err != nil {
		log.Fatal(2, "kafka-mdm failed to create consumer: %s", err)
	}
	log.Info("kafka-mdm consumer created without error")
	k := KafkaMdm{
		consumer:      consumer,
		client:        client,
		stats:         stats,
		StopChan:      make(chan int),
		stopConsuming: make(chan struct{}),
	}

	return &k
}

func (k *KafkaMdm) Start(metrics mdata.Metrics, metricIndex idx.MetricIndex, usg *usage.Usage) {
	k.In = in.New(metrics, metricIndex, usg, "kafka-mdm", k.stats)
	for _, topic := range topics {
		// get partitions.
		partitions, err := k.consumer.Partitions(topic)
		if err != nil {
			log.Fatal(4, "kafka-mdm: Faild to get partitions for topic %s. %s", topic, err)
		}
		for _, partition := range partitions {
			switch offset {
			case "oldest":
				go k.consumePartition(topic, partition, -2)
			case "newest":
				go k.consumePartition(topic, partition, -1)
			case "last":
				o, err := offsetMgr.Last(topic, partition)
				if err != nil {
					log.Fatal(4, "kafka-mdm: Failed to get offset for %s:%d. %s", topic, partition, err)
				}
				go k.consumePartition(topic, partition, o)
			default:
				o, err := k.client.GetOffset(topic, partition, time.Now().Add(-1*offsetDuration).UnixNano()/int64(time.Millisecond))
				if err != nil {
					log.Fatal(4, "kafka-mdm: failed to get offset for %s:%d.  %s", topic, partition, err)
				}
				go k.consumePartition(topic, partition, o)
			}
		}
	}
}

// this will continually consume from the topic until k.stopConsuming is triggered.
func (k *KafkaMdm) consumePartition(topic string, partition int32, partitionOffset int64) {
	k.wg.Add(1)
	defer k.wg.Done()

	pc, err := k.consumer.ConsumePartition(topic, partition, partitionOffset)
	if err != nil {
		log.Fatal(4, "kafka-mdm: failed to start partitionConsumer for %s:%d. %s", topic, partition, err)
	}
	log.Info("kafka-mdm: consuming from %s:%d from offset %d", topic, partition, partitionOffset)
	currentOffset := partitionOffset
	messages := pc.Messages()
	ticker := time.NewTicker(time.Second * 5)
	for {
		select {
		case msg := <-messages:
			if LogLevel < 2 {
				log.Debug("kafka-mdm received message: Topic %s, Partition: %d, Offset: %d, Key: %x", msg.Topic, msg.Partition, msg.Offset, msg.Key)
			}
			k.In.Handle(msg.Value)
			currentOffset = msg.Offset
		case <-ticker.C:
			if err := offsetMgr.Commit(topic, partition, currentOffset); err != nil {
				log.Error(3, "kafka-mdm failed to commit offset for %s:%d, %s", topic, partition, err)
			}
		case <-k.stopConsuming:
			pc.Close()
			offsetMgr.Commit(topic, partition, currentOffset)
			log.Info("kafka-mdm consumer for %s:%d ended.", topic, partition)
			return
		}
	}
}

// Stop will initiate a graceful stop of the Consumer (permanent)
//
// NOTE: receive on StopChan to block until this process completes
func (k *KafkaMdm) Stop() {
	// closes notifications and messages channels, amongst others
	close(k.stopConsuming)
	go func() {
		k.wg.Wait()
		offsetMgr.Close()
		close(k.StopChan)
	}()
}

type OffsetMgr struct {
	db *leveldb.DB
}

func NewOffsetMgr(dir string) (*OffsetMgr, error) {
	dbFile := filepath.Join(dir, "partitionOffsets.db")
	db, err := leveldb.OpenFile(dbFile, &opt.Options{})
	if err != nil {
		if _, ok := err.(*storage.ErrCorrupted); ok {
			log.Warn("partitionOffsets.db is corrupt. Recovering.")
			db, err = leveldb.RecoverFile(dbFile, &opt.Options{})
			if err != nil {
				return nil, err
			}
		} else {
			return nil, err
		}
	}
	log.Info("Opened %s", dbFile)
	return &OffsetMgr{
		db: db,
	}, nil
}

func (o *OffsetMgr) Close() {
	log.Info("Closing partitionsOffset DB.")
	o.db.Close()
}

func (o *OffsetMgr) Commit(topic string, partition int32, offset int64) error {
	key := new(bytes.Buffer)
	key.WriteString(fmt.Sprintf("T:%s-P:%d", topic, partition))
	data := new(bytes.Buffer)
	if err := binary.Write(data, binary.LittleEndian, offset); err != nil {
		return err
	}
	log.Debug("commiting offset %d for %s:%d to partitionsOffset.db", offset, topic, partition)
	return o.db.Put(key.Bytes(), data.Bytes(), &opt.WriteOptions{Sync: true})
}

func (o *OffsetMgr) Last(topic string, partition int32) (int64, error) {
	key := new(bytes.Buffer)
	key.WriteString(fmt.Sprintf("T:%s-P:%d", topic, partition))
	data, err := o.db.Get(key.Bytes(), nil)
	if err != nil {
		if err == leveldb.ErrNotFound {
			log.Debug("no offset recorded for %s:%d", topic, partition)
			return -1, nil
		}
		return 0, err
	}
	var offset int64
	err = binary.Read(bytes.NewBuffer(data), binary.LittleEndian, &offset)
	if err != nil {
		return 0, err
	}
	log.Debug("found saved offset %d for %s:%d", offset, topic, partition)
	return offset, nil
}
<|endoftext|>