<commit_before>// +build !integration
// Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

package engine

import (
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
	"github.com/aws/amazon-ecs-agent/agent/engine/testdata"
	"github.com/aws/amazon-ecs-agent/agent/statemanager"
	"github.com/aws/amazon-ecs-agent/agent/utils/ttime/mocks"
	"github.com/stretchr/testify/assert"

	"github.com/golang/mock/gomock"
)

func TestCleanupTask(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockTime := mock_ttime.NewMockTime(ctrl)
	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
	mockClient := NewMockDockerClient(ctrl)
	mockImageManager := NewMockImageManager(ctrl)
	defer ctrl.Finish()

	taskEngine := &DockerTaskEngine{
		saver:        statemanager.NewNoopStateManager(),
		state:        mockState,
		client:       mockClient,
		imageManager: mockImageManager,
	}
	mTask := &managedTask{
		Task:           testdata.LoadTask("sleep5"),
		_time:          mockTime,
		engine:         taskEngine,
		acsMessages:    make(chan acsTransition),
		dockerMessages: make(chan dockerContainerChange),
	}
	mTask.SetKnownStatus(api.TaskStopped)
	mTask.SetSentStatus(api.TaskStopped)
	container := mTask.Containers[0]
	dockerContainer := &api.DockerContainer{
		DockerName: "dockerContainer",
	}

	// Expectations for triggering cleanup
	now := mTask.GetKnownStatusTime()
	taskStoppedDuration := 1 * time.Minute
	mockTime.EXPECT().Now().Return(now).AnyTimes()
	cleanupTimeTrigger := make(chan time.Time)
	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
	go func() {
		cleanupTimeTrigger <- now
	}()

	// Expectations to verify that the task gets removed
	mockState.EXPECT().ContainerMapByArn(mTask.Arn).Return(map[string]*api.DockerContainer{container.Name: dockerContainer}, true)
	mockClient.EXPECT().RemoveContainer(dockerContainer.DockerName, gomock.Any()).Return(nil)
	mockImageManager.EXPECT().RemoveContainerReferenceFromImageState(container).Return(nil)
	mockState.EXPECT().RemoveTask(mTask.Task)
	mTask.cleanupTask(taskStoppedDuration)
}

func TestCleanupTaskWaitsForStoppedSent(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockTime := mock_ttime.NewMockTime(ctrl)
	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
	mockClient := NewMockDockerClient(ctrl)
	mockImageManager := NewMockImageManager(ctrl)
	defer ctrl.Finish()

	taskEngine := &DockerTaskEngine{
		saver:        statemanager.NewNoopStateManager(),
		state:        mockState,
		client:       mockClient,
		imageManager: mockImageManager,
	}
	mTask := &managedTask{
		Task:           testdata.LoadTask("sleep5"),
		_time:          mockTime,
		engine:         taskEngine,
		acsMessages:    make(chan acsTransition),
		dockerMessages: make(chan dockerContainerChange),
	}
	mTask.SetKnownStatus(api.TaskStopped)
	mTask.SetSentStatus(api.TaskRunning)
	container := mTask.Containers[0]
	dockerContainer := &api.DockerContainer{
		DockerName: "dockerContainer",
	}

	// Expectations for triggering cleanup
	now := mTask.GetKnownStatusTime()
	taskStoppedDuration := 1 * time.Minute
	mockTime.EXPECT().Now().Return(now).AnyTimes()
	cleanupTimeTrigger := make(chan time.Time)
	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
	go func() {
		cleanupTimeTrigger <- now
	}()
	timesCalled := 0
	callsExpected := 3
	mockTime.EXPECT().Sleep(gomock.Any()).AnyTimes().Do(func(_ interface{}) {
		timesCalled++
		if timesCalled == callsExpected {
			mTask.SetSentStatus(api.TaskStopped)
		} else if timesCalled > callsExpected {
			t.Errorf("Sleep called too many times, called %d but expected %d", timesCalled, callsExpected)
		}
	})
	assert.Equal(t, api.TaskRunning, mTask.GetSentStatus())

	// Expectations to verify that the task gets removed
	mockState.EXPECT().ContainerMapByArn(mTask.Arn).Return(map[string]*api.DockerContainer{container.Name: dockerContainer}, true)
	mockClient.EXPECT().RemoveContainer(dockerContainer.DockerName, gomock.Any()).Return(nil)
	mockImageManager.EXPECT().RemoveContainerReferenceFromImageState(container).Return(nil)
	mockState.EXPECT().RemoveTask(mTask.Task)
	mTask.cleanupTask(taskStoppedDuration)
	assert.Equal(t, api.TaskStopped, mTask.GetSentStatus())
}

func TestCleanupTaskGivesUpIfWaitingTooLong(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockTime := mock_ttime.NewMockTime(ctrl)
	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
	mockClient := NewMockDockerClient(ctrl)
	mockImageManager := NewMockImageManager(ctrl)
	defer ctrl.Finish()

	taskEngine := &DockerTaskEngine{
		saver:        statemanager.NewNoopStateManager(),
		state:        mockState,
		client:       mockClient,
		imageManager: mockImageManager,
	}
	mTask := &managedTask{
		Task:           testdata.LoadTask("sleep5"),
		_time:          mockTime,
		engine:         taskEngine,
		acsMessages:    make(chan acsTransition),
		dockerMessages: make(chan dockerContainerChange),
	}
	mTask.SetKnownStatus(api.TaskStopped)
	mTask.SetSentStatus(api.TaskRunning)

	// Expectations for triggering cleanup
	now := mTask.GetKnownStatusTime()
	taskStoppedDuration := 1 * time.Minute
	mockTime.EXPECT().Now().Return(now).AnyTimes()
	cleanupTimeTrigger := make(chan time.Time)
	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
	go func() {
		cleanupTimeTrigger <- now
	}()
	_maxStoppedWaitTimes = 10
	defer func() {
		// reset
		_maxStoppedWaitTimes = int(maxStoppedWaitTimes)
	}()
	mockTime.EXPECT().Sleep(gomock.Any()).Times(_maxStoppedWaitTimes)
	assert.Equal(t, api.TaskRunning, mTask.GetSentStatus())

	// No cleanup expected
	mTask.cleanupTask(taskStoppedDuration)
	assert.Equal(t, api.TaskRunning, mTask.GetSentStatus())
}
<commit_msg>engine: add tests for `managedTask.containerNextState`<commit_after>// +build !integration
// Copyright 2014-2017 Amazon.com, Inc. or its affiliates. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may
// not use this file except in compliance with the License. A copy of the
// License is located at
//
//	http://aws.amazon.com/apache2.0/
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.

package engine

import (
	"testing"
	"time"

	"github.com/aws/amazon-ecs-agent/agent/api"
	"github.com/aws/amazon-ecs-agent/agent/engine/dockerstate/mocks"
	"github.com/aws/amazon-ecs-agent/agent/engine/testdata"
	"github.com/aws/amazon-ecs-agent/agent/statemanager"
	"github.com/aws/amazon-ecs-agent/agent/utils/ttime/mocks"
	"github.com/stretchr/testify/assert"

	"github.com/golang/mock/gomock"
)

func TestContainerNextStateFromNone(t *testing.T) {
	// NONE -> RUNNING transition is allowed and actionable, when desired is Running
	// The exptected next status is Pulled
	testContainerNextStateAssertions(t,
		api.ContainerStatusNone, api.ContainerRunning,
		api.ContainerPulled, true, true)

	// NONE -> NONE transition is not be allowed and is not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerStatusNone, api.ContainerStatusNone,
		api.ContainerStatusNone, false, false)

	// NONE -> STOPPED transition will result in STOPPED and is allowed, but not
	// actionable, when desired is STOPPED
	testContainerNextStateAssertions(t,
		api.ContainerStatusNone, api.ContainerStopped,
		api.ContainerStopped, false, true)
}

func TestContainerNextStateFromPulled(t *testing.T) {
	// PULLED -> RUNNING transition is allowed and actionable, when desired is Running
	// The exptected next status is Created
	testContainerNextStateAssertions(t,
		api.ContainerPulled, api.ContainerRunning,
		api.ContainerCreated, true, true)

	// PULLED -> PULLED transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerPulled, api.ContainerPulled,
		api.ContainerStatusNone, false, false)

	// PULLED -> NONE transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerPulled, api.ContainerStatusNone,
		api.ContainerStatusNone, false, false)

	// PULLED -> STOPPED transition will result in STOPPED and is allowed, but not
	// actionable, when desired is STOPPED
	testContainerNextStateAssertions(t,
		api.ContainerPulled, api.ContainerStopped,
		api.ContainerStopped, false, true)
}

func TestContainerNextStateFromCreated(t *testing.T) {
	// CREATED -> RUNNING transition is allowed and actionable, when desired is Running
	// The exptected next status is Running
	testContainerNextStateAssertions(t,
		api.ContainerCreated, api.ContainerRunning,
		api.ContainerRunning, true, true)

	// CREATED -> CREATED transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerCreated, api.ContainerCreated,
		api.ContainerStatusNone, false, false)

	// CREATED -> NONE transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerCreated, api.ContainerStatusNone,
		api.ContainerStatusNone, false, false)

	// CREATED -> PULLED transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerCreated, api.ContainerPulled,
		api.ContainerStatusNone, false, false)

	// CREATED -> STOPPED transition will result in STOPPED and is allowed, but not
	// actionable, when desired is STOPPED
	testContainerNextStateAssertions(t,
		api.ContainerCreated, api.ContainerStopped,
		api.ContainerStopped, false, true)
}

func TestContainerNextStateFromRunning(t *testing.T) {
	// RUNNING -> STOPPED transition is allowed and actionable, when desired is Running
	// The exptected next status is STOPPED
	testContainerNextStateAssertions(t,
		api.ContainerRunning, api.ContainerStopped,
		api.ContainerStopped, true, true)

	// RUNNING -> RUNNING transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerRunning, api.ContainerRunning,
		api.ContainerStatusNone, false, false)

	// RUNNING -> NONE transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerRunning, api.ContainerStatusNone,
		api.ContainerStatusNone, false, false)

	// RUNNING -> PULLED transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerRunning, api.ContainerPulled,
		api.ContainerStatusNone, false, false)

	// RUNNING -> CREATED transition is not allowed and not actionable,
	// when desired is Running
	testContainerNextStateAssertions(t,
		api.ContainerRunning, api.ContainerCreated,
		api.ContainerStatusNone, false, false)
}

func testContainerNextStateAssertions(t *testing.T,
	containerCurrentStatus api.ContainerStatus,
	containerDesiredStatus api.ContainerStatus,
	expectedContainerStatus api.ContainerStatus,
	expectedTransitionActionable bool,
	expectedTransitionPossible bool) {
	container := &api.Container{
		DesiredStatusUnsafe: containerDesiredStatus,
		KnownStatusUnsafe:   containerCurrentStatus,
	}
	task := &managedTask{
		Task: &api.Task{
			Containers: []*api.Container{
				container,
			},
			DesiredStatusUnsafe: api.TaskRunning,
		},
	}
	nextStatus, actionRequired, possible := task.containerNextState(container)
	assert.Equal(t, nextStatus, expectedContainerStatus,
		"Retrieved next state [%s] != Expected next state [%s]",
		nextStatus.String(), expectedContainerStatus.String())
	assert.Equal(t, actionRequired, expectedTransitionActionable)
	assert.Equal(t, possible, expectedTransitionPossible)
}

func TestCleanupTask(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockTime := mock_ttime.NewMockTime(ctrl)
	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
	mockClient := NewMockDockerClient(ctrl)
	mockImageManager := NewMockImageManager(ctrl)
	defer ctrl.Finish()

	taskEngine := &DockerTaskEngine{
		saver:        statemanager.NewNoopStateManager(),
		state:        mockState,
		client:       mockClient,
		imageManager: mockImageManager,
	}
	mTask := &managedTask{
		Task:           testdata.LoadTask("sleep5"),
		_time:          mockTime,
		engine:         taskEngine,
		acsMessages:    make(chan acsTransition),
		dockerMessages: make(chan dockerContainerChange),
	}
	mTask.SetKnownStatus(api.TaskStopped)
	mTask.SetSentStatus(api.TaskStopped)
	container := mTask.Containers[0]
	dockerContainer := &api.DockerContainer{
		DockerName: "dockerContainer",
	}

	// Expectations for triggering cleanup
	now := mTask.GetKnownStatusTime()
	taskStoppedDuration := 1 * time.Minute
	mockTime.EXPECT().Now().Return(now).AnyTimes()
	cleanupTimeTrigger := make(chan time.Time)
	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
	go func() {
		cleanupTimeTrigger <- now
	}()

	// Expectations to verify that the task gets removed
	mockState.EXPECT().ContainerMapByArn(mTask.Arn).Return(map[string]*api.DockerContainer{container.Name: dockerContainer}, true)
	mockClient.EXPECT().RemoveContainer(dockerContainer.DockerName, gomock.Any()).Return(nil)
	mockImageManager.EXPECT().RemoveContainerReferenceFromImageState(container).Return(nil)
	mockState.EXPECT().RemoveTask(mTask.Task)
	mTask.cleanupTask(taskStoppedDuration)
}

func TestCleanupTaskWaitsForStoppedSent(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockTime := mock_ttime.NewMockTime(ctrl)
	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
	mockClient := NewMockDockerClient(ctrl)
	mockImageManager := NewMockImageManager(ctrl)
	defer ctrl.Finish()

	taskEngine := &DockerTaskEngine{
		saver:        statemanager.NewNoopStateManager(),
		state:        mockState,
		client:       mockClient,
		imageManager: mockImageManager,
	}
	mTask := &managedTask{
		Task:           testdata.LoadTask("sleep5"),
		_time:          mockTime,
		engine:         taskEngine,
		acsMessages:    make(chan acsTransition),
		dockerMessages: make(chan dockerContainerChange),
	}
	mTask.SetKnownStatus(api.TaskStopped)
	mTask.SetSentStatus(api.TaskRunning)
	container := mTask.Containers[0]
	dockerContainer := &api.DockerContainer{
		DockerName: "dockerContainer",
	}

	// Expectations for triggering cleanup
	now := mTask.GetKnownStatusTime()
	taskStoppedDuration := 1 * time.Minute
	mockTime.EXPECT().Now().Return(now).AnyTimes()
	cleanupTimeTrigger := make(chan time.Time)
	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
	go func() {
		cleanupTimeTrigger <- now
	}()
	timesCalled := 0
	callsExpected := 3
	mockTime.EXPECT().Sleep(gomock.Any()).AnyTimes().Do(func(_ interface{}) {
		timesCalled++
		if timesCalled == callsExpected {
			mTask.SetSentStatus(api.TaskStopped)
		} else if timesCalled > callsExpected {
			t.Errorf("Sleep called too many times, called %d but expected %d", timesCalled, callsExpected)
		}
	})
	assert.Equal(t, api.TaskRunning, mTask.GetSentStatus())

	// Expectations to verify that the task gets removed
	mockState.EXPECT().ContainerMapByArn(mTask.Arn).Return(map[string]*api.DockerContainer{container.Name: dockerContainer}, true)
	mockClient.EXPECT().RemoveContainer(dockerContainer.DockerName, gomock.Any()).Return(nil)
	mockImageManager.EXPECT().RemoveContainerReferenceFromImageState(container).Return(nil)
	mockState.EXPECT().RemoveTask(mTask.Task)
	mTask.cleanupTask(taskStoppedDuration)
	assert.Equal(t, api.TaskStopped, mTask.GetSentStatus())
}

func TestCleanupTaskGivesUpIfWaitingTooLong(t *testing.T) {
	ctrl := gomock.NewController(t)
	mockTime := mock_ttime.NewMockTime(ctrl)
	mockState := mock_dockerstate.NewMockTaskEngineState(ctrl)
	mockClient := NewMockDockerClient(ctrl)
	mockImageManager := NewMockImageManager(ctrl)
	defer ctrl.Finish()

	taskEngine := &DockerTaskEngine{
		saver:        statemanager.NewNoopStateManager(),
		state:        mockState,
		client:       mockClient,
		imageManager: mockImageManager,
	}
	mTask := &managedTask{
		Task:           testdata.LoadTask("sleep5"),
		_time:          mockTime,
		engine:         taskEngine,
		acsMessages:    make(chan acsTransition),
		dockerMessages: make(chan dockerContainerChange),
	}
	mTask.SetKnownStatus(api.TaskStopped)
	mTask.SetSentStatus(api.TaskRunning)

	// Expectations for triggering cleanup
	now := mTask.GetKnownStatusTime()
	taskStoppedDuration := 1 * time.Minute
	mockTime.EXPECT().Now().Return(now).AnyTimes()
	cleanupTimeTrigger := make(chan time.Time)
	mockTime.EXPECT().After(gomock.Any()).Return(cleanupTimeTrigger)
	go func() {
		cleanupTimeTrigger <- now
	}()
	_maxStoppedWaitTimes = 10
	defer func() {
		// reset
		_maxStoppedWaitTimes = int(maxStoppedWaitTimes)
	}()
	mockTime.EXPECT().Sleep(gomock.Any()).Times(_maxStoppedWaitTimes)
	assert.Equal(t, api.TaskRunning, mTask.GetSentStatus())

	// No cleanup expected
	mTask.cleanupTask(taskStoppedDuration)
	assert.Equal(t, api.TaskRunning, mTask.GetSentStatus())
}
<|endoftext|>


<commit_before>// Red-Black tree implementation. Each tree node (entry) contains Key and Value.
// Entries in the RbMap are always ordered according to the key value, so
// it can be used as ordered set (std::set in C++) or ordered map (std::map).
// Note: all methods are not goroutine-safe.
package rbt

// import ( "strings" ; "fmt" )

type RbMap struct {
    less       LessFunc
    root       *RbMapNode
    size       int
}

// Red-black tree node, contains key and value. It is safe to overwrite Value
// in-place.
type RbMapNode struct {
    left, right, parent *RbMapNode
    // key
    key          interface{}
    Value        interface{}
    isred        bool         // true == red, false == black
}

// LessFunc is a key comparsion function.
// Must return true if k1 < k2, false otherwise
type LessFunc func(k1, k2 interface{}) bool

// Create new RbMap with provided comparsion function. The latter must
// return true if k1 < k2, false otherwise.
func NewRbMap(lessFunc LessFunc) *RbMap {
    return &RbMap{ less: lessFunc }
}

// Find node by key and return its Value, returns nil if key not found.
func (t *RbMap) Find(key interface{}) interface{} {
    n := t.FindNode(key)
    if n != nil {
        return n.Value
    }
    return nil
}

// Find a node by key, returns nil if not found
func (t *RbMap) FindNode(key interface{}) *RbMapNode {
    x := t.root
    for x != nil {
        if t.less(x.key, key) {
            x = x.right
        } else {
            if t.less(key, x.key) {
                x = x.left
            } else {
                return x
            }
        }
    }
    return nil
}

// Get last node in the tree (with highest key value).
func (t *RbMap) Last() *RbMapNode {
    if nil == t.root {
        return nil
    }
    return t.root.max()
}

// Get first node in the tree (with lowest key value).
func (t *RbMap) First() *RbMapNode {
    if nil == t.root {
        return nil
    }
    return t.root.min()
}

// Returns previous node in the tree, in descending key value order.
func (x *RbMapNode) Next() *RbMapNode {
    if x.right != nil {
        return x.right.min()
    }
    y := x.parent
    for y != nil && x == y.right {
        x = y
        y = y.parent
    }
    return y
}

// Returns key associated with tree node.
func (x *RbMapNode) Key() interface{} {
    return x.key
}

// Returns next node in the tree, in ascending key value order.
func (x *RbMapNode) Prev() *RbMapNode {
    if x.left != nil {
        return x.left.max()
    }
    y := x.parent
    for y != nil && x == y.left {
        x = y
        y = y.parent
    }
    return y
}

// Returns number of entries in the tree. This function returns internal
// counter, therefore it is fast and safe to use in loops.
func (t *RbMap) Size() int {
    return t.size
}

// Remove all entries in the tree.
func (t *RbMap) Clear() {
    t.root = nil
    t.size = 0
}

// Insert key and value into the tree. If new entry is created, returns true.
// If key already exists, value gets replaced and Insert returns false.
func (t *RbMap) Insert(key interface{}, value interface{}) bool {
    x := t.root
    var y *RbMapNode

    for x != nil {
        y = x
        if t.less(x.key, key) {
            x = x.right
        } else if t.less(key, x.key) {
            x = x.left
        } else {
            x.Value = value
            return false // overwrite value
        }
    }
    z := &RbMapNode{parent: y, isred: true, key: key, Value: value}
    if y == nil {
        t.root = z
    } else {
        if t.less(key, y.key) {
            y.left = z
        } else {
            y.right = z
        }
    }
    t.rb_insert_fixup(z)
    t.size++
    return true
}

// Delete tree node by key. Returns true if key was found and deleted.
func (t *RbMap) Delete(key interface{}) bool {
    if z := t.FindNode(key); z != nil {
        t.DeleteNode(z)
        return true
    }
    return false
}

// Delete tree node.
func (t *RbMap) DeleteNode(n *RbMapNode) {
    var x *RbMapNode
    if nil != n.left && nil != n.right {
        x = n.left.max()
        n.key, n.Value = x.key, x.Value
        n = x
    }
    if nil == n.right {
        x = n.left
    } else {
        x = n.right
    }
    if isBlack(n) {
        n.isred = isRed(x)
        if nil != n.parent {
            t.rb_delete_fixup(n)
        }
    }
    t.rbreplace(n, x)
    if isRed(t.root) {
        t.root.isred = false
    }
    t.size--
}

func (t* RbMap) rb_delete_fixup(n *RbMapNode) {
    var s, p *RbMapNode
    for {
        s, p = n.sibling(), n.parent
        if isRed(s) {
            p.isred, s.isred = true, false
            if n == p.left {
                t.left_rotate(p)
                s = p.right
            } else {
                t.right_rotate(p)
                s = p.left
            }
        }
        if isBlack(p) && isBlack(s) && isBlack(s.left) && isBlack(s.right) {
            s.isred = true
            if nil != p.parent {
                n = p
                continue
            }
            return
        } else {
            break
        }
    }
    if isRed(n.parent) && isBlack(s) && isBlack(s.left) && isBlack(s.right) {
        s.isred, n.parent.isred = true, false
    } else {
        if isBlack(s) {
            if n == n.parent.left && isRed(s.left) && isBlack(s.right) {
                s.isred, s.left.isred = true, false
                t.right_rotate(s)
                s = n.parent.right
            } else if n == n.parent.right && isRed(s.right) && isBlack(s.left) {
                s.isred, s.right.isred = true, false
                t.left_rotate(s)
                s = n.parent.left
            }
        }
        s.isred = n.parent.isred
        n.parent.isred = false
        if n == n.parent.left {
            s.right.isred = false
            t.left_rotate(n.parent)
        } else {
            s.left.isred = false
            t.right_rotate(n.parent)
        }
    }
}

func (t *RbMap) rb_insert_fixup(x *RbMapNode) {
    var y *RbMapNode
    for isRed(x.parent) {
        if x.parent == x.parent.parent.left {
            y = x.parent.parent.right
            if isRed(y) {
                x.parent.isred, y.isred = false, false
                x.parent.parent.isred = true
                x = x.parent.parent
            } else {
                if x == x.parent.right {
                    x = x.parent
                    t.left_rotate(x)
                }
                x.parent.isred, x.parent.parent.isred = false, true
                t.right_rotate(x.parent.parent)
            }
        } else {
            y = x.parent.parent.left
            if isRed(y) {
                x.parent.isred, y.isred = false, false
                x.parent.parent.isred = true
                x = x.parent.parent
            } else {
                if x == x.parent.left {
                    x = x.parent
                    t.right_rotate(x)
                }
                x.parent.isred, x.parent.parent.isred = false, true
                t.left_rotate(x.parent.parent)
            }
        }
    }
    t.root.isred = false
}

func (n *RbMapNode) sibling() *RbMapNode {
    if n == n.parent.left {
        return n.parent.right
    } else {
        return n.parent.left
    }
}

func (n *RbMapNode) min() *RbMapNode {
    for n.left != nil {
        n = n.left
    }
    return n
}

func (n *RbMapNode) max() *RbMapNode {
    for n.right != nil {
        n = n.right
    }
    return n
}

func (t *RbMap) left_rotate(n *RbMapNode) {
    r := n.right
    t.rbreplace(n, r)
    n.right = r.left
    if nil != r.left {
        r.left.parent = n
    }
    r.left, n.parent = n, r
}

func (t *RbMap) right_rotate(n *RbMapNode) {
    l := n.left
    t.rbreplace(n, l)
    n.left = l.right
    if nil != l.right {
        l.right.parent = n
    }
    l.right, n.parent = n, l
}

func (t *RbMap) rbreplace(u, v *RbMapNode) {
    parent := u.parent
    if parent == nil {
        t.root = v
    } else if u == parent.left {
        parent.left = v
    } else {
        parent.right = v
    }
    if v != nil {
        v.parent = parent
    }
}

func isBlack(n *RbMapNode) bool {
    return nil == n || !n.isred
}

func isRed(n *RbMapNode) bool {
    return nil != n && n.isred
}

/*

// uncomment for debugging only, because this adds dependency on strings, fmt

func (n *RbMapNode) Dump(indent int, tag string) {
    idn := strings.Repeat(" ", indent*4)
    c := 'B'
    if n.isred { c = 'R' }
    fmt.Printf("%s%s[%v:%v]%c\n", idn, tag, n.Key(), n.Value, c)
    if n.left != nil {
        n.left.Dump(indent + 1, "L:")
    }
    if n.right != nil {
        n.right.Dump(indent + 1, "R:")
    }
}

func (t *RbMap) Dump() {
    if t.root == nil {
        fmt.Printf("<NULL TREE>\n")
    } else {
        t.root.Dump(0, "*")
    }
}
*/

// Internal tree consistency check used by tests.
func (t *RbMap) verify() {
    if nil == t.root { return }
    if isRed(t.root) { panic("root is red") }
    verify1(t.root)
    verify2(t.root)
}

func verify1(n *RbMapNode) {
    if isRed(n) {
        if !isBlack(n.left)   { panic("left is not black") }
        if !isBlack(n.right)  { panic("right is not black") }
        if !isBlack(n.parent) { panic("parent is not black") }
    }
    if nil == n { return }
    verify1(n.left)
    verify1(n.right)
}

func verify2(n *RbMapNode) {
    black_count_path := -1
    verify2h(n, 0, &black_count_path)
}

func verify2h(n *RbMapNode, black_count int, path_black_count *int) {
    if isBlack(n) {
        black_count++
    }
    if n == nil {
        if *path_black_count == -1 {
            *path_black_count = black_count;
        } else {
            if black_count != *path_black_count {
                panic("black count")
            }
        }
        return
    }
    verify2h(n.left,  black_count, path_black_count)
    verify2h(n.right, black_count, path_black_count)
}

<commit_msg>documentation fixes<commit_after>// Red-Black tree implementation. Each tree node (entry) contains Key and Value.
// Entries in the RbMap are always ordered according to the key value, so
// it can be used as ordered set (std::set in C++) or ordered map (std::map).
// Note: all methods are not goroutine-safe.
package rbt

// import ( "strings" ; "fmt" )

type RbMap struct {
    less       LessFunc
    root       *RbMapNode
    size       int
}

// Red-black tree node, contains key and value. It is safe to overwrite Value
// in-place.
type RbMapNode struct {
    left, right, parent *RbMapNode
    // key
    key          interface{}
    Value        interface{}
    isred        bool         // true == red, false == black
}

// LessFunc is a key comparsion function.
// Must return true if k1 < k2, false otherwise.
type LessFunc func(k1, k2 interface{}) bool

// Create new RbMap with provided key comparsion function.
func NewRbMap(lessFunc LessFunc) *RbMap {
    return &RbMap{ less: lessFunc }
}

// Find node by key and return its Value, returns nil if key not found.
func (t *RbMap) Find(key interface{}) interface{} {
    n := t.FindNode(key)
    if n != nil {
        return n.Value
    }
    return nil
}

// Find a node by key, returns nil if not found.
func (t *RbMap) FindNode(key interface{}) *RbMapNode {
    x := t.root
    for x != nil {
        if t.less(x.key, key) {
            x = x.right
        } else {
            if t.less(key, x.key) {
                x = x.left
            } else {
                return x
            }
        }
    }
    return nil
}

// Get last node in the tree (with highest key value).
func (t *RbMap) Last() *RbMapNode {
    if nil == t.root {
        return nil
    }
    return t.root.max()
}

// Get first node in the tree (with lowest key value).
func (t *RbMap) First() *RbMapNode {
    if nil == t.root {
        return nil
    }
    return t.root.min()
}

// Get next node, in ascending key value order.
func (x *RbMapNode) Next() *RbMapNode {
    if x.right != nil {
        return x.right.min()
    }
    y := x.parent
    for y != nil && x == y.right {
        x = y
        y = y.parent
    }
    return y
}

// Returns key associated with tree node.
func (x *RbMapNode) Key() interface{} {
    return x.key
}

// Get previous node, in descending key value order.
func (x *RbMapNode) Prev() *RbMapNode {
    if x.left != nil {
        return x.left.max()
    }
    y := x.parent
    for y != nil && x == y.left {
        x = y
        y = y.parent
    }
    return y
}

// Returns number of entries in the tree. This function returns internal
// counter, therefore it is fast and safe to use in loops.
func (t *RbMap) Size() int {
    return t.size
}

// Remove all entries in the tree.
func (t *RbMap) Clear() {
    t.root = nil
    t.size = 0
}

// Insert key and value into the tree. If new entry is created, returns true.
// If key already exists, value gets replaced and Insert returns false.
func (t *RbMap) Insert(key interface{}, value interface{}) bool {
    x := t.root
    var y *RbMapNode

    for x != nil {
        y = x
        if t.less(x.key, key) {
            x = x.right
        } else if t.less(key, x.key) {
            x = x.left
        } else {
            x.Value = value
            return false // overwrite value
        }
    }
    z := &RbMapNode{parent: y, isred: true, key: key, Value: value}
    if y == nil {
        t.root = z
    } else {
        if t.less(key, y.key) {
            y.left = z
        } else {
            y.right = z
        }
    }
    t.rb_insert_fixup(z)
    t.size++
    return true
}

// Delete tree node by key. Returns true if key was found and deleted.
func (t *RbMap) Delete(key interface{}) bool {
    if z := t.FindNode(key); z != nil {
        t.DeleteNode(z)
        return true
    }
    return false
}

// Delete tree node.
func (t *RbMap) DeleteNode(n *RbMapNode) {
    var x *RbMapNode
    if nil != n.left && nil != n.right {
        x = n.left.max()
        n.key, n.Value = x.key, x.Value
        n = x
    }
    if nil == n.right {
        x = n.left
    } else {
        x = n.right
    }
    if isBlack(n) {
        n.isred = isRed(x)
        if nil != n.parent {
            t.rb_delete_fixup(n)
        }
    }
    t.rbreplace(n, x)
    if isRed(t.root) {
        t.root.isred = false
    }
    t.size--
}

func (t* RbMap) rb_delete_fixup(n *RbMapNode) {
    var s, p *RbMapNode
    for {
        s, p = n.sibling(), n.parent
        if isRed(s) {
            p.isred, s.isred = true, false
            if n == p.left {
                t.left_rotate(p)
                s = p.right
            } else {
                t.right_rotate(p)
                s = p.left
            }
        }
        if isBlack(p) && isBlack(s) && isBlack(s.left) && isBlack(s.right) {
            s.isred = true
            if nil != p.parent {
                n = p
                continue
            }
            return
        } else {
            break
        }
    }
    if isRed(n.parent) && isBlack(s) && isBlack(s.left) && isBlack(s.right) {
        s.isred, n.parent.isred = true, false
    } else {
        if isBlack(s) {
            if n == n.parent.left && isRed(s.left) && isBlack(s.right) {
                s.isred, s.left.isred = true, false
                t.right_rotate(s)
                s = n.parent.right
            } else if n == n.parent.right && isRed(s.right) && isBlack(s.left) {
                s.isred, s.right.isred = true, false
                t.left_rotate(s)
                s = n.parent.left
            }
        }
        s.isred = n.parent.isred
        n.parent.isred = false
        if n == n.parent.left {
            s.right.isred = false
            t.left_rotate(n.parent)
        } else {
            s.left.isred = false
            t.right_rotate(n.parent)
        }
    }
}

func (t *RbMap) rb_insert_fixup(x *RbMapNode) {
    var y *RbMapNode
    for isRed(x.parent) {
        if x.parent == x.parent.parent.left {
            y = x.parent.parent.right
            if isRed(y) {
                x.parent.isred, y.isred = false, false
                x.parent.parent.isred = true
                x = x.parent.parent
            } else {
                if x == x.parent.right {
                    x = x.parent
                    t.left_rotate(x)
                }
                x.parent.isred, x.parent.parent.isred = false, true
                t.right_rotate(x.parent.parent)
            }
        } else {
            y = x.parent.parent.left
            if isRed(y) {
                x.parent.isred, y.isred = false, false
                x.parent.parent.isred = true
                x = x.parent.parent
            } else {
                if x == x.parent.left {
                    x = x.parent
                    t.right_rotate(x)
                }
                x.parent.isred, x.parent.parent.isred = false, true
                t.left_rotate(x.parent.parent)
            }
        }
    }
    t.root.isred = false
}

func (n *RbMapNode) sibling() *RbMapNode {
    if n == n.parent.left {
        return n.parent.right
    } else {
        return n.parent.left
    }
}

func (n *RbMapNode) min() *RbMapNode {
    for n.left != nil {
        n = n.left
    }
    return n
}

func (n *RbMapNode) max() *RbMapNode {
    for n.right != nil {
        n = n.right
    }
    return n
}

func (t *RbMap) left_rotate(n *RbMapNode) {
    r := n.right
    t.rbreplace(n, r)
    n.right = r.left
    if nil != r.left {
        r.left.parent = n
    }
    r.left, n.parent = n, r
}

func (t *RbMap) right_rotate(n *RbMapNode) {
    l := n.left
    t.rbreplace(n, l)
    n.left = l.right
    if nil != l.right {
        l.right.parent = n
    }
    l.right, n.parent = n, l
}

func (t *RbMap) rbreplace(u, v *RbMapNode) {
    parent := u.parent
    if parent == nil {
        t.root = v
    } else if u == parent.left {
        parent.left = v
    } else {
        parent.right = v
    }
    if v != nil {
        v.parent = parent
    }
}

func isBlack(n *RbMapNode) bool {
    return nil == n || !n.isred
}

func isRed(n *RbMapNode) bool {
    return nil != n && n.isred
}

/*

// uncomment for debugging only, because this adds dependency on strings, fmt

func (n *RbMapNode) Dump(indent int, tag string) {
    idn := strings.Repeat(" ", indent*4)
    c := 'B'
    if n.isred { c = 'R' }
    fmt.Printf("%s%s[%v:%v]%c\n", idn, tag, n.Key(), n.Value, c)
    if n.left != nil {
        n.left.Dump(indent + 1, "L:")
    }
    if n.right != nil {
        n.right.Dump(indent + 1, "R:")
    }
}

func (t *RbMap) Dump() {
    if t.root == nil {
        fmt.Printf("<NULL TREE>\n")
    } else {
        t.root.Dump(0, "*")
    }
}
*/

// Internal tree consistency check used by tests.
func (t *RbMap) verify() {
    if nil == t.root { return }
    if isRed(t.root) { panic("root is red") }
    verify1(t.root)
    verify2(t.root)
}

func verify1(n *RbMapNode) {
    if isRed(n) {
        if !isBlack(n.left)   { panic("left is not black") }
        if !isBlack(n.right)  { panic("right is not black") }
        if !isBlack(n.parent) { panic("parent is not black") }
    }
    if nil == n { return }
    verify1(n.left)
    verify1(n.right)
}

func verify2(n *RbMapNode) {
    black_count_path := -1
    verify2h(n, 0, &black_count_path)
}

func verify2h(n *RbMapNode, black_count int, path_black_count *int) {
    if isBlack(n) {
        black_count++
    }
    if n == nil {
        if *path_black_count == -1 {
            *path_black_count = black_count;
        } else {
            if black_count != *path_black_count {
                panic("black count")
            }
        }
        return
    }
    verify2h(n.left,  black_count, path_black_count)
    verify2h(n.right, black_count, path_black_count)
}

<|endoftext|>


<commit_before>package cmd

import (
	"math/rand"
	"sync"
	"time"
)

var (
	random  *rand.Rand
	getRand sync.Once

	getLocation sync.Once

	now    time.Time
	getNow sync.Once
)

func GetRand() *rand.Rand {
	getRand.Do(func() {
		random = rand.New(rand.NewSource(time.Now().UnixNano()))
	})
	return random
}

func GetLocation() *time.Location {
	getLocation.Do(func() {
		loc, _ := time.LoadLocation(GetFlags().Location)
		time.Local = loc
	})
	return time.Local
}

func Now() time.Time {
	getNow.Do(func() {
		timeString := GetFlags().Now
		if len(timeString) < 1 {
			now = time.Now()
		} else {
			t, _ := time.ParseInLocation("2006-01-02 15:04:05.999999999", timeString, GetLocation())
			now = t
		}
	})
	return now
}
<commit_msg>Release v0.3.3<commit_after>package cmd

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

var (
	random  *rand.Rand
	getRand sync.Once

	getLocation sync.Once

	now    time.Time
	getNow sync.Once
)

func GetRand() *rand.Rand {
	getRand.Do(func() {
		random = rand.New(rand.NewSource(time.Now().UnixNano()))
	})
	return random
}

func GetLocation() *time.Location {
	getLocation.Do(func() {
		loc, _ := time.LoadLocation(GetFlags().Location)
		time.Local = loc
		fmt.Println(GetFlags().Location)
	})
	return time.Local
}

func Now() time.Time {
	getNow.Do(func() {
		timeString := GetFlags().Now
		if len(timeString) < 1 {
			GetLocation()
			now = time.Now()
		} else {
			t, _ := time.ParseInLocation("2006-01-02 15:04:05.999999999", timeString, GetLocation())
			now = t
		}
	})
	return now
}
<|endoftext|>


<commit_before>package lib

import (
	"fmt"
	. "github.com/ahmetalpbalkan/go-linq"
)

func getCollection(allPosts []*LongPost) map[string][]*LongPost {
	collection, err := From(allPosts).GroupBy(func(post T) T { return post.(*LongPost).Category }, func(post T) T { return post.(*LongPost) })
	if err != nil {
		fmt.Errorf("Error", err)
	}
	return coverMapType(collection)
}

func coverMapType(in map[T][]T) (out map[string][]*LongPost) {
	out = make(map[string][]*LongPost, len(in))
	for k, _ := range in {
		if key, ok := k.(string); ok {
			v := in[k]
			out[key] = make([]*LongPost, 0, len(v))
			for i := range v {
				if value, ok := v[i].(*LongPost); ok {
					out[key] = append(out[key], value)
				}
			}
		}
	}
	return
}
<commit_msg>todo: remove go-linq<commit_after>package lib

import (
	"fmt"
	. "github.com/ahmetalpbalkan/go-linq"
)

// TODO: remove go-linq
func getCollection(allPosts []*LongPost) map[string][]*LongPost {
	collection, err := From(allPosts).GroupBy(func(post T) T { return post.(*LongPost).Category }, func(post T) T { return post.(*LongPost) })
	if err != nil {
		fmt.Errorf("Error", err)
	}
	return coverMapType(collection)
}

func coverMapType(in map[T][]T) (out map[string][]*LongPost) {
	out = make(map[string][]*LongPost, len(in))
	for k, _ := range in {
		if key, ok := k.(string); ok {
			v := in[k]
			out[key] = make([]*LongPost, 0, len(v))
			for i := range v {
				if value, ok := v[i].(*LongPost); ok {
					out[key] = append(out[key], value)
				}
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>// Copyright (C) 2015 Foursquare Labs Inc.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"
	"math/rand"
	"strings"

	"github.com/foursquare/fsgo/net/httpthrift"
	"github.com/foursquare/fsgo/report"
	"github.com/foursquare/quiver/gen"
	"github.com/foursquare/quiver/hfile"
	"github.com/foursquare/quiver/util"
)

func WrapHttpRpcHandler(cs *hfile.CollectionSet, stats *report.Recorder) *httpthrift.ThriftOverHTTPHandler {
	return httpthrift.NewThriftOverHTTPHandler(gen.NewHFileServiceProcessor(&ThriftRpcImpl{cs}), stats)
}

type ThriftRpcImpl struct {
	*hfile.CollectionSet
}

func (cs *ThriftRpcImpl) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if Settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	// TODO: clients should request strict during dev/testing?
	reader.EnforceKeyOrder = false
	defer reader.Release()

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if Settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		if idx > 0 && bytes.Equal(req.SortedKeys[idx-1], key) {
			if prev, ok := res.Values[int32(idx-1)]; ok {
				res.Values[int32(idx)] = prev
				found++
			}
			continue
		}

		if !hfile.MightContain(key) {
			continue
		}

		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			if !req.GetCountOnly() {
				if req.PerKeyValueLimit != nil {
					res.Values[int32(idx)] = value[:req.GetPerKeyValueLimit()]
				} else {
					res.Values[int32(idx)] = value
				}
			}
		}
	}

	if Settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *ThriftRpcImpl) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	if Settings.debug {
		log.Println("[GetValuesMulti]", len(req.SortedKeys))
	}

	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	res := new(gen.MultiHFileKeyResponse)
	res.Values = make(map[int32][][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if !hfile.MightContain(key) {
			continue
		}
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			res.Values[int32(idx)] = values
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *ThriftRpcImpl) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.ReaderFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.GetIterator()
		defer i.Release()
		limit := int32(0)
		if req.ValueLimit != nil {
			limit = *req.ValueLimit
		}
		if res.Values, res.LastKey, err = i.AllForPrefixes(req.SortedKeys, limit, req.LastKey); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *ThriftRpcImpl) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	scanner := reader.GetScanner()
	defer scanner.Release()

	for _, parts := range util.RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *ThriftRpcImpl) GetIterator(req *gen.IteratorRequest) (*gen.IteratorResponse, error) {
	// 	HfileName     *string `thrift:"hfileName,1" json:"hfileName"`
	// 	IncludeValues *bool   `thrift:"includeValues,2" json:"includeValues"`
	// 	LastKey       []byte  `thrift:"lastKey,3" json:"lastKey"`
	// 	SkipKeys      *int32  `thrift:"skipKeys,4" json:"skipKeys"`
	// 	ResponseLimit *int32  `thrift:"responseLimit,5" json:"responseLimit"`
	// 	EndKey        []byte  `thrift:"endKey,6" json:"endKey"`
	var err error

	if req.ResponseLimit == nil {
		return nil, fmt.Errorf("Missing limit.")
	}
	limit := int(*req.ResponseLimit)

	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	it := reader.GetIterator()
	defer it.Release()

	remaining := false

	if req.LastKey != nil {
		remaining, err = it.Seek(req.LastKey)
	} else {
		remaining, err = it.Next()
	}

	if err != nil {
		return nil, err
	}

	res := new(gen.IteratorResponse)

	if !remaining {
		return res, nil
	}

	skipKeys := int32(0)
	lastKey := it.Key()

	if toSkip := req.GetSkipKeys(); toSkip > 0 {
		for i := int32(0); i < toSkip && remaining; i++ {
			if bytes.Equal(lastKey, it.Key()) {
				skipKeys = skipKeys + 1
			} else {
				skipKeys = 0
			}

			lastKey = it.Key()

			remaining, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		if !remaining {
			return res, nil
		}
	}

	if req.EndKey != nil {
		remaining = remaining && !hfile.After(it.Key(), req.EndKey)
	}

	r := make([]*gen.KeyValueItem, 0)
	for i := 0; i < limit && remaining; i++ {
		v := []byte{}
		if req.IncludeValues == nil || *req.IncludeValues {
			v = it.Value()
		}
		r = append(r, &gen.KeyValueItem{it.Key(), v})

		if bytes.Equal(lastKey, it.Key()) {
			skipKeys = skipKeys + 1
		} else {
			skipKeys = 1
		}
		lastKey = it.Key()

		remaining, err = it.Next()
		if err != nil {
			return nil, err
		}
		if req.EndKey != nil {
			remaining = remaining && !hfile.After(it.Key(), req.EndKey)
		}
	}
	return &gen.IteratorResponse{r, lastKey, &skipKeys}, nil
}

func GetCollectionInfo(r *hfile.Reader, keySampleSize int) (*gen.HFileInfo, error) {
	i := new(gen.HFileInfo)
	i.Name = &r.Name
	i.Path = &r.SourcePath
	c := int64(r.EntryCount)
	i.NumElements = &c
	i.FirstKey, _ = r.FirstKey()

	if keySampleSize > 0 {
		it := r.GetIterator()
		defer it.Release()

		pr := float64(keySampleSize) / float64(c)
		buf := make([][]byte, keySampleSize)
		found := 0
		next, err := it.Next()
		for next && found < keySampleSize {
			if rand.Float64() < pr {
				buf[found] = it.Key()
				found++
			}
			next, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		buf = buf[:found]
		i.RandomKeys = buf
	}
	if Settings.debug {
		log.Printf("[GetCollectionInfo] %v (%d keys)\n", i, len(i.RandomKeys))
	}
	return i, nil
}

func (cs *ThriftRpcImpl) getInfo(req *gen.InfoRequest, allowRandom bool) (r []*gen.HFileInfo, err error) {
	require := ""
	if req.IsSetHfileName() {
		require := req.GetHfileName()
		if require != "" && !strings.ContainsRune(require, '/') {
			require = require + "/"
		}
	}

	sample := 0
	if req.IsSetNumRandomKeys() && allowRandom {
		sample = int(*req.NumRandomKeys)
	}

	for name, reader := range cs.Collections {
		if require == "" || strings.HasPrefix(name, require) {
			if i, err := GetCollectionInfo(reader, sample); err != nil {
				return nil, err
			} else {
				r = append(r, i)
			}
		}
	}

	return r, nil
}

func (cs *ThriftRpcImpl) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return cs.getInfo(req, false)
}

func (cs *ThriftRpcImpl) ScanCollectionAndSampleKeys(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return cs.getInfo(req, true)
}

func (cs *ThriftRpcImpl) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, fmt.Errorf("Not implemented")
}
<commit_msg>whoops, valueLimit should apply to multi, not single<commit_after>// Copyright (C) 2015 Foursquare Labs Inc.

package main

import (
	"bytes"
	"encoding/hex"
	"fmt"
	"log"
	"math/rand"
	"strings"

	"github.com/foursquare/fsgo/net/httpthrift"
	"github.com/foursquare/fsgo/report"
	"github.com/foursquare/quiver/gen"
	"github.com/foursquare/quiver/hfile"
	"github.com/foursquare/quiver/util"
)

func WrapHttpRpcHandler(cs *hfile.CollectionSet, stats *report.Recorder) *httpthrift.ThriftOverHTTPHandler {
	return httpthrift.NewThriftOverHTTPHandler(gen.NewHFileServiceProcessor(&ThriftRpcImpl{cs}), stats)
}

type ThriftRpcImpl struct {
	*hfile.CollectionSet
}

func (cs *ThriftRpcImpl) GetValuesSingle(req *gen.SingleHFileKeyRequest) (r *gen.SingleHFileKeyResponse, err error) {
	if Settings.debug {
		log.Printf("[GetValuesSingle] %s (%d keys)\n", *req.HfileName, len(req.SortedKeys))
	}
	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	// TODO: clients should request strict during dev/testing?
	reader.EnforceKeyOrder = false
	defer reader.Release()

	res := new(gen.SingleHFileKeyResponse)
	res.Values = make(map[int32][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if Settings.debug {
			log.Printf("[GetValuesSingle] key: %s\n", hex.EncodeToString(key))
		}
		if idx > 0 && bytes.Equal(req.SortedKeys[idx-1], key) {
			if prev, ok := res.Values[int32(idx-1)]; ok {
				res.Values[int32(idx)] = prev
				found++
			}
			continue
		}

		if !hfile.MightContain(key) {
			continue
		}

		value, err, ok := reader.GetFirst(key)
		if err != nil {
			return nil, err
		}
		if ok {
			found++
			if !req.GetCountOnly() {
				res.Values[int32(idx)] = value
			}
		}
	}

	if Settings.debug {
		log.Printf("[GetValuesSingle] %s found %d of %d.\n", *req.HfileName, found, len(req.SortedKeys))
	}
	res.KeyCount = &found
	return res, nil
}

func (cs *ThriftRpcImpl) GetValuesMulti(req *gen.SingleHFileKeyRequest) (r *gen.MultiHFileKeyResponse, err error) {
	if Settings.debug {
		log.Println("[GetValuesMulti]", len(req.SortedKeys))
	}

	hfile, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	reader := hfile.GetScanner()
	defer reader.Release()

	res := new(gen.MultiHFileKeyResponse)
	res.Values = make(map[int32][][]byte)
	found := int32(0)

	for idx, key := range req.SortedKeys {
		if !hfile.MightContain(key) {
			continue
		}
		values, err := reader.GetAll(key)
		if err != nil {
			return nil, err
		}
		if len(values) > 0 {
			found += int32(len(values))
			if req.PerKeyValueLimit != nil {
				res.Values[int32(idx)] = values[:req.GetPerKeyValueLimit()]
			} else {
				res.Values[int32(idx)] = values
			}
		}
	}

	res.KeyCount = &found
	return res, nil

}

func (cs *ThriftRpcImpl) GetValuesForPrefixes(req *gen.PrefixRequest) (r *gen.PrefixResponse, err error) {
	res := new(gen.PrefixResponse)
	if reader, err := cs.ReaderFor(*req.HfileName); err != nil {
		return nil, err
	} else {
		i := reader.GetIterator()
		defer i.Release()
		limit := int32(0)
		if req.ValueLimit != nil {
			limit = *req.ValueLimit
		}
		if res.Values, res.LastKey, err = i.AllForPrefixes(req.SortedKeys, limit, req.LastKey); err != nil {
			return nil, err
		} else {
			return res, nil
		}
	}
}

func (cs *ThriftRpcImpl) GetValuesMultiSplitKeys(req *gen.MultiHFileSplitKeyRequest) (r *gen.KeyToValuesResponse, err error) {
	res := make(map[string][][]byte)
	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	scanner := reader.GetScanner()
	defer scanner.Release()

	for _, parts := range util.RevProduct(req.SplitKey) {
		// TODO(davidt): avoid allocing concated key by adding split-key search lower down.
		key := bytes.Join(parts, nil)

		if values, err := scanner.GetAll(key); err != nil {
			return nil, err
		} else if len(values) > 0 {
			res[string(key)] = values
		}
	}
	return &gen.KeyToValuesResponse{res}, nil
}

func (cs *ThriftRpcImpl) GetIterator(req *gen.IteratorRequest) (*gen.IteratorResponse, error) {
	// 	HfileName     *string `thrift:"hfileName,1" json:"hfileName"`
	// 	IncludeValues *bool   `thrift:"includeValues,2" json:"includeValues"`
	// 	LastKey       []byte  `thrift:"lastKey,3" json:"lastKey"`
	// 	SkipKeys      *int32  `thrift:"skipKeys,4" json:"skipKeys"`
	// 	ResponseLimit *int32  `thrift:"responseLimit,5" json:"responseLimit"`
	// 	EndKey        []byte  `thrift:"endKey,6" json:"endKey"`
	var err error

	if req.ResponseLimit == nil {
		return nil, fmt.Errorf("Missing limit.")
	}
	limit := int(*req.ResponseLimit)

	reader, err := cs.ReaderFor(*req.HfileName)
	if err != nil {
		return nil, err
	}
	it := reader.GetIterator()
	defer it.Release()

	remaining := false

	if req.LastKey != nil {
		remaining, err = it.Seek(req.LastKey)
	} else {
		remaining, err = it.Next()
	}

	if err != nil {
		return nil, err
	}

	res := new(gen.IteratorResponse)

	if !remaining {
		return res, nil
	}

	skipKeys := int32(0)
	lastKey := it.Key()

	if toSkip := req.GetSkipKeys(); toSkip > 0 {
		for i := int32(0); i < toSkip && remaining; i++ {
			if bytes.Equal(lastKey, it.Key()) {
				skipKeys = skipKeys + 1
			} else {
				skipKeys = 0
			}

			lastKey = it.Key()

			remaining, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		if !remaining {
			return res, nil
		}
	}

	if req.EndKey != nil {
		remaining = remaining && !hfile.After(it.Key(), req.EndKey)
	}

	r := make([]*gen.KeyValueItem, 0)
	for i := 0; i < limit && remaining; i++ {
		v := []byte{}
		if req.IncludeValues == nil || *req.IncludeValues {
			v = it.Value()
		}
		r = append(r, &gen.KeyValueItem{it.Key(), v})

		if bytes.Equal(lastKey, it.Key()) {
			skipKeys = skipKeys + 1
		} else {
			skipKeys = 1
		}
		lastKey = it.Key()

		remaining, err = it.Next()
		if err != nil {
			return nil, err
		}
		if req.EndKey != nil {
			remaining = remaining && !hfile.After(it.Key(), req.EndKey)
		}
	}
	return &gen.IteratorResponse{r, lastKey, &skipKeys}, nil
}

func GetCollectionInfo(r *hfile.Reader, keySampleSize int) (*gen.HFileInfo, error) {
	i := new(gen.HFileInfo)
	i.Name = &r.Name
	i.Path = &r.SourcePath
	c := int64(r.EntryCount)
	i.NumElements = &c
	i.FirstKey, _ = r.FirstKey()

	if keySampleSize > 0 {
		it := r.GetIterator()
		defer it.Release()

		pr := float64(keySampleSize) / float64(c)
		buf := make([][]byte, keySampleSize)
		found := 0
		next, err := it.Next()
		for next && found < keySampleSize {
			if rand.Float64() < pr {
				buf[found] = it.Key()
				found++
			}
			next, err = it.Next()
			if err != nil {
				return nil, err
			}
		}
		buf = buf[:found]
		i.RandomKeys = buf
	}
	if Settings.debug {
		log.Printf("[GetCollectionInfo] %v (%d keys)\n", i, len(i.RandomKeys))
	}
	return i, nil
}

func (cs *ThriftRpcImpl) getInfo(req *gen.InfoRequest, allowRandom bool) (r []*gen.HFileInfo, err error) {
	require := ""
	if req.IsSetHfileName() {
		require := req.GetHfileName()
		if require != "" && !strings.ContainsRune(require, '/') {
			require = require + "/"
		}
	}

	sample := 0
	if req.IsSetNumRandomKeys() && allowRandom {
		sample = int(*req.NumRandomKeys)
	}

	for name, reader := range cs.Collections {
		if require == "" || strings.HasPrefix(name, require) {
			if i, err := GetCollectionInfo(reader, sample); err != nil {
				return nil, err
			} else {
				r = append(r, i)
			}
		}
	}

	return r, nil
}

func (cs *ThriftRpcImpl) GetInfo(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return cs.getInfo(req, false)
}

func (cs *ThriftRpcImpl) ScanCollectionAndSampleKeys(req *gen.InfoRequest) (r []*gen.HFileInfo, err error) {
	return cs.getInfo(req, true)
}

func (cs *ThriftRpcImpl) TestTimeout(waitInMillis int32) (r int32, err error) {
	return 0, fmt.Errorf("Not implemented")
}
<|endoftext|>


<commit_before>package rss

import (
	"bytes"
	"encoding/xml"
	"io/ioutil"
	"net/http"
)

type Rss struct {
	Channel Channel `xml:"channel"`
}
type Item struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Enclosure		Enclosure `xml:"enclosure"`
}
type Enclosure struct {
	Url					string `xml:"url,attr"`
	Type				string `xml:"type,attr"`
}
type Channel struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Items       []Item `xml:"item"`
}

func ResponseToRss(r *http.Response) (Rss, error) {
	rss := Rss{}
	XMLdata, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return rss, err
	}
	buffer := bytes.NewBuffer(XMLdata)
	decoded := xml.NewDecoder(buffer)
	err = decoded.Decode(&rss)
	return rss, err
}
/*
func main() {

         response, err := http.Get("http://www.thestar.com.my/RSS/Metro/Community/")

         if err != nil {
                 fmt.Println(err)
                 os.Exit(1)
         }

         defer response.Body.Close()

         XMLdata, err := ioutil.ReadAll(response.Body)

         if err != nil {
                 fmt.Println(err)
                 os.Exit(1)
         }

         rss := new(Rss)

         buffer := bytes.NewBuffer(XMLdata)

         decoded := xml.NewDecoder(buffer)

         err = decoded.Decode(rss)

         if err != nil {
                 fmt.Println(err)
                 os.Exit(1)
         }

         fmt.Printf("Title : %s\n", rss.Channel.Title)
         fmt.Printf("Description : %s\n", rss.Channel.Description)
         fmt.Printf("Link : %s\n", rss.Channel.Link)

         total := len(rss.Channel.Items)

         fmt.Printf("Total items : %v\n", total)

         for i := 0; i < total; i++ {
                 fmt.Printf("[%d] item title : %s\n", i, rss.Channel.Items[i].Title)
                 fmt.Printf("[%d] item description : %s\n", i, rss.Channel.Items[i].Description)
                 fmt.Printf("[%d] item link : %s\n\n", i, rss.Channel.Items[i].Link)
         }

 }
 */
<commit_msg>remove old notes<commit_after>package rss

import (
	"bytes"
	"encoding/xml"
	"io/ioutil"
	"net/http"
)

type Rss struct {
	Channel Channel `xml:"channel"`
}
type Item struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Enclosure		Enclosure `xml:"enclosure"`
}
type Enclosure struct {
	Url					string `xml:"url,attr"`
	Type				string `xml:"type,attr"`
}
type Channel struct {
	Title       string `xml:"title"`
	Link        string `xml:"link"`
	Description string `xml:"description"`
	Items       []Item `xml:"item"`
}

func ResponseToRss(r *http.Response) (Rss, error) {
	rss := Rss{}
	XMLdata, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return rss, err
	}
	buffer := bytes.NewBuffer(XMLdata)
	decoded := xml.NewDecoder(buffer)
	err = decoded.Decode(&rss)
	return rss, err
}
<|endoftext|>


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

package main

import (
	"bufio"
	"net"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/osrg/gobgp/packet"
)

const rtrProtocolVersion uint8 = 0

type rtrConn struct {
	conn       *net.TCPConn
	sessionId  uint16
	remoteAddr net.Addr
}

type rtrServer struct {
	connCh     chan *rtrConn
	listenPort int
}

func newRTRServer(port int) *rtrServer {
	s := &rtrServer{
		connCh:     make(chan *rtrConn, 1),
		listenPort: port,
	}
	return s
}

func (s *rtrServer) run() {
	service := ":" + strconv.Itoa(s.listenPort)
	addr, _ := net.ResolveTCPAddr("tcp", service)

	l, err := net.ListenTCP("tcp", addr)
	checkError(err)

	for i := 0; ; {
		conn, err := l.AcceptTCP()
		if err != nil {
			continue
		}
		i++
		c := &rtrConn{
			conn:       conn,
			sessionId:  uint16(i),
			remoteAddr: conn.RemoteAddr(),
		}
		s.connCh <- c
	}
}

func (rtr *rtrConn) sendPDU(msg bgp.RTRMessage) error {
	pdu, _ := msg.Serialize()
	_, err := rtr.conn.Write(pdu)
	if err != nil {
		return err
	}
	return nil
}

func (rtr *rtrConn) cacheResponse(currentSN uint32, lists FakeROATable) error {
	if err := rtr.sendPDU(bgp.NewRTRCacheResponse(rtr.sessionId)); err != nil {
		return err
	}
	log.Infof("Sent Cache Response PDU to %v (ID: %v)", rtr.remoteAddr, rtr.sessionId)

	var counter uint32
	for _, rf := range []bgp.RouteFamily{bgp.RF_IPv4_UC, bgp.RF_IPv6_UC} {
		for _, flag := range []uint8{bgp.ANNOUNCEMENT, bgp.WITHDRAWAL} {
			counter = 0
			for _, v := range lists[rf][flag] {
				if err := rtr.sendPDU(bgp.NewRTRIPPrefix(v.Prefix, v.PrefixLen, v.MaxLen, v.AS, flag)); err != nil {
					return err
				}
				counter++
				log.Debugf("Sent %s Prefix PDU to %v (Prefix: %v/%v, Maxlen: %v, AS: %v, flags: %v)", RFToIPVer(rf), rtr.remoteAddr, v.Prefix, v.PrefixLen, v.MaxLen, v.AS, flag)
			}
			if !commandOpts.Debug && counter != 0 {
				log.Infof("Sent %s Prefix PDU(s) to %v (%d ROA(s), flags: %v)", RFToIPVer(rf), rtr.remoteAddr, counter, flag)
			}
		}
	}

	if err := rtr.sendPDU(bgp.NewRTREndOfData(rtr.sessionId, currentSN)); err != nil {
		return err
	}
	log.Infof("Sent End of Data PDU to %v (ID: %v, SN: %v)", rtr.remoteAddr, rtr.sessionId, currentSN)

	return nil
}

func (rtr *rtrConn) noIncrementalUpdateAvailable() error {
	if err := rtr.sendPDU(bgp.NewRTRCacheReset()); err != nil {
		return err
	}
	log.Infof("Sent Cache Reset PDU to %v", rtr.remoteAddr)

	return nil
}

func (rtr *rtrConn) cacheHasNoDataAvailable() error {
	if err := rtr.sendPDU(bgp.NewRTRErrorReport(bgp.NO_DATA_AVAILABLE, nil, nil)); err != nil {
		return err
	}
	log.Infof("Sent Error Report PDU to %v (ID: %v, ErrorCode: %v)", rtr.remoteAddr, rtr.sessionId, bgp.NO_DATA_AVAILABLE)

	return nil
}

func RFToIPVer(rf bgp.RouteFamily) string {
	return strings.Split(rf.String(), "_")[1]
}

type errMsg struct {
	code uint16
	data []byte
}

type resourceResponse struct {
	sn   uint32
	list FakeROATable
}

func handleRTR(rtr *rtrConn, mgr *ResourceManager) {
	bcastReceiver := mgr.serialNotify.Join()
	scanner := bufio.NewScanner(bufio.NewReader(rtr.conn))
	scanner.Split(bgp.SplitRTR)

	msgCh := make(chan bgp.RTRMessage, 1)
	errCh := make(chan *errMsg, 1)
	go func() {
		defer func() {
			log.Infof("Connection to %v was closed. (ID: %v)", rtr.remoteAddr, rtr.sessionId)
			rtr.conn.Close()
		}()

		for scanner.Scan() {
			buf := scanner.Bytes()
			if buf[0] != rtrProtocolVersion {
				errCh <- &errMsg{code: bgp.UNSUPPORTED_PROTOCOL_VERSION, data: buf}
				continue
			}
			m, err := bgp.ParseRTR(buf)
			if err != nil {
				errCh <- &errMsg{code: bgp.INVALID_REQUEST, data: buf}
				continue
			}
			msgCh <- m
		}
	}()

LOOP:
	for {
		select {
		case <-bcastReceiver.In:
			currentSN := mgr.CurrentSerial()
			if err := rtr.sendPDU(bgp.NewRTRSerialNotify(rtr.sessionId, currentSN)); err != nil {
				break LOOP
			}
			log.Infof("Sent Serial Notify PDU to %v (ID: %v, SN: %v)", rtr.remoteAddr, rtr.sessionId, currentSN)
		case msg := <-errCh:
			rtr.sendPDU(bgp.NewRTRErrorReport(msg.code, msg.data, nil))
			log.Infof("Sent Error Report PDU to %v (ID: %v, ErrorCode: %v)", rtr.remoteAddr, rtr.sessionId, msg.code)
			return
		case m := <-msgCh:
			switch msg := m.(type) {
			case *bgp.RTRSerialQuery:
				peerSN := msg.SerialNumber
				log.Infof("Received Serial Query PDU from %v (ID: %v, SN: %d)", rtr.remoteAddr, msg.SessionID, peerSN)

				timeoutCh := make(chan bool, 1)
				resourceResponseCh := make(chan *resourceResponse, 1)

				go func(tCh chan bool) {
					time.Sleep(10 * time.Second)
					tCh <- true
				}(timeoutCh)

				go func(rrCh chan *resourceResponse, peerSN uint32) {
					trans := mgr.BeginTransaction()
					defer trans.EndTransaction()
					if trans.HasKey(peerSN) {
						rrCh <- &resourceResponse{
							sn:   trans.CurrentSerial(),
							list: trans.DeltaList(peerSN),
						}
					} else {
						rrCh <- nil
					}
				}(resourceResponseCh, peerSN)

				select {
				case rr := <-resourceResponseCh:
					if rr != nil {
						if err := rtr.cacheResponse(rr.sn, rr.list); err == nil {
							continue
						}
					} else {
						if err := rtr.noIncrementalUpdateAvailable(); err == nil {
							continue
						}
					}
				case <-timeoutCh:
					if err := rtr.cacheHasNoDataAvailable(); err == nil {
						continue
					}
				}
				break LOOP
			case *bgp.RTRResetQuery:
				log.Infof("Received Reset Query PDU from %v", rtr.remoteAddr)

				timeoutCh := make(chan bool, 1)
				resourceResponseCh := make(chan *resourceResponse, 1)

				go func(tCh chan bool) {
					time.Sleep(10 * time.Second)
					tCh <- true
				}(timeoutCh)

				go func(rrCh chan *resourceResponse) {
					trans := mgr.BeginTransaction()
					defer trans.EndTransaction()
					rrCh <- &resourceResponse{
						sn:   trans.CurrentSerial(),
						list: trans.CurrentList(),
					}
				}(resourceResponseCh)

				select {
				case rr := <-resourceResponseCh:
					if err := rtr.cacheResponse(rr.sn, rr.list); err == nil {
						continue
					}
				case <-timeoutCh:
					if err := rtr.cacheHasNoDataAvailable(); err == nil {
						continue
					}
				}

				break LOOP
			case *bgp.RTRErrorReport:
				log.Warnf("Received Error Report PDU from %v (%#v)", rtr.remoteAddr, msg)
				return
			default:
				pdu, _ := msg.Serialize()
				log.Warnf("Received unsupported PDU (type %d) from %v (%#v)", pdu[1], rtr.remoteAddr, msg)
				rtr.sendPDU(bgp.NewRTRErrorReport(bgp.UNSUPPORTED_PDU_TYPE, pdu, nil))
				return
			}
		}
	}
	rtr.sendPDU(bgp.NewRTRErrorReport(bgp.INTERNAL_ERROR, nil, nil))
	return
}
<commit_msg>Use length of slice instead of counter to count prefixes<commit_after>// Copyright (C) 2015 Eiichiro Watanabe
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
// implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"bufio"
	"net"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/osrg/gobgp/packet"
)

const rtrProtocolVersion uint8 = 0

type rtrConn struct {
	conn       *net.TCPConn
	sessionId  uint16
	remoteAddr net.Addr
}

type rtrServer struct {
	connCh     chan *rtrConn
	listenPort int
}

func newRTRServer(port int) *rtrServer {
	s := &rtrServer{
		connCh:     make(chan *rtrConn, 1),
		listenPort: port,
	}
	return s
}

func (s *rtrServer) run() {
	service := ":" + strconv.Itoa(s.listenPort)
	addr, _ := net.ResolveTCPAddr("tcp", service)

	l, err := net.ListenTCP("tcp", addr)
	checkError(err)

	for i := 0; ; {
		conn, err := l.AcceptTCP()
		if err != nil {
			continue
		}
		i++
		c := &rtrConn{
			conn:       conn,
			sessionId:  uint16(i),
			remoteAddr: conn.RemoteAddr(),
		}
		s.connCh <- c
	}
}

func (rtr *rtrConn) sendPDU(msg bgp.RTRMessage) error {
	pdu, _ := msg.Serialize()
	_, err := rtr.conn.Write(pdu)
	if err != nil {
		return err
	}
	return nil
}

func (rtr *rtrConn) cacheResponse(currentSN uint32, lists FakeROATable) error {
	if err := rtr.sendPDU(bgp.NewRTRCacheResponse(rtr.sessionId)); err != nil {
		return err
	}
	log.Infof("Sent Cache Response PDU to %v (ID: %v)", rtr.remoteAddr, rtr.sessionId)

	for _, rf := range []bgp.RouteFamily{bgp.RF_IPv4_UC, bgp.RF_IPv6_UC} {
		for _, flag := range []uint8{bgp.ANNOUNCEMENT, bgp.WITHDRAWAL} {
			for _, v := range lists[rf][flag] {
				if err := rtr.sendPDU(bgp.NewRTRIPPrefix(v.Prefix, v.PrefixLen, v.MaxLen, v.AS, flag)); err != nil {
					return err
				}
				log.Debugf("Sent %s Prefix PDU to %v (Prefix: %v/%v, Maxlen: %v, AS: %v, flags: %v)", RFToIPVer(rf), rtr.remoteAddr, v.Prefix, v.PrefixLen, v.MaxLen, v.AS, flag)
			}
			prefixes := len(lists[rf][flag])
			if !commandOpts.Debug && prefixes != 0 {
				log.Infof("Sent %s Prefix PDU(s) to %v (%d ROA(s), flags: %v)", RFToIPVer(rf), rtr.remoteAddr, prefixes, flag)
			}
		}
	}

	if err := rtr.sendPDU(bgp.NewRTREndOfData(rtr.sessionId, currentSN)); err != nil {
		return err
	}
	log.Infof("Sent End of Data PDU to %v (ID: %v, SN: %v)", rtr.remoteAddr, rtr.sessionId, currentSN)

	return nil
}

func (rtr *rtrConn) noIncrementalUpdateAvailable() error {
	if err := rtr.sendPDU(bgp.NewRTRCacheReset()); err != nil {
		return err
	}
	log.Infof("Sent Cache Reset PDU to %v", rtr.remoteAddr)

	return nil
}

func (rtr *rtrConn) cacheHasNoDataAvailable() error {
	if err := rtr.sendPDU(bgp.NewRTRErrorReport(bgp.NO_DATA_AVAILABLE, nil, nil)); err != nil {
		return err
	}
	log.Infof("Sent Error Report PDU to %v (ID: %v, ErrorCode: %v)", rtr.remoteAddr, rtr.sessionId, bgp.NO_DATA_AVAILABLE)

	return nil
}

func RFToIPVer(rf bgp.RouteFamily) string {
	return strings.Split(rf.String(), "_")[1]
}

type errMsg struct {
	code uint16
	data []byte
}

type resourceResponse struct {
	sn   uint32
	list FakeROATable
}

func handleRTR(rtr *rtrConn, mgr *ResourceManager) {
	bcastReceiver := mgr.serialNotify.Join()
	scanner := bufio.NewScanner(bufio.NewReader(rtr.conn))
	scanner.Split(bgp.SplitRTR)

	msgCh := make(chan bgp.RTRMessage, 1)
	errCh := make(chan *errMsg, 1)
	go func() {
		defer func() {
			log.Infof("Connection to %v was closed. (ID: %v)", rtr.remoteAddr, rtr.sessionId)
			rtr.conn.Close()
		}()

		for scanner.Scan() {
			buf := scanner.Bytes()
			if buf[0] != rtrProtocolVersion {
				errCh <- &errMsg{code: bgp.UNSUPPORTED_PROTOCOL_VERSION, data: buf}
				continue
			}
			m, err := bgp.ParseRTR(buf)
			if err != nil {
				errCh <- &errMsg{code: bgp.INVALID_REQUEST, data: buf}
				continue
			}
			msgCh <- m
		}
	}()

LOOP:
	for {
		select {
		case <-bcastReceiver.In:
			currentSN := mgr.CurrentSerial()
			if err := rtr.sendPDU(bgp.NewRTRSerialNotify(rtr.sessionId, currentSN)); err != nil {
				break LOOP
			}
			log.Infof("Sent Serial Notify PDU to %v (ID: %v, SN: %v)", rtr.remoteAddr, rtr.sessionId, currentSN)
		case msg := <-errCh:
			rtr.sendPDU(bgp.NewRTRErrorReport(msg.code, msg.data, nil))
			log.Infof("Sent Error Report PDU to %v (ID: %v, ErrorCode: %v)", rtr.remoteAddr, rtr.sessionId, msg.code)
			return
		case m := <-msgCh:
			switch msg := m.(type) {
			case *bgp.RTRSerialQuery:
				peerSN := msg.SerialNumber
				log.Infof("Received Serial Query PDU from %v (ID: %v, SN: %d)", rtr.remoteAddr, msg.SessionID, peerSN)

				timeoutCh := make(chan bool, 1)
				resourceResponseCh := make(chan *resourceResponse, 1)

				go func(tCh chan bool) {
					time.Sleep(10 * time.Second)
					tCh <- true
				}(timeoutCh)

				go func(rrCh chan *resourceResponse, peerSN uint32) {
					trans := mgr.BeginTransaction()
					defer trans.EndTransaction()
					if trans.HasKey(peerSN) {
						rrCh <- &resourceResponse{
							sn:   trans.CurrentSerial(),
							list: trans.DeltaList(peerSN),
						}
					} else {
						rrCh <- nil
					}
				}(resourceResponseCh, peerSN)

				select {
				case rr := <-resourceResponseCh:
					if rr != nil {
						if err := rtr.cacheResponse(rr.sn, rr.list); err == nil {
							continue
						}
					} else {
						if err := rtr.noIncrementalUpdateAvailable(); err == nil {
							continue
						}
					}
				case <-timeoutCh:
					if err := rtr.cacheHasNoDataAvailable(); err == nil {
						continue
					}
				}
				break LOOP
			case *bgp.RTRResetQuery:
				log.Infof("Received Reset Query PDU from %v", rtr.remoteAddr)

				timeoutCh := make(chan bool, 1)
				resourceResponseCh := make(chan *resourceResponse, 1)

				go func(tCh chan bool) {
					time.Sleep(10 * time.Second)
					tCh <- true
				}(timeoutCh)

				go func(rrCh chan *resourceResponse) {
					trans := mgr.BeginTransaction()
					defer trans.EndTransaction()
					rrCh <- &resourceResponse{
						sn:   trans.CurrentSerial(),
						list: trans.CurrentList(),
					}
				}(resourceResponseCh)

				select {
				case rr := <-resourceResponseCh:
					if err := rtr.cacheResponse(rr.sn, rr.list); err == nil {
						continue
					}
				case <-timeoutCh:
					if err := rtr.cacheHasNoDataAvailable(); err == nil {
						continue
					}
				}

				break LOOP
			case *bgp.RTRErrorReport:
				log.Warnf("Received Error Report PDU from %v (%#v)", rtr.remoteAddr, msg)
				return
			default:
				pdu, _ := msg.Serialize()
				log.Warnf("Received unsupported PDU (type %d) from %v (%#v)", pdu[1], rtr.remoteAddr, msg)
				rtr.sendPDU(bgp.NewRTRErrorReport(bgp.UNSUPPORTED_PDU_TYPE, pdu, nil))
				return
			}
		}
	}
	rtr.sendPDU(bgp.NewRTRErrorReport(bgp.INTERNAL_ERROR, nil, nil))
	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"io"
	"log"
	"net"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/remind101/emp/Godeps/_workspace/src/github.com/bgentry/heroku-go"
	"github.com/remind101/emp/Godeps/_workspace/src/github.com/docker/docker/pkg/term"
)

var (
	detachedRun bool
	dynoSize    string
)

var cmdRun = &Command{
	Run:      runRun,
	Usage:    "run [-s <size>] [-d] <command> [<argument>...]",
	NeedsApp: true,
	Category: "dyno",
	Short:    "run a process in a dyno",
	Long: `
Run a process on Heroku. Flags such as` + " `-a` " + `may be parsed out of
the command unless the command is quoted or provided after a
double-dash (--).

Options:

    -s <size>  set the size for this dyno (e.g. 2X)
    -d         run in detached mode instead of attached to terminal

Examples:

    $ emp run echo "hello"
    Running ` + "`echo \"hello\"`" + ` on myapp as run.1234:
    "hello"

    $ emp run console
    Running ` + "`console`" + ` on myapp as run.5678:
    Loading production environment (Rails 3.2.14)
    irb(main):001:0> ...

    $ emp run -d -s 2X bin/my_worker
    Ran ` + "`bin/my_worker`" + ` on myapp as run.4321, detached.

    $ emp run -a myapp -- ls -a /
    Running ` + "`ls -a bin /`" + ` on myapp as run.8650:
    /:
    .  ..  app  bin  dev  etc  home  lib  lib64  lost+found  proc  sbin  tmp  usr  var
`,
}

func init() {
	cmdRun.Flag.BoolVarP(&detachedRun, "detached", "d", false, "detached")
	cmdRun.Flag.StringVarP(&dynoSize, "size", "s", "", "dyno size")
}

func runRun(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.PrintUsage()
		os.Exit(2)
	}
	appname := mustApp()

	w, err := term.GetWinsize(inFd)
	if err != nil {
		// If syscall.TIOCGWINSZ is not supported by the device, we're
		// probably trying to run tests. Set w to some sensible default.
		if err.Error() == "operation not supported by device" {
			w = &term.Winsize{
				Height: 20,
				Width:  80,
			}
		} else {
			printFatal(err.Error())
		}
	}

	attached := !detachedRun
	opts := heroku.DynoCreateOpts{Attach: &attached}
	if attached {
		env := map[string]string{
			"COLUMNS": strconv.Itoa(int(w.Width)),
			"LINES":   strconv.Itoa(int(w.Height)),
			"TERM":    os.Getenv("TERM"),
		}
		opts.Env = &env
	}
	if dynoSize != "" {
		if !strings.HasSuffix(dynoSize, "X") {
			cmd.PrintUsage()
			os.Exit(2)
		}
		opts.Size = &dynoSize
	}

	command := strings.Join(args, " ")
	if detachedRun {
		dyno, err := client.DynoCreate(appname, command, &opts)
		must(err)

		log.Printf("Ran `%s` on %s as %s, detached.", dyno.Command, appname, dyno.Name)
		return
	}

	params := struct {
		Command string             `json:"command"`
		Attach  *bool              `json:"attach,omitempty"`
		Env     *map[string]string `json:"env,omitempty"`
		Size    *string            `json:"size,omitempty"`
	}{
		Command: command,
		Attach:  opts.Attach,
		Env:     opts.Env,
		Size:    opts.Size,
	}
	req, err := client.NewRequest("POST", "/apps/"+appname+"/dynos", params)
	must(err)

	u, err := url.Parse(apiURL)
	must(err)

	protocol := u.Scheme
	address := u.Path
	if protocol != "unix" {
		protocol = "tcp"
		address = u.Host + ":80"
	}

	if u.Scheme == "https" {
		address = address + ":443"
	}

	var dial net.Conn
	if u.Scheme == "https" {
		dial, err = tlsDial(protocol, address, &tls.Config{})
		if err != nil {
			printFatal(err.Error())
		}
	} else {
		dial, err = net.Dial(protocol, address)
		if err != nil {
			printFatal(err.Error())
		}
	}

	clientconn := httputil.NewClientConn(dial, nil)
	defer clientconn.Close()
	_, err = clientconn.Do(req)
	if err != nil && err != httputil.ErrPersistEOF {
		printFatal(err.Error())
	}
	rwc, br := clientconn.Hijack()
	defer rwc.Close()

	if isTerminalIn && isTerminalOut {
		state, err := term.SetRawTerminal(inFd)
		if err != nil {
			printFatal(err.Error())
		}
		defer term.RestoreTerminal(inFd, state)
	}

	errChanOut := make(chan error, 1)
	errChanIn := make(chan error, 1)
	exit := make(chan bool)
	go func() {
		defer close(exit)
		defer close(errChanOut)
		var err error
		_, err = io.Copy(os.Stdout, br)
		errChanOut <- err
	}()
	go func() {
		_, err := io.Copy(rwc, os.Stdin)
		errChanIn <- err
		rwc.(interface {
			CloseWrite() error
		}).CloseWrite()
	}()
	<-exit
	select {
	case err = <-errChanIn:
		must(err)
	case err = <-errChanOut:
		must(err)
	}
}
<commit_msg>Fix double port.<commit_after>package main

import (
	"crypto/tls"
	"io"
	"log"
	"net"
	"net/http/httputil"
	"net/url"
	"os"
	"strconv"
	"strings"

	"github.com/remind101/emp/Godeps/_workspace/src/github.com/bgentry/heroku-go"
	"github.com/remind101/emp/Godeps/_workspace/src/github.com/docker/docker/pkg/term"
)

var (
	detachedRun bool
	dynoSize    string
)

var cmdRun = &Command{
	Run:      runRun,
	Usage:    "run [-s <size>] [-d] <command> [<argument>...]",
	NeedsApp: true,
	Category: "dyno",
	Short:    "run a process in a dyno",
	Long: `
Run a process on Heroku. Flags such as` + " `-a` " + `may be parsed out of
the command unless the command is quoted or provided after a
double-dash (--).

Options:

    -s <size>  set the size for this dyno (e.g. 2X)
    -d         run in detached mode instead of attached to terminal

Examples:

    $ emp run echo "hello"
    Running ` + "`echo \"hello\"`" + ` on myapp as run.1234:
    "hello"

    $ emp run console
    Running ` + "`console`" + ` on myapp as run.5678:
    Loading production environment (Rails 3.2.14)
    irb(main):001:0> ...

    $ emp run -d -s 2X bin/my_worker
    Ran ` + "`bin/my_worker`" + ` on myapp as run.4321, detached.

    $ emp run -a myapp -- ls -a /
    Running ` + "`ls -a bin /`" + ` on myapp as run.8650:
    /:
    .  ..  app  bin  dev  etc  home  lib  lib64  lost+found  proc  sbin  tmp  usr  var
`,
}

func init() {
	cmdRun.Flag.BoolVarP(&detachedRun, "detached", "d", false, "detached")
	cmdRun.Flag.StringVarP(&dynoSize, "size", "s", "", "dyno size")
}

func runRun(cmd *Command, args []string) {
	if len(args) == 0 {
		cmd.PrintUsage()
		os.Exit(2)
	}
	appname := mustApp()

	w, err := term.GetWinsize(inFd)
	if err != nil {
		// If syscall.TIOCGWINSZ is not supported by the device, we're
		// probably trying to run tests. Set w to some sensible default.
		if err.Error() == "operation not supported by device" {
			w = &term.Winsize{
				Height: 20,
				Width:  80,
			}
		} else {
			printFatal(err.Error())
		}
	}

	attached := !detachedRun
	opts := heroku.DynoCreateOpts{Attach: &attached}
	if attached {
		env := map[string]string{
			"COLUMNS": strconv.Itoa(int(w.Width)),
			"LINES":   strconv.Itoa(int(w.Height)),
			"TERM":    os.Getenv("TERM"),
		}
		opts.Env = &env
	}
	if dynoSize != "" {
		if !strings.HasSuffix(dynoSize, "X") {
			cmd.PrintUsage()
			os.Exit(2)
		}
		opts.Size = &dynoSize
	}

	command := strings.Join(args, " ")
	if detachedRun {
		dyno, err := client.DynoCreate(appname, command, &opts)
		must(err)

		log.Printf("Ran `%s` on %s as %s, detached.", dyno.Command, appname, dyno.Name)
		return
	}

	params := struct {
		Command string             `json:"command"`
		Attach  *bool              `json:"attach,omitempty"`
		Env     *map[string]string `json:"env,omitempty"`
		Size    *string            `json:"size,omitempty"`
	}{
		Command: command,
		Attach:  opts.Attach,
		Env:     opts.Env,
		Size:    opts.Size,
	}
	req, err := client.NewRequest("POST", "/apps/"+appname+"/dynos", params)
	must(err)

	u, err := url.Parse(apiURL)
	must(err)

	protocol := u.Scheme
	address := u.Path
	if protocol != "unix" {
		protocol = "tcp"
		address = u.Host + ":80"
	}

	if u.Scheme == "https" {
		address = u.Host + ":443"
	}

	var dial net.Conn
	if u.Scheme == "https" {
		dial, err = tlsDial(protocol, address, &tls.Config{})
		if err != nil {
			printFatal(err.Error())
		}
	} else {
		dial, err = net.Dial(protocol, address)
		if err != nil {
			printFatal(err.Error())
		}
	}

	clientconn := httputil.NewClientConn(dial, nil)
	defer clientconn.Close()
	_, err = clientconn.Do(req)
	if err != nil && err != httputil.ErrPersistEOF {
		printFatal(err.Error())
	}
	rwc, br := clientconn.Hijack()
	defer rwc.Close()

	if isTerminalIn && isTerminalOut {
		state, err := term.SetRawTerminal(inFd)
		if err != nil {
			printFatal(err.Error())
		}
		defer term.RestoreTerminal(inFd, state)
	}

	errChanOut := make(chan error, 1)
	errChanIn := make(chan error, 1)
	exit := make(chan bool)
	go func() {
		defer close(exit)
		defer close(errChanOut)
		var err error
		_, err = io.Copy(os.Stdout, br)
		errChanOut <- err
	}()
	go func() {
		_, err := io.Copy(rwc, os.Stdin)
		errChanIn <- err
		rwc.(interface {
			CloseWrite() error
		}).CloseWrite()
	}()
	<-exit
	select {
	case err = <-errChanIn:
		must(err)
	case err = <-errChanOut:
		must(err)
	}
}
<|endoftext|>


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

package elastigo

import (
	"fmt"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	hostpool "github.com/bitly/go-hostpool"
)

const (
	Version         = "0.0.2"
	DefaultProtocol = "http"
	DefaultDomain   = "localhost"
	DefaultPort     = "9200"
	// A decay duration of zero results in the default behaviour
	DefaultDecayDuration = 0
)

type Conn struct {
	// Maintain these for backwards compatibility
	Protocol       string
	Domain         string
	ClusterDomains []string
	Port           string
	Username       string
	Password       string
	Hosts          []string
	Client         *http.Client
	hp             hostpool.HostPool
	once           sync.Once

	// To compute the weighting scores, we perform a weighted average of recent response times,
	// over the course of `DecayDuration`. DecayDuration may be set to 0 to use the default
	// value of 5 minutes. The EpsilonValueCalculator uses this to calculate a score
	// from the weighted average response time.
	DecayDuration time.Duration
}

func NewConn() *Conn {
	return &Conn{
		// Maintain these for backwards compatibility
		Protocol:       DefaultProtocol,
		Domain:         DefaultDomain,
		ClusterDomains: []string{DefaultDomain},
		Port:           DefaultPort,
		DecayDuration:  time.Duration(DefaultDecayDuration * time.Second),
		Client:         http.DefaultClient,
	}
}

func (c *Conn) SetPort(port string) {
	c.Port = port
}

func (c *Conn) SetHosts(newhosts []string) {

	// Store the new host list
	c.Hosts = newhosts

	// Reinitialise the host pool Pretty naive as this will nuke the current
	// hostpool, and therefore reset any scoring
	c.initializeHostPool()
}

// Set up the host pool to be used
func (c *Conn) initializeHostPool() {

	// If no hosts are set, fallback to defaults
	if len(c.Hosts) == 0 {
		c.Hosts = append(c.Hosts, fmt.Sprintf("%s:%s", c.Domain, c.Port))
	}

	// Epsilon Greedy is an algorithm that allows HostPool not only to
	// track failure state, but also to learn about "better" options in
	// terms of speed, and to pick from available hosts based on how well
	// they perform. This gives a weighted request rate to better
	// performing hosts, while still distributing requests to all hosts
	// (proportionate to their performance).  The interface is the same as
	// the standard HostPool, but be sure to mark the HostResponse
	// immediately after executing the request to the host, as that will
	// stop the implicitly running request timer.
	//
	// A good overview of Epsilon Greedy is here http://stevehanov.ca/blog/index.php?id=132
	if c.hp != nil {
		c.hp.Close()
	}
	c.hp = hostpool.NewEpsilonGreedy(
		c.Hosts, c.DecayDuration, &hostpool.LinearEpsilonValueCalculator{})
}

func (c *Conn) Close() {
	c.hp.Close()
}

func (c *Conn) NewRequest(method, path, query string) (*Request, error) {
	// Setup the hostpool on our first run
	c.once.Do(c.initializeHostPool)

	// Get a host from the host pool
	hr := c.hp.Get()

	// Get the final host and port
	host, portNum := splitHostnamePartsFromHost(hr.Host(), c.Port)

	// Build request
	var uri string
	// If query parameters are provided, the add them to the URL,
	// otherwise, leave them out
	if len(query) > 0 {
		uri = fmt.Sprintf("%s://%s:%s%s?%s", c.Protocol, host, portNum, path, query)
	} else {
		uri = fmt.Sprintf("%s://%s:%s%s", c.Protocol, host, portNum, path)
	}
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Accept", "application/json")
	req.Header.Add("User-Agent", "elasticSearch/"+Version+" ("+runtime.GOOS+"-"+runtime.GOARCH+")")

	if c.Username != "" || c.Password != "" {
		req.SetBasicAuth(c.Username, c.Password)
	}

	newRequest := &Request{
		Request:      req,
		hostResponse: hr,
		client:       c.Client,
	}
	return newRequest, nil
}

// Split apart the hostname on colon
// Return the host and a default port if there is no separator
func splitHostnamePartsFromHost(fullHost string, defaultPortNum string) (string, string) {

	h := strings.Split(fullHost, ":")

	if len(h) == 2 {
		return h[0], h[1]
	}

	return h[0], defaultPortNum
}
<commit_msg>Add content type header for requests<commit_after>// Copyright 2013 Matthew Baird
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package elastigo

import (
	"fmt"
	"net/http"
	"runtime"
	"strings"
	"sync"
	"time"

	hostpool "github.com/bitly/go-hostpool"
)

const (
	Version         = "0.0.2"
	DefaultProtocol = "http"
	DefaultDomain   = "localhost"
	DefaultPort     = "9200"
	// A decay duration of zero results in the default behaviour
	DefaultDecayDuration = 0
)

type Conn struct {
	// Maintain these for backwards compatibility
	Protocol       string
	Domain         string
	ClusterDomains []string
	Port           string
	Username       string
	Password       string
	Hosts          []string
	Client         *http.Client
	hp             hostpool.HostPool
	once           sync.Once

	// To compute the weighting scores, we perform a weighted average of recent response times,
	// over the course of `DecayDuration`. DecayDuration may be set to 0 to use the default
	// value of 5 minutes. The EpsilonValueCalculator uses this to calculate a score
	// from the weighted average response time.
	DecayDuration time.Duration
}

func NewConn() *Conn {
	return &Conn{
		// Maintain these for backwards compatibility
		Protocol:       DefaultProtocol,
		Domain:         DefaultDomain,
		ClusterDomains: []string{DefaultDomain},
		Port:           DefaultPort,
		DecayDuration:  time.Duration(DefaultDecayDuration * time.Second),
		Client:         http.DefaultClient,
	}
}

func (c *Conn) SetPort(port string) {
	c.Port = port
}

func (c *Conn) SetHosts(newhosts []string) {

	// Store the new host list
	c.Hosts = newhosts

	// Reinitialise the host pool Pretty naive as this will nuke the current
	// hostpool, and therefore reset any scoring
	c.initializeHostPool()
}

// Set up the host pool to be used
func (c *Conn) initializeHostPool() {

	// If no hosts are set, fallback to defaults
	if len(c.Hosts) == 0 {
		c.Hosts = append(c.Hosts, fmt.Sprintf("%s:%s", c.Domain, c.Port))
	}

	// Epsilon Greedy is an algorithm that allows HostPool not only to
	// track failure state, but also to learn about "better" options in
	// terms of speed, and to pick from available hosts based on how well
	// they perform. This gives a weighted request rate to better
	// performing hosts, while still distributing requests to all hosts
	// (proportionate to their performance).  The interface is the same as
	// the standard HostPool, but be sure to mark the HostResponse
	// immediately after executing the request to the host, as that will
	// stop the implicitly running request timer.
	//
	// A good overview of Epsilon Greedy is here http://stevehanov.ca/blog/index.php?id=132
	if c.hp != nil {
		c.hp.Close()
	}
	c.hp = hostpool.NewEpsilonGreedy(
		c.Hosts, c.DecayDuration, &hostpool.LinearEpsilonValueCalculator{})
}

func (c *Conn) Close() {
	c.hp.Close()
}

func (c *Conn) NewRequest(method, path, query string) (*Request, error) {
	// Setup the hostpool on our first run
	c.once.Do(c.initializeHostPool)

	// Get a host from the host pool
	hr := c.hp.Get()

	// Get the final host and port
	host, portNum := splitHostnamePartsFromHost(hr.Host(), c.Port)

	// Build request
	var uri string
	// If query parameters are provided, the add them to the URL,
	// otherwise, leave them out
	if len(query) > 0 {
		uri = fmt.Sprintf("%s://%s:%s%s?%s", c.Protocol, host, portNum, path, query)
	} else {
		uri = fmt.Sprintf("%s://%s:%s%s", c.Protocol, host, portNum, path)
	}
	req, err := http.NewRequest(method, uri, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Add("Accept", "application/json")
	req.Header.Add("Content-Type", "application/json")
	req.Header.Add("User-Agent", "elasticSearch/"+Version+" ("+runtime.GOOS+"-"+runtime.GOARCH+")")

	if c.Username != "" || c.Password != "" {
		req.SetBasicAuth(c.Username, c.Password)
	}

	newRequest := &Request{
		Request:      req,
		hostResponse: hr,
		client:       c.Client,
	}
	return newRequest, nil
}

// Split apart the hostname on colon
// Return the host and a default port if there is no separator
func splitHostnamePartsFromHost(fullHost string, defaultPortNum string) (string, string) {

	h := strings.Split(fullHost, ":")

	if len(h) == 2 {
		return h[0], h[1]
	}

	return h[0], defaultPortNum
}
<|endoftext|>