<commit_before>package main

import (
	"bytes"
	"github.com/bmizerany/assert"
	"math/rand"
	"regexp"
	"strconv"
	"testing"
	"time"
)

var commonPercentiles = Percentiles{
	&Percentile{
		99,
		"99",
	},
}

func TestPacketParse(t *testing.T) {
	d := []byte("gaugor:333|g")
	packets := parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet := packets[0]
	assert.Equal(t, "gaugor", packet.Bucket)
	assert.Equal(t, uint64(333), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:2|c|@0.1")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(2), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(0.1), packet.Sampling)

	d = []byte("gorets:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:-4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(-4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("glork:320|ms")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "glork", packet.Bucket)
	assert.Equal(t, uint64(320), packet.Value.(uint64))
	assert.Equal(t, "ms", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c\ngauge:3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 2)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	packet = packets[1]
	assert.Equal(t, "gauge", packet.Bucket)
	assert.Equal(t, uint64(3), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4\ngauge3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("a.key.with-0.dash:4")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5m")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|mg")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|ms|@")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
}

func TestMean(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("response_time:0|ms\nresponse_time:30|ms\nresponse_time:30|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()
	pattern := `response_time\.mean 20\.[0-9]+ `
	meanRegexp := regexp.MustCompile(pattern)

	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func TestUpperPercentile(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("time:0|ms\ntime:1|ms\ntime:2|ms\ntime:3|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{
		&Percentile{
			75,
			"75",
		},
	})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()

	meanRegexp := regexp.MustCompile(`time\.upper_75 2 `)
	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func TestLowerPercentile(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("time:0|ms\ntime:1|ms\ntime:2|ms\ntime:3|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{
		&Percentile{
			-75,
			"-75",
		},
	})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()

	meanRegexp := regexp.MustCompile(`time\.upper_75 1 `)
	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, false)

	meanRegexp = regexp.MustCompile(`time\.lower_75 1 `)
	matched = meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func BenchmarkManyDifferentSensors(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "count" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := int64(r.Uint32() % 1000)
			counters[bucket] = a
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "gauge" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			gauges[bucket] = a
		}
	}

	var buff bytes.Buffer
	now := time.Now().Unix()
	t.ResetTimer()
	processTimers(&buff, now, commonPercentiles)
	processCounters(&buff, now)
	processGauges(&buff, now)
}

func BenchmarkOneBigTimer(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	bucket := "response_time"
	for i := 0; i < 10000000; i++ {
		a := uint64(r.Uint32() % 1000)
		timers[bucket] = append(timers[bucket], a)
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkLotsOfTimers(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkParseMessage(b *testing.B) {
	d := []byte("a.key.with-0.dash:4|c|@0.5")
	for i := 0; i < b.N; i++ {
		parseMessage(d)
	}
}
<commit_msg>added tests for packetHandler() function<commit_after>package main

import (
	"bytes"
	"github.com/bmizerany/assert"
	"math/rand"
	"regexp"
	"strconv"
	"testing"
	"time"
)

var commonPercentiles = Percentiles{
	&Percentile{
		99,
		"99",
	},
}

func TestPacketParse(t *testing.T) {
	d := []byte("gaugor:333|g")
	packets := parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet := packets[0]
	assert.Equal(t, "gaugor", packet.Bucket)
	assert.Equal(t, uint64(333), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:2|c|@0.1")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(2), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(0.1), packet.Sampling)

	d = []byte("gorets:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:-4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(-4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("glork:320|ms")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "glork", packet.Bucket)
	assert.Equal(t, uint64(320), packet.Value.(uint64))
	assert.Equal(t, "ms", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c\ngauge:3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 2)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	packet = packets[1]
	assert.Equal(t, "gauge", packet.Bucket)
	assert.Equal(t, uint64(3), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4\ngauge3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("a.key.with-0.dash:4")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5m")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|mg")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|ms|@")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
}

func TestReceiveCounterPacketHandling(t *testing.T) {

	counters = make(map[string]int64)
	*receiveCounter = "countme"

	p := &Packet {
		Bucket: "gorets",
		Value:	int64(100),
		Modifier: "c",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, counters["countme"], int64(1))

	packetHandler(p)
	assert.Equal(t, counters["countme"], int64(2))
}

func TestCountPacketHandling(t *testing.T) {

	counters = make(map[string]int64)

	p := &Packet {
		Bucket: "gorets",
		Value:	int64(100),
		Modifier: "c",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(100))

	p.Value = int64(3)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(103))

	p.Value = int64(-4)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(99))

}

func TestGaugePacketHandling(t *testing.T) {

	gauges = make(map[string]uint64)

	p := &Packet {
		Bucket: "gaugor",
		Value:	uint64(333),
		Modifier: "g",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, gauges["gaugor"], uint64(333))

}

func TestTimerPacketHandling(t *testing.T) {

	timers = make(map[string]Uint64Slice)

	p := &Packet {
		Bucket: "glork",
		Value:	uint64(320),
		Modifier: "ms",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, len(timers["glork"]), 1)
	assert.Equal(t, timers["glork"][0], uint64(320))

	p.Value = uint64(100)
	packetHandler(p)
	assert.Equal(t, len(timers["glork"]), 2)
	assert.Equal(t, timers["glork"][1], uint64(100))
}

func TestMean(t *testing.T) {

	timers = make(map[string]Uint64Slice)

	// Some data with expected mean of 20
	d := []byte("response_time:0|ms\nresponse_time:30|ms\nresponse_time:30|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()
	pattern := `response_time\.mean 20\.[0-9]+ `
	meanRegexp := regexp.MustCompile(pattern)

	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func TestUpperPercentile(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("time:0|ms\ntime:1|ms\ntime:2|ms\ntime:3|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{
		&Percentile{
			75,
			"75",
		},
	})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()

	meanRegexp := regexp.MustCompile(`time\.upper_75 2 `)
	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func TestLowerPercentile(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("time:0|ms\ntime:1|ms\ntime:2|ms\ntime:3|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{
		&Percentile{
			-75,
			"-75",
		},
	})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()

	meanRegexp := regexp.MustCompile(`time\.upper_75 1 `)
	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, false)

	meanRegexp = regexp.MustCompile(`time\.lower_75 1 `)
	matched = meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func BenchmarkManyDifferentSensors(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "count" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := int64(r.Uint32() % 1000)
			counters[bucket] = a
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "gauge" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			gauges[bucket] = a
		}
	}

	var buff bytes.Buffer
	now := time.Now().Unix()
	t.ResetTimer()
	processTimers(&buff, now, commonPercentiles)
	processCounters(&buff, now)
	processGauges(&buff, now)
}

func BenchmarkOneBigTimer(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	bucket := "response_time"
	for i := 0; i < 10000000; i++ {
		a := uint64(r.Uint32() % 1000)
		timers[bucket] = append(timers[bucket], a)
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkLotsOfTimers(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkParseMessage(b *testing.B) {
	d := []byte("a.key.with-0.dash:4|c|@0.5")
	for i := 0; i < b.N; i++ {
		parseMessage(d)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"github.com/bmizerany/assert"
	"math/rand"
	"regexp"
	"strconv"
	"testing"
	"time"
)

var commonPercentiles = Percentiles{
	&Percentile{
		99,
		"99",
	},
}

func TestPacketParse(t *testing.T) {
	d := []byte("gaugor:333|g")
	packets := parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet := packets[0]
	assert.Equal(t, "gaugor", packet.Bucket)
	assert.Equal(t, uint64(333), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:2|c|@0.1")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(2), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(0.1), packet.Sampling)

	d = []byte("gorets:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:-4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(-4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("glork:320|ms")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "glork", packet.Bucket)
	assert.Equal(t, uint64(320), packet.Value.(uint64))
	assert.Equal(t, "ms", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c\ngauge:3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 2)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	packet = packets[1]
	assert.Equal(t, "gauge", packet.Bucket)
	assert.Equal(t, uint64(3), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4\ngauge3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("a.key.with-0.dash:4")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5m")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|mg")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|ms|@")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
}

func TestReceiveCounterPacketHandling(t *testing.T) {
	counters = make(map[string]int64)
	*receiveCounter = "countme"

	p := &Packet{
		Bucket:   "gorets",
		Value:    int64(100),
		Modifier: "c",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, counters["countme"], int64(1))

	packetHandler(p)
	assert.Equal(t, counters["countme"], int64(2))
}

func TestCountPacketHandling(t *testing.T) {
	counters = make(map[string]int64)

	p := &Packet{
		Bucket:   "gorets",
		Value:    int64(100),
		Modifier: "c",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(100))

	p.Value = int64(3)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(103))

	p.Value = int64(-4)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(99))

	p.Value = int64(-100)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(-1))
}

func TestGaugePacketHandling(t *testing.T) {
	gauges = make(map[string]uint64)

	p := &Packet{
		Bucket:   "gaugor",
		Value:    uint64(333),
		Modifier: "g",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, gauges["gaugor"], uint64(333))
}

func TestTimerPacketHandling(t *testing.T) {
	timers = make(map[string]Uint64Slice)

	p := &Packet{
		Bucket:   "glork",
		Value:    uint64(320),
		Modifier: "ms",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, len(timers["glork"]), 1)
	assert.Equal(t, timers["glork"][0], uint64(320))

	p.Value = uint64(100)
	packetHandler(p)
	assert.Equal(t, len(timers["glork"]), 2)
	assert.Equal(t, timers["glork"][1], uint64(100))
}

func TestProcessCounters(t *testing.T) {

	*persistCountKeys = int64(10)
	counters = make(map[string]int64)
	var buffer bytes.Buffer
	now := int64(1418052649)

	counters["gorets"] = int64(123)

	num := processCounters(&buffer, now)
	assert.Equal(t, num, int64(1))
	assert.Equal(t, buffer.String(), "gorets 123 1418052649\n")

	// run processCounters() enough times to make sure it purges items
	for i := 0; i < int(*persistCountKeys)+10; i++ {
		num = processCounters(&buffer, now)
	}
	lines := bytes.Split(buffer.Bytes(), []byte("\n"))

	// expect two more lines - the good one and an empty one at the end
	assert.Equal(t, len(lines), int(*persistCountKeys+2))
	assert.Equal(t, string(lines[0]), "gorets 123 1418052649")
	assert.Equal(t, string(lines[*persistCountKeys]), "gorets 0 1418052649")
}

func TestMean(t *testing.T) {
	timers = make(map[string]Uint64Slice)

	// Some data with expected mean of 20
	d := []byte("response_time:0|ms\nresponse_time:30|ms\nresponse_time:30|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()
	pattern := `response_time\.mean 20\.[0-9]+ `
	meanRegexp := regexp.MustCompile(pattern)

	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func TestUpperPercentile(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("time:0|ms\ntime:1|ms\ntime:2|ms\ntime:3|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{
		&Percentile{
			75,
			"75",
		},
	})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()

	meanRegexp := regexp.MustCompile(`time\.upper_75 2 `)
	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func TestLowerPercentile(t *testing.T) {
	// Some data with expected mean of 20
	d := []byte("time:0|ms\ntime:1|ms\ntime:2|ms\ntime:3|ms")
	packets := parseMessage(d)

	for _, s := range packets {
		timers[s.Bucket] = append(timers[s.Bucket], s.Value.(uint64))
	}

	var buff bytes.Buffer
	var num int64
	num += processTimers(&buff, time.Now().Unix(), Percentiles{
		&Percentile{
			-75,
			"-75",
		},
	})
	assert.Equal(t, num, int64(1))
	dataForGraphite := buff.String()

	meanRegexp := regexp.MustCompile(`time\.upper_75 1 `)
	matched := meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, false)

	meanRegexp = regexp.MustCompile(`time\.lower_75 1 `)
	matched = meanRegexp.MatchString(dataForGraphite)
	assert.Equal(t, matched, true)
}

func BenchmarkManyDifferentSensors(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "count" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := int64(r.Uint32() % 1000)
			counters[bucket] = a
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "gauge" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			gauges[bucket] = a
		}
	}

	var buff bytes.Buffer
	now := time.Now().Unix()
	t.ResetTimer()
	processTimers(&buff, now, commonPercentiles)
	processCounters(&buff, now)
	processGauges(&buff, now)
}

func BenchmarkOneBigTimer(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	bucket := "response_time"
	for i := 0; i < 10000000; i++ {
		a := uint64(r.Uint32() % 1000)
		timers[bucket] = append(timers[bucket], a)
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkLotsOfTimers(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkParseMessage(b *testing.B) {
	d := []byte("a.key.with-0.dash:4|c|@0.5")
	for i := 0; i < b.N; i++ {
		parseMessage(d)
	}
}
<commit_msg>Simplified tests:<commit_after>package main

import (
	"bytes"
	"github.com/bmizerany/assert"
	"math/rand"
	"strconv"
	"testing"
	"time"
)

var commonPercentiles = Percentiles{
	&Percentile{
		99,
		"99",
	},
}

func TestPacketParse(t *testing.T) {
	d := []byte("gaugor:333|g")
	packets := parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet := packets[0]
	assert.Equal(t, "gaugor", packet.Bucket)
	assert.Equal(t, uint64(333), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:2|c|@0.1")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(2), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(0.1), packet.Sampling)

	d = []byte("gorets:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("gorets:-4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "gorets", packet.Bucket)
	assert.Equal(t, int64(-4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("glork:320|ms")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "glork", packet.Bucket)
	assert.Equal(t, uint64(320), packet.Value.(uint64))
	assert.Equal(t, "ms", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4|c\ngauge:3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 2)
	packet = packets[0]
	assert.Equal(t, "a.key.with-0.dash", packet.Bucket)
	assert.Equal(t, int64(4), packet.Value.(int64))
	assert.Equal(t, "c", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	packet = packets[1]
	assert.Equal(t, "gauge", packet.Bucket)
	assert.Equal(t, uint64(3), packet.Value.(uint64))
	assert.Equal(t, "g", packet.Modifier)
	assert.Equal(t, float32(1), packet.Sampling)

	d = []byte("a.key.with-0.dash:4\ngauge3|g")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("a.key.with-0.dash:4")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5m")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|mg")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 0)

	d = []byte("gorets:5|ms|@")
	packets = parseMessage(d)
	assert.Equal(t, len(packets), 1)
}

func TestReceiveCounterPacketHandling(t *testing.T) {
	counters = make(map[string]int64)
	*receiveCounter = "countme"

	p := &Packet{
		Bucket:   "gorets",
		Value:    int64(100),
		Modifier: "c",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, counters["countme"], int64(1))

	packetHandler(p)
	assert.Equal(t, counters["countme"], int64(2))
}

func TestCountPacketHandling(t *testing.T) {
	counters = make(map[string]int64)

	p := &Packet{
		Bucket:   "gorets",
		Value:    int64(100),
		Modifier: "c",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(100))

	p.Value = int64(3)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(103))

	p.Value = int64(-4)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(99))

	p.Value = int64(-100)
	packetHandler(p)
	assert.Equal(t, counters["gorets"], int64(-1))
}

func TestGaugePacketHandling(t *testing.T) {
	gauges = make(map[string]uint64)

	p := &Packet{
		Bucket:   "gaugor",
		Value:    uint64(333),
		Modifier: "g",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, gauges["gaugor"], uint64(333))
}

func TestTimerPacketHandling(t *testing.T) {
	timers = make(map[string]Uint64Slice)

	p := &Packet{
		Bucket:   "glork",
		Value:    uint64(320),
		Modifier: "ms",
		Sampling: float32(1),
	}
	packetHandler(p)
	assert.Equal(t, len(timers["glork"]), 1)
	assert.Equal(t, timers["glork"][0], uint64(320))

	p.Value = uint64(100)
	packetHandler(p)
	assert.Equal(t, len(timers["glork"]), 2)
	assert.Equal(t, timers["glork"][1], uint64(100))
}

func TestProcessCounters(t *testing.T) {

	*persistCountKeys = int64(10)
	counters = make(map[string]int64)
	var buffer bytes.Buffer
	now := int64(1418052649)

	counters["gorets"] = int64(123)

	num := processCounters(&buffer, now)
	assert.Equal(t, num, int64(1))
	assert.Equal(t, buffer.String(), "gorets 123 1418052649\n")

	// run processCounters() enough times to make sure it purges items
	for i := 0; i < int(*persistCountKeys)+10; i++ {
		num = processCounters(&buffer, now)
	}
	lines := bytes.Split(buffer.Bytes(), []byte("\n"))

	// expect two more lines - the good one and an empty one at the end
	assert.Equal(t, len(lines), int(*persistCountKeys+2))
	assert.Equal(t, string(lines[0]), "gorets 123 1418052649")
	assert.Equal(t, string(lines[*persistCountKeys]), "gorets 0 1418052649")
}

func TestProcessTimers(t *testing.T) {
	// Some data with expected mean of 20
	timers = make(map[string]Uint64Slice)
	timers["response_time"] = []uint64{0, 30, 30}

	now := int64(1418052649)

	var buffer bytes.Buffer
	var num int64
	num += processTimers(&buffer, now, Percentiles{})

	lines := bytes.Split(buffer.Bytes(), []byte("\n"))

	assert.Equal(t, num, int64(1))
	assert.Equal(t, string(lines[0]), "response_time.mean 20.000000 1418052649")
	assert.Equal(t, string(lines[1]), "response_time.upper 30 1418052649")
	assert.Equal(t, string(lines[2]), "response_time.lower 0 1418052649")
	assert.Equal(t, string(lines[3]), "response_time.count 3 1418052649")
}

func TestProcessTimersUpperPercentile(t *testing.T) {
	// Some data with expected 75% of 2
	timers = make(map[string]Uint64Slice)
	timers["response_time"] = []uint64{0, 1, 2, 3}

	now := int64(1418052649)

	var buffer bytes.Buffer
	var num int64
	num += processTimers(&buffer, now, Percentiles{
		&Percentile{
			75,
			"75",
		},
	})

	lines := bytes.Split(buffer.Bytes(), []byte("\n"))

	assert.Equal(t, num, int64(1))
	assert.Equal(t, string(lines[0]), "response_time.upper_75 2 1418052649")
}

func TestProcessTimesLowerPercentile(t *testing.T) {
	timers = make(map[string]Uint64Slice)
	timers["time"] = []uint64{0, 1, 2, 3}

	now := int64(1418052649)

	var buffer bytes.Buffer
	var num int64
	num += processTimers(&buffer, now, Percentiles{
		&Percentile{
			-75,
			"-75",
		},
	})

	lines := bytes.Split(buffer.Bytes(), []byte("\n"))

	assert.Equal(t, num, int64(1))
	assert.Equal(t, string(lines[0]), "time.lower_75 1 1418052649")
}

func BenchmarkManyDifferentSensors(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "count" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := int64(r.Uint32() % 1000)
			counters[bucket] = a
		}
	}

	for i := 0; i < 1000; i++ {
		bucket := "gauge" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			gauges[bucket] = a
		}
	}

	var buff bytes.Buffer
	now := time.Now().Unix()
	t.ResetTimer()
	processTimers(&buff, now, commonPercentiles)
	processCounters(&buff, now)
	processGauges(&buff, now)
}

func BenchmarkOneBigTimer(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	bucket := "response_time"
	for i := 0; i < 10000000; i++ {
		a := uint64(r.Uint32() % 1000)
		timers[bucket] = append(timers[bucket], a)
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkLotsOfTimers(t *testing.B) {
	r := rand.New(rand.NewSource(438))
	for i := 0; i < 1000; i++ {
		bucket := "response_time" + strconv.Itoa(i)
		for i := 0; i < 10000; i++ {
			a := uint64(r.Uint32() % 1000)
			timers[bucket] = append(timers[bucket], a)
		}
	}

	var buff bytes.Buffer
	t.ResetTimer()
	processTimers(&buff, time.Now().Unix(), commonPercentiles)
}

func BenchmarkParseMessage(b *testing.B) {
	d := []byte("a.key.with-0.dash:4|c|@0.5")
	for i := 0; i < b.N; i++ {
		parseMessage(d)
	}
}
<|endoftext|>


<commit_before>package stm

import (
	"compress/zlib"
	"log"
	"os"
	"regexp"
)

var gzipPtn = regexp.MustCompile(".gz$")

func NewFileAdapter() *FileAdapter {
	adapter := &FileAdapter{}
	return adapter
}

type FileAdapter struct{}

func (adp *FileAdapter) Write(loc *Location, data []byte) {
	dir := loc.Directory()
	fi, err := os.Stat(dir)
	if err != nil {
		_ = os.MkdirAll(dir, 0755)
	} else if !fi.IsDir() {
		log.Fatalf("%s should be a directory", dir)
	}

	file, _ := os.OpenFile(loc.Path(), os.O_RDWR|os.O_CREATE, 0666)
	fi, err = file.Stat()
	if err != nil {
		log.Fatalf("%s file not exists", loc.Path())
	} else if !fi.Mode().IsRegular() {
		log.Fatalf("%s should be a filename", loc.Path())
	}

	if gzipPtn.MatchString(loc.Path()) {
		adp.gzip(file, data)
	} else {
		adp.plain(file, data)
	}
}

func (adp *FileAdapter) gzip(file *os.File, data []byte) {
	gz := zlib.NewWriter(file)
	defer gz.Close()
	gz.Write(data)
}

func (adp *FileAdapter) plain(file *os.File, data []byte) {
	file.Write(data)
	defer file.Close()
}
<commit_msg>fix gzip<commit_after>package stm

import (
	"compress/gzip"
	"log"
	"os"
	"regexp"
)

var gzipPtn = regexp.MustCompile(".gz$")

func NewFileAdapter() *FileAdapter {
	adapter := &FileAdapter{}
	return adapter
}

type FileAdapter struct{}

func (adp *FileAdapter) Write(loc *Location, data []byte) {
	dir := loc.Directory()
	fi, err := os.Stat(dir)
	if err != nil {
		_ = os.MkdirAll(dir, 0755)
	} else if !fi.IsDir() {
		log.Fatalf("%s should be a directory", dir)
	}

	file, _ := os.OpenFile(loc.Path(), os.O_RDWR|os.O_CREATE, 0666)
	fi, err = file.Stat()
	if err != nil {
		log.Fatalf("%s file not exists", loc.Path())
	} else if !fi.Mode().IsRegular() {
		log.Fatalf("%s should be a filename", loc.Path())
	}

	if gzipPtn.MatchString(loc.Path()) {
		adp.gzip(file, data)
	} else {
		adp.plain(file, data)
	}
}

func (adp *FileAdapter) gzip(file *os.File, data []byte) {
	gz := gzip.NewWriter(file)
	defer gz.Close()
	gz.Write(data)
}

func (adp *FileAdapter) plain(file *os.File, data []byte) {
	file.Write(data)
	defer file.Close()
}
<|endoftext|>


<commit_before>package storage

import (
	"github.com/golang/glog"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
)

func TestOpenWriteAheadFile(t *testing.T) {
	assert := assert.New(t)

	//Create temp directory
	dir, err := ioutil.TempDir("", "IosWALTests")
	if err != nil {
		glog.Fatal(err)
	}
	defer os.RemoveAll(dir) // clean up

	//create file
	testFile := dir + "/test.temp"
	wal := openWriteAheadFile(testFile, "fsync", 2560)
	actualBytes, err := ioutil.ReadFile(testFile)
	assert.Nil(err)
	// TODO; check file is expected size (maybe system dependant)
	glog.Info("File created of size ", len(actualBytes))
	// TODO: check actualBytes are empty

	//verfiy that write ahead logging works
	start := 0
	history := make([]byte, 0, 1000)
	for size := 1; size < 100; size += 10 {
		expectedBytes := make([]byte, size)
		rand.Read(expectedBytes)
		wal.writeAhead(expectedBytes)
		actualBytes, err = ioutil.ReadFile(testFile)
		assert.Nil(err)
		assert.Equal(expectedBytes, actualBytes[start:start+size], "Bytes read are not same as written")
		assert.Equal([]byte{0xa}, actualBytes[start+size:start+size+1], "Delim missing from end of write")
		assert.Equal(history[:start], actualBytes[:start], "Past writes have been corrupted")
		// TODO: check that rest of file is empty
		// update state for next run
		history = append(history, expectedBytes...)
		history = append(history, 0xa)
		start = start + size + 1
	}

	//verfiy that write ahead logging continues after failures
	walf := openWriteAheadFile(testFile, "fsync", 2560)
	actualBytesf, err := ioutil.ReadFile(testFile)
	assert.Nil(err)
	glog.Info("File now of size ", len(actualBytes))
	assert.Equal(history, actualBytesf[:start], "File has not recovered")

	// continue logging
	// for size := 1; size < 100; size += 10 {
	// 	expectedBytes := make([]byte, size)
	// 	rand.Read(expectedBytes)
	// 	walf.writeAhead(expectedBytes)
	// 	actualBytes, err = ioutil.ReadFile(testFile)
	// 	assert.Nil(err)
	// 	assert.Equal(expectedBytes, actualBytes[start:start+size], "Bytes read are not same as written")
	// 	assert.Equal([]byte{0xa}, actualBytes[start+size:start+size+1], "Delim missing from end of write")
	// 	assert.Equal(history[:start], actualBytes[:start], "Past writes have been corrupted")
	// 	// update state for next run
	// 	history = append(history, expectedBytes...)
	// 	history = append(history, 0xa)
	// 	start = start + size + 1
	// }

}
<commit_msg>quick fix - removing unused variable<commit_after>package storage

import (
	"github.com/golang/glog"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"math/rand"
	"os"
	"testing"
)

func TestOpenWriteAheadFile(t *testing.T) {
	assert := assert.New(t)

	//Create temp directory
	dir, err := ioutil.TempDir("", "IosWALTests")
	if err != nil {
		glog.Fatal(err)
	}
	defer os.RemoveAll(dir) // clean up

	//create file
	testFile := dir + "/test.temp"
	wal := openWriteAheadFile(testFile, "fsync", 2560)
	actualBytes, err := ioutil.ReadFile(testFile)
	assert.Nil(err)
	// TODO; check file is expected size (maybe system dependant)
	glog.Info("File created of size ", len(actualBytes))
	// TODO: check actualBytes are empty

	//verfiy that write ahead logging works
	start := 0
	history := make([]byte, 0, 1000)
	for size := 1; size < 100; size += 10 {
		expectedBytes := make([]byte, size)
		rand.Read(expectedBytes)
		wal.writeAhead(expectedBytes)
		actualBytes, err = ioutil.ReadFile(testFile)
		assert.Nil(err)
		assert.Equal(expectedBytes, actualBytes[start:start+size], "Bytes read are not same as written")
		assert.Equal([]byte{0xa}, actualBytes[start+size:start+size+1], "Delim missing from end of write")
		assert.Equal(history[:start], actualBytes[:start], "Past writes have been corrupted")
		// TODO: check that rest of file is empty
		// update state for next run
		history = append(history, expectedBytes...)
		history = append(history, 0xa)
		start = start + size + 1
	}

	//verfiy that write ahead logging continues after failures
	openWriteAheadFile(testFile, "fsync", 2560)
	actualBytesf, err := ioutil.ReadFile(testFile)
	assert.Nil(err)
	glog.Info("File now of size ", len(actualBytes))
	assert.Equal(history, actualBytesf[:start], "File has not recovered")

	// continue logging
	// for size := 1; size < 100; size += 10 {
	// 	expectedBytes := make([]byte, size)
	// 	rand.Read(expectedBytes)
	// 	walf.writeAhead(expectedBytes)
	// 	actualBytes, err = ioutil.ReadFile(testFile)
	// 	assert.Nil(err)
	// 	assert.Equal(expectedBytes, actualBytes[start:start+size], "Bytes read are not same as written")
	// 	assert.Equal([]byte{0xa}, actualBytes[start+size:start+size+1], "Delim missing from end of write")
	// 	assert.Equal(history[:start], actualBytes[:start], "Past writes have been corrupted")
	// 	// update state for next run
	// 	history = append(history, expectedBytes...)
	// 	history = append(history, 0xa)
	// 	start = start + size + 1
	// }

}
<|endoftext|>


<commit_before>package ar

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strconv"
	"strings"
	"testing"
	"time"
)

// test whether fileInfo implements os.FileInfo
var _ os.FileInfo = new(fileInfo)

var testCommon = "!<arch>\n" +
	"debian-binary   1385068169  0     0     100644  4         `\n" +
	"2.0\n" +
	"control.tar.gz  1385068169  0     0     100644  0         `\n"

var testCommonFileHeaders = []struct {
	in   string
	want *fileInfo
	err  string
}{
	{
		in: "debian-binary   1385068169  0     0     100644  4         `\n",
		want: &fileInfo{
			name:  "debian-binary",
			mtime: time.Unix(1385068169, 0),
			mode:  os.FileMode(0100644) & os.ModePerm,
			size:  4,
		},
	},
	{
		in: "debian-binary   1385068169  0     0     644     4         `\n",
		want: &fileInfo{
			name:  "debian-binary",
			mtime: time.Unix(1385068169, 0),
			mode:  os.FileMode(0644),
			size:  4,
		},
	},
	{
		in:  "debian-binary   1385068169  0     0     120644  4         `\n",
		err: "feature not implemented: non-regular files",
	},
	{
		in:  "debian-binary   1385068169  0     0     220644  4         `\n",
		err: "corrupt archive: invalid file mode",
	},
}

func TestReadFileHeader(t *testing.T) {
	for i, test := range testCommonFileHeaders {
		got, err := readFileHeader(strings.NewReader(test.in))
		switch {
		case err == nil && test.err != "":
			t.Errorf("%d: got no err, expected err %v", i, test.err)
			continue
		case err != nil && test.err != err.Error():
			t.Errorf("%d: got err %q, expected err %q", i, err, test.err)
			continue
		case err == nil && test.err == "":
			// no error as expected
		case err != nil && test.err == err.Error():
			t.Logf("%d: got expected error %q", i, err)
			continue
		}

		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("%d: got %#v, expected %+v", i, got, test.want)
		} else {
			t.Logf("%d: got %#v", i, got)
		}
	}
}

func BenchmarkReadFileHeader(b *testing.B) {
	r := bytes.NewReader([]byte("debian-binary   1385068169  0     0     100644  4         `\n"))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = readFileHeader(r)
		b.SetBytes(60)
		r.Seek(0, 0)
	}

}

var testMagic = []struct {
	in   string
	want error
}{
	{
		in: magic,
	},
	{
		in:   strings.Repeat("a", len(magic)),
		want: CorruptArchiveError("global archive header not found"),
	},
	{
		in:   "!",
		want: io.ErrUnexpectedEOF,
	},
	{
		in:   "",
		want: io.EOF,
	},
}

func TestReadMagic(t *testing.T) {
	for i, test := range testMagic {
		got := checkMagic(strings.NewReader(test.in))
		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("%d: got %#v, expected %+v", i, got, test.want)
		} else {
			t.Logf("%d: got %#v", i, got)
		}
	}
}

func TestFileInfo(t *testing.T) {
	test := &fileInfo{
		name:  "debian-binary",
		mtime: time.Unix(1385068169, 0),
		mode:  os.FileMode(0644),
		size:  4,
	}

	if test.IsDir() != false {
		t.Error("IsDir")
	}
	if test.Mode() != os.FileMode(0644) {
		t.Error("Mode")
	}
	if test.ModTime() != time.Unix(1385068169, 0) {
		t.Error("ModTime")
	}
	if test.Name() != "debian-binary" {
		t.Error("Name")
	}
	if test.Size() != 4 {
		t.Error("Size")
	}
	if test.Sys() != nil {
		t.Error("Sys")
	}
}

func TestReaderBasics(t *testing.T) {
	test := strings.NewReader(testCommon)
	r := NewReader(test)
	fi, err := r.Next()
	if err != nil {
		t.Fatal(err)
		return
	}
	if fi.Mode() != os.FileMode(0644) {
		t.Error("Mode")
	}
	if fi.Name() != "debian-binary" {
		t.Error("Name")
	}
	if fi.Size() != 4 {
		t.Error("Size")
	}
	if fi.ModTime() != time.Unix(1385068169, 0) {
		t.Error("ModTime")
	}

	if content, err := ioutil.ReadAll(r); err != nil {
		t.Error(err)
	} else if string(content) != "2.0\n" {
		t.Error("Content")
	}
	fi, err = r.Next()
	if err != nil {
		t.Fatal(err)
		return
	}
	if fi.Mode() != os.FileMode(0644) {
		t.Error("Mode2")
	}
	if fi.Name() != "control.tar.gz" {
		t.Error("Name2")
	}
	if fi.Size() != 0 {
		t.Error("Size2")
	}
	if fi.ModTime() != time.Unix(1385068169, 0) {
		t.Error("ModTime2")
	}

	if content, err := ioutil.ReadAll(r); err != nil {
		t.Error(err)
	} else if string(content) != "" {
		t.Error("Content2")
	}

	fi, err = r.Next()
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}
}

func BenchmarkReaderBigFiles(b *testing.B) {
	benchmarkReader(b, 8, 8*1024*1024)
}

func BenchmarkReaderManySmallFiles(b *testing.B) {
	benchmarkReader(b, 1024, 8)
}

func genArchiveFile(meta os.FileInfo) []byte {
	ret := make([]byte, 60+meta.Size(), 60+meta.Size())
	for i := 0; i < 60; i++ {
		ret[i] = ' '
	}
	copy(ret[0:], meta.Name())
	copy(ret[16:], strconv.FormatInt(meta.ModTime().Unix(), 10))
	copy(ret[28:], "0")
	copy(ret[34:], "0")
	copy(ret[40:], strconv.FormatUint(uint64(meta.Mode()), 8))
	copy(ret[48:], strconv.FormatInt(meta.Size(), 10))
	copy(ret[58:], filemagic)
	return ret
}

func benchmarkReader(b *testing.B, numFiles int, sizeFiles int64) {
	buf := bytes.NewBufferString(magic)
	for i := 0; i < numFiles; i++ {
		buf.Write(genArchiveFile(&fileInfo{
			name:  strconv.Itoa(1000 + i),
			mtime: time.Unix(int64(i), 0),
			size:  sizeFiles,
			mode:  os.FileMode(0640),
		}))

	}
	test := bytes.NewReader(buf.Bytes())
	r := NewReader(test)

	var err error
	var read int64

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for j := 0; j < numFiles; j++ {
			_, err = r.Next()
			if err != nil {
				b.Fatal(err)
			}
			read += 60
			n, err := io.Copy(ioutil.Discard, r)
			if err != nil {
				b.Fatal(err)
			}
			read += n
		}
		b.SetBytes(read)
		read = 0
		test.Seek(0, 0)
		r.Reset(test)
	}
}
<commit_msg>Preallocate<commit_after>package ar

import (
	"bytes"
	"io"
	"io/ioutil"
	"os"
	"reflect"
	"strconv"
	"strings"
	"testing"
	"time"
)

// test whether fileInfo implements os.FileInfo
var _ os.FileInfo = new(fileInfo)

var testCommon = "!<arch>\n" +
	"debian-binary   1385068169  0     0     100644  4         `\n" +
	"2.0\n" +
	"control.tar.gz  1385068169  0     0     100644  0         `\n"

var testCommonFileHeaders = []struct {
	in   string
	want *fileInfo
	err  string
}{
	{
		in: "debian-binary   1385068169  0     0     100644  4         `\n",
		want: &fileInfo{
			name:  "debian-binary",
			mtime: time.Unix(1385068169, 0),
			mode:  os.FileMode(0100644) & os.ModePerm,
			size:  4,
		},
	},
	{
		in: "debian-binary   1385068169  0     0     644     4         `\n",
		want: &fileInfo{
			name:  "debian-binary",
			mtime: time.Unix(1385068169, 0),
			mode:  os.FileMode(0644),
			size:  4,
		},
	},
	{
		in:  "debian-binary   1385068169  0     0     120644  4         `\n",
		err: "feature not implemented: non-regular files",
	},
	{
		in:  "debian-binary   1385068169  0     0     220644  4         `\n",
		err: "corrupt archive: invalid file mode",
	},
}

func TestReadFileHeader(t *testing.T) {
	for i, test := range testCommonFileHeaders {
		got, err := readFileHeader(strings.NewReader(test.in))
		switch {
		case err == nil && test.err != "":
			t.Errorf("%d: got no err, expected err %v", i, test.err)
			continue
		case err != nil && test.err != err.Error():
			t.Errorf("%d: got err %q, expected err %q", i, err, test.err)
			continue
		case err == nil && test.err == "":
			// no error as expected
		case err != nil && test.err == err.Error():
			t.Logf("%d: got expected error %q", i, err)
			continue
		}

		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("%d: got %#v, expected %+v", i, got, test.want)
		} else {
			t.Logf("%d: got %#v", i, got)
		}
	}
}

func BenchmarkReadFileHeader(b *testing.B) {
	r := bytes.NewReader([]byte("debian-binary   1385068169  0     0     100644  4         `\n"))
	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, _ = readFileHeader(r)
		b.SetBytes(60)
		r.Seek(0, 0)
	}

}

var testMagic = []struct {
	in   string
	want error
}{
	{
		in: magic,
	},
	{
		in:   strings.Repeat("a", len(magic)),
		want: CorruptArchiveError("global archive header not found"),
	},
	{
		in:   "!",
		want: io.ErrUnexpectedEOF,
	},
	{
		in:   "",
		want: io.EOF,
	},
}

func TestReadMagic(t *testing.T) {
	for i, test := range testMagic {
		got := checkMagic(strings.NewReader(test.in))
		if !reflect.DeepEqual(got, test.want) {
			t.Errorf("%d: got %#v, expected %+v", i, got, test.want)
		} else {
			t.Logf("%d: got %#v", i, got)
		}
	}
}

func TestFileInfo(t *testing.T) {
	test := &fileInfo{
		name:  "debian-binary",
		mtime: time.Unix(1385068169, 0),
		mode:  os.FileMode(0644),
		size:  4,
	}

	if test.IsDir() != false {
		t.Error("IsDir")
	}
	if test.Mode() != os.FileMode(0644) {
		t.Error("Mode")
	}
	if test.ModTime() != time.Unix(1385068169, 0) {
		t.Error("ModTime")
	}
	if test.Name() != "debian-binary" {
		t.Error("Name")
	}
	if test.Size() != 4 {
		t.Error("Size")
	}
	if test.Sys() != nil {
		t.Error("Sys")
	}
}

func TestReaderBasics(t *testing.T) {
	test := strings.NewReader(testCommon)
	r := NewReader(test)
	fi, err := r.Next()
	if err != nil {
		t.Fatal(err)
		return
	}
	if fi.Mode() != os.FileMode(0644) {
		t.Error("Mode")
	}
	if fi.Name() != "debian-binary" {
		t.Error("Name")
	}
	if fi.Size() != 4 {
		t.Error("Size")
	}
	if fi.ModTime() != time.Unix(1385068169, 0) {
		t.Error("ModTime")
	}

	if content, err := ioutil.ReadAll(r); err != nil {
		t.Error(err)
	} else if string(content) != "2.0\n" {
		t.Error("Content")
	}
	fi, err = r.Next()
	if err != nil {
		t.Fatal(err)
		return
	}
	if fi.Mode() != os.FileMode(0644) {
		t.Error("Mode2")
	}
	if fi.Name() != "control.tar.gz" {
		t.Error("Name2")
	}
	if fi.Size() != 0 {
		t.Error("Size2")
	}
	if fi.ModTime() != time.Unix(1385068169, 0) {
		t.Error("ModTime2")
	}

	if content, err := ioutil.ReadAll(r); err != nil {
		t.Error(err)
	} else if string(content) != "" {
		t.Error("Content2")
	}

	fi, err = r.Next()
	if err != io.EOF {
		t.Errorf("expected EOF, got %v", err)
	}
}

func BenchmarkReaderBigFiles(b *testing.B) {
	benchmarkReader(b, 8, 8*1024*1024)
}

func BenchmarkReaderManySmallFiles(b *testing.B) {
	benchmarkReader(b, 1024, 8)
}

func genArchiveFile(meta os.FileInfo) []byte {
	ret := make([]byte, 60+meta.Size(), 60+meta.Size())
	for i := 0; i < 60; i++ {
		ret[i] = ' '
	}
	copy(ret[0:], meta.Name())
	copy(ret[16:], strconv.FormatInt(meta.ModTime().Unix(), 10))
	copy(ret[28:], "0")
	copy(ret[34:], "0")
	copy(ret[40:], strconv.FormatUint(uint64(meta.Mode()), 8))
	copy(ret[48:], strconv.FormatInt(meta.Size(), 10))
	copy(ret[58:], filemagic)
	return ret
}

func benchmarkReader(b *testing.B, numFiles int, sizeFiles int64) {
	buf := bytes.NewBufferString(magic)
	buf.Grow(numFiles * (int(sizeFiles) + 60))
	for i := 0; i < numFiles; i++ {
		buf.Write(genArchiveFile(&fileInfo{
			name:  strconv.Itoa(1000 + i),
			mtime: time.Unix(int64(i), 0),
			size:  sizeFiles,
			mode:  os.FileMode(0640),
		}))

	}
	test := bytes.NewReader(buf.Bytes())
	r := NewReader(test)

	var err error
	var read int64

	b.ReportAllocs()
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		for j := 0; j < numFiles; j++ {
			_, err = r.Next()
			if err != nil {
				b.Fatal(err)
			}
			read += 60
			n, err := io.Copy(ioutil.Discard, r)
			if err != nil {
				b.Fatal(err)
			}
			read += n
		}
		b.SetBytes(read)
		read = 0
		test.Seek(0, 0)
		r.Reset(test)
	}
}
<|endoftext|>


<commit_before>package xml

import (
	"encoding/xml"
)

type X struct {
	A int
	B float64
	C bool
	D string
	E []byte
	F *X
}

func Fuzz(data []byte) int {
	score := 0
	if xml.Unmarshal(data, nil) == nil {
		score++
	}
	var s string
	if xml.Unmarshal(data, &s) == nil {
		score++
	}
	var a []string
	if xml.Unmarshal(data, &a) == nil {
		score++
	}
	var x X
	if xml.Unmarshal(data, &x) == nil {
		score++
	}
	return score
}
<commit_msg>improve and extend xml test<commit_after>package xml

import (
	"bytes"
	"encoding/xml"
)

type X struct {
	A int     `xml:"a,attr"`
	B float64 `xml:"B"`
	C bool    `xml:"C>CC"`
	D string  `xml:",comment"`
	E []byte  `xml:",innerxml"`
	F *X
	G string `xml:",any"`
	H string `xml:"-"`
	J []byte `xml:",chardata"`
}

func Fuzz(data []byte) int {
	score := 0
	if test(data, nil) {
		score = 1
	}
	var s string
	if test(data, &s) {
		score = 1
	}
	var a []string
	if test(data, &a) {
		score = 1
	}
	var x X
	if test(data, &x) {
		score = 1
		xml.Marshal(&x)
	}
	return score
}

func test(data []byte, v interface{}) (valid bool) {
	dec := xml.NewDecoder(bytes.NewReader(data))
	if dec.Decode(v) == nil {
		valid = true
	}
	dec1 := xml.NewDecoder(bytes.NewReader(data))
	dec1.Strict = false
	dec1.AutoClose = xml.HTMLAutoClose
	dec1.Entity = xml.HTMLEntity
	if dec1.Decode(v) == nil {
		valid = true
	} else if valid {
		panic("non-strict mode is weaker than strict")
	}
	return
}
<|endoftext|>


<commit_before>/* Copyright (c) 2016-2017 Jason Ish
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
 * IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

package elasticsearch

import (
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/eve"
	"github.com/jasonish/evebox/log"
	"github.com/pkg/errors"
	"time"
)

type DataStore struct {
	es *ElasticSearch
	core.UnimplementedDatastore
}

func NewDataStore(es *ElasticSearch) (*DataStore, error) {
	datastore := DataStore{
		es: es,
	}
	return &datastore, nil
}

func (d *DataStore) GetEveEventSink() core.EveEventSink {
	return NewIndexer(d.es)
}

func (s *DataStore) asKeyword(keyword string) string {
	return s.es.FormatKeyword(keyword)
}

// FindFlow finds the flow events matching the query parameters in options.
func (d *DataStore) FindFlow(flowId uint64, proto string, timestamp string,
	srcIp string, destIp string) (interface{}, error) {

	query := NewEventQuery()

	query.EventType("flow")
	query.AddFilter(TermQuery("flow_id", flowId))
	query.AddFilter(TermQuery("proto", proto))
	query.AddFilter(RangeLte("flow.start", timestamp))
	query.AddFilter(RangeGte("flow.end", timestamp))
	query.ShouldHaveIp(srcIp, d.es.GetKeyword())
	query.ShouldHaveIp(destIp, d.es.GetKeyword())

	response, err := d.es.Search(query)
	if err != nil {
		log.Error("%v", err)
		return nil, err
	}

	return response.Hits.Hits, nil
}

// FindNetflow finds netflow events matching the parameters in options.
func (s *DataStore) FindNetflow(options core.EventQueryOptions, sortBy string,
	order string) (interface{}, error) {

	size := int64(10)

	if options.Size > 0 {
		size = options.Size
	}

	if order == "" {
		order = "desc"
	}

	query := NewEventQuery()
	query.AddFilter(TermQuery("event_type", "netflow"))

	if options.TimeRange != "" {
		query.AddTimeRangeFilter(options.TimeRange)
	}

	if options.QueryString != "" {
		query.AddFilter(QueryString(options.QueryString))
	}

	if sortBy != "" {
		query.Aggs["agg"] = TopHitsAgg(sortBy, order, size)
	} else {
		query.SetSize(size)
	}

	response, err := s.es.Search(query)
	if err != nil {
		return nil, err
	}

	// Unwrap response.
	hits := response.Aggregations.GetMap("agg").GetMap("hits").Get("hits")

	return map[string]interface{}{
		"data": hits,
	}, nil
}

const ACTION_ARCHIVED = "archived"

const ACTION_ESCALATED = "escalated"

const ACTION_DEESCALATED = "de-escalated"

const ACTION_COMMENT = "comment"

type HistoryEntry struct {
	Timestamp string `json:"timestamp"`
	Username  string `json:"username"`
	Action    string `json:"action"`
	Comment   string `json:"comment,omitempty"`
}

func (s *DataStore) buildAlertGroupQuery(p core.AlertGroupQueryParams) *EventQuery {
	q := EventQuery{}
	q.AddFilter(ExistsQuery("event_type"))
	q.AddFilter(KeywordTermQuery("event_type", "alert", s.es.GetKeyword()))
	q.AddFilter(NewRangeQuery("@timestamp", eve.FormatTimestampUTC(p.MinTs),
		eve.FormatTimestampUTC(p.MaxTs)))
	q.AddFilter(KeywordTermQuery("src_ip", p.SrcIP, s.es.GetKeyword()))
	q.AddFilter(KeywordTermQuery("dest_ip", p.DstIP, s.es.GetKeyword()))
	q.AddFilter(TermQuery("alert.signature_id", p.SignatureID))
	return &q
}

// ArchiveAlertGroupByQuery uses the Elastic Search update_by_query API to
// archive events with a query instead of updating each document. This is
// only available in Elastic Search v5+.
func (s *DataStore) AddTagsToAlertGroupsByQuery(p core.AlertGroupQueryParams, tags []string, action HistoryEntry) error {
	log.Println("AddTagsToAlertGroupsByQuery")
	var mustNot []interface{}
	for _, tag := range tags {
		mustNot = append(mustNot, TermQuery("tags", tag))
	}

	query := s.buildAlertGroupQuery(p)
	if len(mustNot) > 0 {
		query.Query.Bool.MustNot = mustNot
	}
	query.Script = &Script{
		Lang: "painless",
		Inline: `
		        if (params.tags != null) {
			        if (ctx._source.tags == null) {
			            ctx._source.tags = new ArrayList();
			        }
			        for (tag in params.tags) {
			            if (!ctx._source.tags.contains(tag)) {
			                ctx._source.tags.add(tag);
			            }
			        }
			    }
			    if (ctx._source.evebox == null) {
			        ctx._source.evebox = new HashMap();
			    }
			    if (ctx._source.evebox.history == null) {
			        ctx._source.evebox.history = new ArrayList();
			    }
			    ctx._source.evebox.history.add(params.action);
		`,
		Params: map[string]interface{}{
			"tags":   tags,
			"action": action,
		},
	}

	response, err := s.es.doUpdateByQuery(query)
	if err != nil {
		log.Error("failed to update by query: %v", err)
		return err
	}
	log.Info("Events updated: %v; failures=%d",
		response.Get("updated"), len(response.GetMapList("failures")))

	return nil
}

// ArchiveAlertGroup is a specialization of AddTagsToAlertGroup.
func (s *DataStore) ArchiveAlertGroup(p core.AlertGroupQueryParams, user core.User) error {
	tags := []string{"archived", "evebox.archived"}
	return s.AddTagsToAlertGroupsByQuery(p, tags, HistoryEntry{
		Action:    ACTION_ARCHIVED,
		Timestamp: FormatTimestampUTC(time.Now()),
		Username:  user.Username,
	})
}

// EscalateAlertGroup is a specialization of AddTagsToAlertGroup.
func (s *DataStore) EscalateAlertGroup(p core.AlertGroupQueryParams, user core.User) error {
	tags := []string{"escalated", "evebox.escalated"}
	history := HistoryEntry{
		Username:  user.Username,
		Action:    ACTION_ESCALATED,
		Timestamp: FormatTimestampUTC(time.Now()),
	}
	return s.AddTagsToAlertGroupsByQuery(p, tags, history)
}

func (s *DataStore) DeEscalateAlertGroup(p core.AlertGroupQueryParams, user core.User) error {
	tags := []string{"escalated", "evebox.escalated"}
	return s.RemoveTagsFromAlertGroupsByQuery(p, tags, HistoryEntry{
		Username:  user.Username,
		Timestamp: FormatTimestampUTC(time.Now()),
		Action:    ACTION_DEESCALATED,
	})
}

func (s *DataStore) RemoveTagsFromAlertGroupsByQuery(p core.AlertGroupQueryParams,
	tags []string, action HistoryEntry) error {
	var should []interface{}
	for _, tag := range tags {
		should = append(should, TermQuery("tags", tag))
	}

	query := s.buildAlertGroupQuery(p)
	query.Query.Bool.Should = should
	query.Script = &Script{
		Lang: "painless",
		Inline: `
			    if (ctx._source.tags != null) {
			        for (tag in params.tags) {
			            ctx._source.tags.removeIf(entry -> entry == tag);
			        }
			    }
			    if (ctx._source.evebox == null) {
			        ctx._source.evebox = new HashMap();
			    }
			    if (ctx._source.evebox.history == null) {
			        ctx._source.evebox.history = new ArrayList();
			    }
			    ctx._source.evebox.history.add(params.action);
		`,
		Params: map[string]interface{}{
			"tags":   tags,
			"action": action,
		},
	}

	response, err := s.es.doUpdateByQuery(query)
	if err != nil {
		log.Error("failed to update by query: %v", err)
		return err
	}
	log.Info("Events updated: %v; failures=%d",
		response.Get("updated"), len(response.GetMapList("failures")))

	return nil
}

func (s *DataStore) CommentOnAlertGroup(p core.AlertGroupQueryParams, user core.User, comment string) error {
	history := HistoryEntry{
		Username:  user.Username,
		Action:    ACTION_COMMENT,
		Comment:   comment,
		Timestamp: FormatTimestampUTC(time.Now()),
	}
	return s.AddTagsToAlertGroupsByQuery(p, nil, history)
}

func (s *DataStore) CommentOnEventId(eventId string, user core.User, comment string) error {

	event, err := s.GetEventById(eventId)
	if err != nil {
		return errors.Wrapf(err, "failed to find event with ID %s", eventId)
	}
	doc := Document{event}

	action := HistoryEntry{
		Username:  user.Username,
		Action:    ACTION_COMMENT,
		Comment:   comment,
		Timestamp: FormatTimestampUTC(time.Now()),
	}

	query := EventQuery{}
	query.Script = &Script{
		Lang: "painless",
		Inline: `
			    if (ctx._source.evebox == null) {
			        ctx._source.evebox = new HashMap();
			    }
			    if (ctx._source.evebox.history == null) {
			        ctx._source.evebox.history = new ArrayList();
			    }
			    ctx._source.evebox.history.add(params.action);
		`,
		Params: map[string]interface{}{
			"action": action,
		},
	}

	_, err = s.es.Update(doc.Index(), doc.Type(), doc.Id(), query)
	return err
}
<commit_msg>es: if failures on archive, print the es error<commit_after>/* Copyright (c) 2016-2017 Jason Ish
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
 * INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
 * IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 */

package elasticsearch

import (
	"github.com/jasonish/evebox/core"
	"github.com/jasonish/evebox/eve"
	"github.com/jasonish/evebox/log"
	"github.com/pkg/errors"
	"time"
)

type DataStore struct {
	es *ElasticSearch
	core.UnimplementedDatastore
}

func NewDataStore(es *ElasticSearch) (*DataStore, error) {
	datastore := DataStore{
		es: es,
	}
	return &datastore, nil
}

func (d *DataStore) GetEveEventSink() core.EveEventSink {
	return NewIndexer(d.es)
}

func (s *DataStore) asKeyword(keyword string) string {
	return s.es.FormatKeyword(keyword)
}

// FindFlow finds the flow events matching the query parameters in options.
func (d *DataStore) FindFlow(flowId uint64, proto string, timestamp string,
	srcIp string, destIp string) (interface{}, error) {

	query := NewEventQuery()

	query.EventType("flow")
	query.AddFilter(TermQuery("flow_id", flowId))
	query.AddFilter(TermQuery("proto", proto))
	query.AddFilter(RangeLte("flow.start", timestamp))
	query.AddFilter(RangeGte("flow.end", timestamp))
	query.ShouldHaveIp(srcIp, d.es.GetKeyword())
	query.ShouldHaveIp(destIp, d.es.GetKeyword())

	response, err := d.es.Search(query)
	if err != nil {
		log.Error("%v", err)
		return nil, err
	}

	return response.Hits.Hits, nil
}

// FindNetflow finds netflow events matching the parameters in options.
func (s *DataStore) FindNetflow(options core.EventQueryOptions, sortBy string,
	order string) (interface{}, error) {

	size := int64(10)

	if options.Size > 0 {
		size = options.Size
	}

	if order == "" {
		order = "desc"
	}

	query := NewEventQuery()
	query.AddFilter(TermQuery("event_type", "netflow"))

	if options.TimeRange != "" {
		query.AddTimeRangeFilter(options.TimeRange)
	}

	if options.QueryString != "" {
		query.AddFilter(QueryString(options.QueryString))
	}

	if sortBy != "" {
		query.Aggs["agg"] = TopHitsAgg(sortBy, order, size)
	} else {
		query.SetSize(size)
	}

	response, err := s.es.Search(query)
	if err != nil {
		return nil, err
	}

	// Unwrap response.
	hits := response.Aggregations.GetMap("agg").GetMap("hits").Get("hits")

	return map[string]interface{}{
		"data": hits,
	}, nil
}

const ACTION_ARCHIVED = "archived"

const ACTION_ESCALATED = "escalated"

const ACTION_DEESCALATED = "de-escalated"

const ACTION_COMMENT = "comment"

type HistoryEntry struct {
	Timestamp string `json:"timestamp"`
	Username  string `json:"username"`
	Action    string `json:"action"`
	Comment   string `json:"comment,omitempty"`
}

func (s *DataStore) buildAlertGroupQuery(p core.AlertGroupQueryParams) *EventQuery {
	q := EventQuery{}
	q.AddFilter(ExistsQuery("event_type"))
	q.AddFilter(KeywordTermQuery("event_type", "alert", s.es.GetKeyword()))
	q.AddFilter(NewRangeQuery("@timestamp", eve.FormatTimestampUTC(p.MinTs),
		eve.FormatTimestampUTC(p.MaxTs)))
	q.AddFilter(KeywordTermQuery("src_ip", p.SrcIP, s.es.GetKeyword()))
	q.AddFilter(KeywordTermQuery("dest_ip", p.DstIP, s.es.GetKeyword()))
	q.AddFilter(TermQuery("alert.signature_id", p.SignatureID))
	return &q
}

// ArchiveAlertGroupByQuery uses the Elastic Search update_by_query API to
// archive events with a query instead of updating each document. This is
// only available in Elastic Search v5+.
func (s *DataStore) AddTagsToAlertGroupsByQuery(p core.AlertGroupQueryParams, tags []string, action HistoryEntry) error {
	var mustNot []interface{}
	for _, tag := range tags {
		mustNot = append(mustNot, TermQuery("tags", tag))
	}

	query := s.buildAlertGroupQuery(p)
	if len(mustNot) > 0 {
		query.Query.Bool.MustNot = mustNot
	}
	query.Script = &Script{
		Lang: "painless",
		Inline: `
		        if (params.tags != null) {
			        if (ctx._source.tags == null) {
			            ctx._source.tags = new ArrayList();
			        }
			        for (tag in params.tags) {
			            if (!ctx._source.tags.contains(tag)) {
			                ctx._source.tags.add(tag);
			            }
			        }
			    }
			    if (ctx._source.evebox == null) {
			        ctx._source.evebox = new HashMap();
			    }
			    if (ctx._source.evebox.history == null) {
			        ctx._source.evebox.history = new ArrayList();
			    }
			    ctx._source.evebox.history.add(params.action);
		`,
		Params: map[string]interface{}{
			"tags":   tags,
			"action": action,
		},
	}

	response, err := s.es.doUpdateByQuery(query)
	if err != nil {
		log.Error("failed to update by query: %v", err)
		return err
	}
	failures := response.GetMapList("failures")
	log.Info("Events updated: %v; failures=%d {%+v}",
		response.Get("updated"), len(failures), failures)

	return nil
}

// ArchiveAlertGroup is a specialization of AddTagsToAlertGroup.
func (s *DataStore) ArchiveAlertGroup(p core.AlertGroupQueryParams, user core.User) error {
	tags := []string{"archived", "evebox.archived"}
	return s.AddTagsToAlertGroupsByQuery(p, tags, HistoryEntry{
		Action:    ACTION_ARCHIVED,
		Timestamp: FormatTimestampUTC(time.Now()),
		Username:  user.Username,
	})
}

// EscalateAlertGroup is a specialization of AddTagsToAlertGroup.
func (s *DataStore) EscalateAlertGroup(p core.AlertGroupQueryParams, user core.User) error {
	tags := []string{"escalated", "evebox.escalated"}
	history := HistoryEntry{
		Username:  user.Username,
		Action:    ACTION_ESCALATED,
		Timestamp: FormatTimestampUTC(time.Now()),
	}
	return s.AddTagsToAlertGroupsByQuery(p, tags, history)
}

func (s *DataStore) DeEscalateAlertGroup(p core.AlertGroupQueryParams, user core.User) error {
	tags := []string{"escalated", "evebox.escalated"}
	return s.RemoveTagsFromAlertGroupsByQuery(p, tags, HistoryEntry{
		Username:  user.Username,
		Timestamp: FormatTimestampUTC(time.Now()),
		Action:    ACTION_DEESCALATED,
	})
}

func (s *DataStore) RemoveTagsFromAlertGroupsByQuery(p core.AlertGroupQueryParams,
	tags []string, action HistoryEntry) error {
	var should []interface{}
	for _, tag := range tags {
		should = append(should, TermQuery("tags", tag))
	}

	query := s.buildAlertGroupQuery(p)
	query.Query.Bool.Should = should
	query.Script = &Script{
		Lang: "painless",
		Inline: `
			    if (ctx._source.tags != null) {
			        for (tag in params.tags) {
			            ctx._source.tags.removeIf(entry -> entry == tag);
			        }
			    }
			    if (ctx._source.evebox == null) {
			        ctx._source.evebox = new HashMap();
			    }
			    if (ctx._source.evebox.history == null) {
			        ctx._source.evebox.history = new ArrayList();
			    }
			    ctx._source.evebox.history.add(params.action);
		`,
		Params: map[string]interface{}{
			"tags":   tags,
			"action": action,
		},
	}

	response, err := s.es.doUpdateByQuery(query)
	if err != nil {
		log.Error("failed to update by query: %v", err)
		return err
	}
	failures := response.GetMapList("failures")
	log.Info("Events updated: %v; failures=%d {%+v}",
		response.Get("updated"), len(failures), failures)

	return nil
}

func (s *DataStore) CommentOnAlertGroup(p core.AlertGroupQueryParams, user core.User, comment string) error {
	history := HistoryEntry{
		Username:  user.Username,
		Action:    ACTION_COMMENT,
		Comment:   comment,
		Timestamp: FormatTimestampUTC(time.Now()),
	}
	return s.AddTagsToAlertGroupsByQuery(p, nil, history)
}

func (s *DataStore) CommentOnEventId(eventId string, user core.User, comment string) error {

	event, err := s.GetEventById(eventId)
	if err != nil {
		return errors.Wrapf(err, "failed to find event with ID %s", eventId)
	}
	doc := Document{event}

	action := HistoryEntry{
		Username:  user.Username,
		Action:    ACTION_COMMENT,
		Comment:   comment,
		Timestamp: FormatTimestampUTC(time.Now()),
	}

	query := EventQuery{}
	query.Script = &Script{
		Lang: "painless",
		Inline: `
			    if (ctx._source.evebox == null) {
			        ctx._source.evebox = new HashMap();
			    }
			    if (ctx._source.evebox.history == null) {
			        ctx._source.evebox.history = new ArrayList();
			    }
			    ctx._source.evebox.history.add(params.action);
		`,
		Params: map[string]interface{}{
			"action": action,
		},
	}

	_, err = s.es.Update(doc.Index(), doc.Type(), doc.Id(), query)
	return err
}
<|endoftext|>


<commit_before><commit_msg>refs #13 Cookie 保存時に Domain と Path を設定する<commit_after><|endoftext|>


<commit_before><commit_msg>remove dummy checkin<commit_after><|endoftext|>


<commit_before>package beater

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

	"github.com/appcelerator/amp/api/rpc/logs"
	"github.com/appcelerator/amp/api/rpc/stats"
	"github.com/appcelerator/amp/cmd/ampbeat/config"
	"github.com/appcelerator/amp/pkg/nats-streaming"
	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/golang/protobuf/proto"
	"github.com/nats-io/go-nats-streaming"
)

// Ampbeat amp libbeat type
type Ampbeat struct {
	done          chan struct{}
	config        config.Config
	client        publisher.Client
	natsStreaming *ns.NatsStreaming
}

var bt = &Ampbeat{
	done: make(chan struct{}),
}

// New Creates beater
func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
	bt.config = config.DefaultConfig
	if err := cfg.Unpack(&bt.config); err != nil {
		return nil, fmt.Errorf("Error reading config file: %v", err)
	}

	// NATS Connect
	hostname, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("Unable to get hostname: %v", err)
	}
	bt.natsStreaming = ns.NewClient(ns.DefaultURL, ns.ClusterID, b.Name+"-"+hostname, time.Minute)
	if err = bt.natsStreaming.Connect(); err != nil {
		return nil, fmt.Errorf("Unable to connect to NATS: %v", err)
	}

	return bt, nil
}

//Run run the beat
func (bt *Ampbeat) Run(b *beat.Beat) error {
	logp.Info("ampbeat is running! Hit CTRL-C to stop it.")
	bt.client = b.Publisher.Connect()

	// logs subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.LogsSubject, logMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to logs subject")

	// metrics subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.MetricsSubject, metricsMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to metrics subject")

	select {
	case <-bt.done:
		return nil
	}
}

//Stop stop the beat
func (bt *Ampbeat) Stop() {
	bt.natsStreaming.Close()
	bt.client.Close()
	close(bt.done)
}

func logMessageHandler(msg *stan.Msg) {
	e := logs.LogEntry{}
	err := proto.Unmarshal(msg.Data, &e)
	if err != nil {
		logp.Err("Error unmarshalling log entry: %s", err.Error())
		return
	}
	timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
	if err != nil {
		timestamp = time.Now()
	}
	event := common.MapStr{
		"@timestamp":           common.Time(timestamp),
		"type":                 "logs",
		"container_id":         e.ContainerId,
		"container_name":       e.ContainerName,
		"container_short_name": e.ContainerShortName,
		"container_state":      e.ContainerState,
		"service_name":         e.ServiceName,
		"service_id":           e.ServiceId,
		"task_id":              e.TaskId,
		"stack_name":           e.StackName,
		"node_id":              e.NodeId,
		"msg":                  e.Msg,
	}
	if len(e.Labels) > 0 {
		event["labels"] = getLabels(e.Labels)
	}
	bt.client.PublishEvent(event)
}

func metricsMessageHandler(msg *stan.Msg) {
	e := stats.MetricsEntry{}
	err := proto.Unmarshal(msg.Data, &e)
	if err != nil {
		logp.Err("Error unmarshalling metrics entry: %s", err.Error())
		return
	}
	timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
	if err != nil {
		timestamp = time.Now()
	}
	event := common.MapStr{
		"@timestamp":           common.Time(timestamp),
		"type":                 "metrics",
		"container_id":         e.ContainerId,
		"container_name":       e.ContainerName,
		"container_short_name": e.ContainerShortName,
		"container_state":      e.ContainerState,
		"service_name":         e.ServiceName,
		"service_id":           e.ServiceId,
		"task_id":              e.TaskId,
		"stack_name":           e.StackName,
		"node_id":              e.NodeId,
	}
	if len(e.Labels) > 0 {
		event["labels"] = getLabels(e.Labels)
	}
	if e.Cpu != nil {
		event["cpu"] = common.MapStr{
			"total_usage":          e.Cpu.TotalUsage,
			"usage_in_kernel_mode": e.Cpu.UsageInKernelMode,
			"usage_in_user_mode":   e.Cpu.UsageInUserMode,
		}
	}
	if e.Io != nil {
		event["io"] = common.MapStr{
			"read":  e.Io.Read,
			"write": e.Io.Write,
			"total": e.Io.Total,
		}
	}
	if e.Mem != nil {
		event["mem"] = common.MapStr{
			"fail_count": e.Mem.Failcnt,
			"limit":      e.Mem.Limit,
			"max_usage":  e.Mem.Maxusage,
			"usage":      e.Mem.Usage,
			"usage_pct":  e.Mem.UsageP,
		}
	}
	if e.Net != nil {
		event["net"] = common.MapStr{
			"total_bytes": e.Net.TotalBytes,
			"rx_bytes":    e.Net.RxBytes,
			"rx_dropped":  e.Net.RxDropped,
			"rx_errors":   e.Net.RxErrors,
			"rx_packets":  e.Net.RxPackets,
			"tx_bytes":    e.Net.TxBytes,
			"tx_dropped":  e.Net.TxDropped,
			"tx_errors":   e.Net.TxErrors,
			"tx_packets":  e.Net.TxPackets,
		}
	}
	bt.client.PublishEvent(event)
}

// Need to replace '.'' by '-'', becuse '.' is used by ES for object operator it can't be part of a field name.
func getLabels(labels map[string]string) common.MapStr {
	labmap := common.MapStr{}
	for name, value := range labels {
		if !strings.HasPrefix(name, "com.docker.") {
			name = strings.Replace(name, ".", "-", -1)
			labmap[name] = value
		}
	}
	return labmap
}
<commit_msg>ampbet side of the agent improvement<commit_after>package beater

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

	"github.com/appcelerator/amp/api/rpc/logs"
	"github.com/appcelerator/amp/api/rpc/stats"
	"github.com/appcelerator/amp/cmd/ampbeat/config"
	"github.com/appcelerator/amp/pkg/nats-streaming"
	"github.com/elastic/beats/libbeat/beat"
	"github.com/elastic/beats/libbeat/common"
	"github.com/elastic/beats/libbeat/logp"
	"github.com/elastic/beats/libbeat/publisher"
	"github.com/golang/protobuf/proto"
	"github.com/nats-io/go-nats-streaming"
)

// Ampbeat amp libbeat type
type Ampbeat struct {
	done          chan struct{}
	config        config.Config
	client        publisher.Client
	natsStreaming *ns.NatsStreaming
}

var bt = &Ampbeat{
	done: make(chan struct{}),
}

// New Creates beater
func New(b *beat.Beat, cfg *common.Config) (beat.Beater, error) {
	bt.config = config.DefaultConfig
	if err := cfg.Unpack(&bt.config); err != nil {
		return nil, fmt.Errorf("Error reading config file: %v", err)
	}

	// NATS Connect
	hostname, err := os.Hostname()
	if err != nil {
		return nil, fmt.Errorf("Unable to get hostname: %v", err)
	}
	bt.natsStreaming = ns.NewClient(ns.DefaultURL, ns.ClusterID, b.Name+"-"+hostname, time.Minute)
	if err = bt.natsStreaming.Connect(); err != nil {
		return nil, fmt.Errorf("Unable to connect to NATS: %v", err)
	}

	return bt, nil
}

//Run run the beat
func (bt *Ampbeat) Run(b *beat.Beat) error {
	logp.Info("ampbeat is running! Hit CTRL-C to stop it.")
	bt.client = b.Publisher.Connect()

	// logs subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.LogsSubject, logMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to logs subject")

	// metrics subscription
	if _, err := bt.natsStreaming.GetClient().Subscribe(ns.MetricsSubject, metricsMessageHandler, stan.DeliverAllAvailable()); err != nil {
		return fmt.Errorf("Unable to subscribe to subject: %v", err)
	}
	logp.Info("Succesfully subscribed to metrics subject")

	select {
	case <-bt.done:
		return nil
	}
}

//Stop stop the beat
func (bt *Ampbeat) Stop() {
	bt.natsStreaming.Close()
	bt.client.Close()
	close(bt.done)
}

func logMessageHandler(msg *stan.Msg) {
	list := logs.GetReply{}
	err := proto.Unmarshal(msg.Data, &list)
	if err != nil {
		logp.Err("Error unmarshalling log entry: %s", err.Error())
		return
	}
	for _, e := range list.Entries {
		timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
		if err != nil {
			timestamp = time.Now()
		}
		event := common.MapStr{
			"@timestamp":           common.Time(timestamp),
			"type":                 "logs",
			"container_id":         e.ContainerId,
			"container_name":       e.ContainerName,
			"container_short_name": e.ContainerShortName,
			"container_state":      e.ContainerState,
			"service_name":         e.ServiceName,
			"service_id":           e.ServiceId,
			"task_id":              e.TaskId,
			"stack_name":           e.StackName,
			"node_id":              e.NodeId,
			"msg":                  e.Msg,
		}
		if len(e.Labels) > 0 {
			event["labels"] = getLabels(e.Labels)
		}
		bt.client.PublishEvent(event)
	}
}

func metricsMessageHandler(msg *stan.Msg) {
	list := stats.StatsReply{}
	err := proto.Unmarshal(msg.Data, &list)
	if err != nil {
		logp.Err("Error unmarshalling metrics entry: %s", err.Error())
		return
	}
	for _, e := range list.Entries {
		timestamp, err := time.Parse(time.RFC3339Nano, e.Timestamp)
		if err != nil {
			timestamp = time.Now()
		}
		event := common.MapStr{
			"@timestamp":           common.Time(timestamp),
			"type":                 "metrics",
			"container_id":         e.ContainerId,
			"container_name":       e.ContainerName,
			"container_short_name": e.ContainerShortName,
			"container_state":      e.ContainerState,
			"service_name":         e.ServiceName,
			"service_id":           e.ServiceId,
			"task_id":              e.TaskId,
			"stack_name":           e.StackName,
			"node_id":              e.NodeId,
		}
		if len(e.Labels) > 0 {
			event["labels"] = getLabels(e.Labels)
		}
		if e.Cpu != nil {
			event["cpu"] = common.MapStr{
				"total_usage":          e.Cpu.TotalUsage,
				"usage_in_kernel_mode": e.Cpu.UsageInKernelMode,
				"usage_in_user_mode":   e.Cpu.UsageInUserMode,
			}
		}
		if e.Io != nil {
			event["io"] = common.MapStr{
				"read":  e.Io.Read,
				"write": e.Io.Write,
				"total": e.Io.Total,
			}
		}
		if e.Mem != nil {
			event["mem"] = common.MapStr{
				"fail_count": e.Mem.Failcnt,
				"limit":      e.Mem.Limit,
				"max_usage":  e.Mem.Maxusage,
				"usage":      e.Mem.Usage,
				"usage_pct":  e.Mem.UsageP,
			}
		}
		if e.Net != nil {
			event["net"] = common.MapStr{
				"total_bytes": e.Net.TotalBytes,
				"rx_bytes":    e.Net.RxBytes,
				"rx_dropped":  e.Net.RxDropped,
				"rx_errors":   e.Net.RxErrors,
				"rx_packets":  e.Net.RxPackets,
				"tx_bytes":    e.Net.TxBytes,
				"tx_dropped":  e.Net.TxDropped,
				"tx_errors":   e.Net.TxErrors,
				"tx_packets":  e.Net.TxPackets,
			}
		}
		bt.client.PublishEvent(event)
	}
}

// Need to replace '.'' by '-'', becuse '.' is used by ES for object operator it can't be part of a field name.
func getLabels(labels map[string]string) common.MapStr {
	labmap := common.MapStr{}
	for name, value := range labels {
		if !strings.HasPrefix(name, "com.docker.") {
			name = strings.Replace(name, ".", "-", -1)
			labmap[name] = value
		}
	}
	return labmap
}
<|endoftext|>