nebula/timeout_system_test.go

package nebula

import (
	"net"
	"testing"
	"time"

	"github.com/slackhq/nebula/iputil"
	"github.com/stretchr/testify/assert"
)

func TestNewSystemTimerWheel(t *testing.T) {
	// Make sure we get an object we expect
	tw := NewSystemTimerWheel(time.Second, time.Second*10)
	assert.Equal(t, 12, tw.wheelLen)
	assert.Equal(t, 0, tw.current)
	assert.Nil(t, tw.lastTick)
	assert.Equal(t, time.Second*1, tw.tickDuration)
	assert.Equal(t, time.Second*10, tw.wheelDuration)
	assert.Len(t, tw.wheel, 12)

	// Assert the math is correct
	tw = NewSystemTimerWheel(time.Second*3, time.Second*10)
	assert.Equal(t, 5, tw.wheelLen)

	tw = NewSystemTimerWheel(time.Second*120, time.Minute*10)
	assert.Equal(t, 7, tw.wheelLen)
}

func TestSystemTimerWheel_findWheel(t *testing.T) {
	tw := NewSystemTimerWheel(time.Second, time.Second*10)
	assert.Len(t, tw.wheel, 12)

	// Current + tick + 1 since we don't know how far into current we are
	assert.Equal(t, 2, tw.findWheel(time.Second*1))

	// Scale up to min duration
	assert.Equal(t, 2, tw.findWheel(time.Millisecond*1))

	// Make sure we hit that last index
	assert.Equal(t, 11, tw.findWheel(time.Second*10))

	// Scale down to max duration
	assert.Equal(t, 11, tw.findWheel(time.Second*11))

	tw.current = 1
	// Make sure we account for the current position properly
	assert.Equal(t, 3, tw.findWheel(time.Second*1))
	assert.Equal(t, 0, tw.findWheel(time.Second*10))

	// Ensure that all configurations of a wheel does not result in calculating an overflow of the wheel
	for min := time.Duration(1); min < 100; min++ {
		for max := min; max < 100; max++ {
			tw = NewSystemTimerWheel(min, max)

			for current := 0; current < tw.wheelLen; current++ {
				tw.current = current
				for timeout := time.Duration(0); timeout <= tw.wheelDuration; timeout++ {
					tick := tw.findWheel(timeout)
					if tick >= tw.wheelLen {
						t.Errorf("Min: %v; Max: %v; Wheel len: %v; Current Tick: %v; Insert timeout: %v; Calc tick: %v", min, max, tw.wheelLen, current, timeout, tick)
					}
				}
			}
		}
	}
}

func TestSystemTimerWheel_Add(t *testing.T) {
	tw := NewSystemTimerWheel(time.Second, time.Second*10)

	fp1 := iputil.Ip2VpnIp(net.ParseIP("1.2.3.4"))
	tw.Add(fp1, time.Second*1)

	// Make sure we set head and tail properly
	assert.NotNil(t, tw.wheel[2])
	assert.Equal(t, fp1, tw.wheel[2].Head.Item)
	assert.Nil(t, tw.wheel[2].Head.Next)
	assert.Equal(t, fp1, tw.wheel[2].Tail.Item)
	assert.Nil(t, tw.wheel[2].Tail.Next)

	// Make sure we only modify head
	fp2 := iputil.Ip2VpnIp(net.ParseIP("1.2.3.4"))
	tw.Add(fp2, time.Second*1)
	assert.Equal(t, fp2, tw.wheel[2].Head.Item)
	assert.Equal(t, fp1, tw.wheel[2].Head.Next.Item)
	assert.Equal(t, fp1, tw.wheel[2].Tail.Item)
	assert.Nil(t, tw.wheel[2].Tail.Next)

	// Make sure we use free'd items first
	tw.itemCache = &SystemTimeoutItem{}
	tw.itemsCached = 1
	tw.Add(fp2, time.Second*1)
	assert.Nil(t, tw.itemCache)
	assert.Equal(t, 0, tw.itemsCached)
}

func TestSystemTimerWheel_Purge(t *testing.T) {
	// First advance should set the lastTick and do nothing else
	tw := NewSystemTimerWheel(time.Second, time.Second*10)
	assert.Nil(t, tw.lastTick)
	tw.advance(time.Now())
	assert.NotNil(t, tw.lastTick)
	assert.Equal(t, 0, tw.current)

	fps := []iputil.VpnIp{9, 10, 11, 12}

	//fp1 := ip2int(net.ParseIP("1.2.3.4"))

	tw.Add(fps[0], time.Second*1)
	tw.Add(fps[1], time.Second*1)
	tw.Add(fps[2], time.Second*2)
	tw.Add(fps[3], time.Second*2)

	ta := time.Now().Add(time.Second * 3)
	lastTick := *tw.lastTick
	tw.advance(ta)
	assert.Equal(t, 3, tw.current)
	assert.True(t, tw.lastTick.After(lastTick))

	// Make sure we get all 4 packets back
	for i := 0; i < 4; i++ {
		assert.Contains(t, fps, tw.Purge())
	}

	// Make sure there aren't any leftover
	assert.Nil(t, tw.Purge())
	assert.Nil(t, tw.expired.Head)
	assert.Nil(t, tw.expired.Tail)

	// Make sure we cached the free'd items
	assert.Equal(t, 4, tw.itemsCached)
	ci := tw.itemCache
	for i := 0; i < 4; i++ {
		assert.NotNil(t, ci)
		ci = ci.Next
	}
	assert.Nil(t, ci)

	// Lets make sure we roll over properly
	ta = ta.Add(time.Second * 5)
	tw.advance(ta)
	assert.Equal(t, 8, tw.current)

	ta = ta.Add(time.Second * 2)
	tw.advance(ta)
	assert.Equal(t, 10, tw.current)

	ta = ta.Add(time.Second * 1)
	tw.advance(ta)
	assert.Equal(t, 11, tw.current)

	ta = ta.Add(time.Second * 1)
	tw.advance(ta)
	assert.Equal(t, 0, tw.current)
}