Merge pull request #226 from cstyan/callum-decouple-prom-pkg

decouple tsdb main.go (and tests) from prometheus/pkg

cmd/tsdb/main.go

@@ -14,6 +14,7 @@
 package main
 
 import (
+	"bufio"
 	"flag"
 	"fmt"
 	"io"
@@ -22,15 +23,14 @@ import (
 	"path/filepath"
 	"runtime"
 	"runtime/pprof"
+	"sort"
+	"strings"
 	"sync"
 	"text/tabwriter"
 	"time"
-	"unsafe"
 
 	"github.com/go-kit/kit/log"
 	"github.com/pkg/errors"
-	promlabels "github.com/prometheus/prometheus/pkg/labels"
-	"github.com/prometheus/prometheus/pkg/textparse"
 	"github.com/prometheus/tsdb"
 	"github.com/prometheus/tsdb/labels"
 	"gopkg.in/alecthomas/kingpin.v2"
@@ -302,21 +302,32 @@ func measureTime(stage string, f func()) time.Duration {
 	return time.Since(start)
 }
 
-func readPrometheusLabels(r io.Reader, n int) ([]labels.Labels, error) {
-	b, err := ioutil.ReadAll(r)
-	if err != nil {
-		return nil, err
+func mapToLabels(m map[string]interface{}, l *labels.Labels) {
+	for k, v := range m {
+		*l = append(*l, labels.Label{Name: k, Value: v.(string)})
 	}
+}
+
+func readPrometheusLabels(r io.Reader, n int) ([]labels.Labels, error) {
+	scanner := bufio.NewScanner(r)
 
-	p := textparse.New(b)
-	i := 0
 	var mets []labels.Labels
 	hashes := map[uint64]struct{}{}
+	i := 0
 
-	for p.Next() && i < n {
+	for scanner.Scan() && i < n {
 		m := make(labels.Labels, 0, 10)
-		p.Metric((*promlabels.Labels)(unsafe.Pointer(&m)))
 
+		r := strings.NewReplacer("\"", "", "{", "", "}", "")
+		s := r.Replace(scanner.Text())
+
+		labelChunks := strings.Split(s, ",")
+		for _, labelChunk := range labelChunks {
+			split := strings.Split(labelChunk, ":")
+			m = append(m, labels.Label{Name: split[0], Value: split[1]})
+		}
+		// Order of the k/v labels matters, don't assume we'll always receive them already sorted.
+		sort.Sort(m)
 		h := m.Hash()
 		if _, ok := hashes[h]; ok {
 			continue
@@ -325,7 +336,7 @@ func readPrometheusLabels(r io.Reader, n int) ([]labels.Labels, error) {
 		hashes[h] = struct{}{}
 		i++
 	}
-	return mets, p.Err()
+	return mets, nil
 }
 
 func exitWithError(err error) {

head_test.go

@@ -14,18 +14,17 @@
 package tsdb
 
 import (
-	"io/ioutil"
+	"bufio"
+	"fmt"
 	"math/rand"
 	"os"
+	"sort"
+	"strings"
 	"testing"
-	"unsafe"
 
 	"github.com/pkg/errors"
 	"github.com/prometheus/tsdb/chunks"
 	"github.com/prometheus/tsdb/labels"
-
-	promlabels "github.com/prometheus/prometheus/pkg/labels"
-	"github.com/prometheus/prometheus/pkg/textparse"
 	"github.com/prometheus/tsdb/testutil"
 )
 
@@ -54,20 +53,26 @@ func readPrometheusLabels(fn string, n int) ([]labels.Labels, error) {
 	}
 	defer f.Close()
 
-	b, err := ioutil.ReadAll(f)
-	if err != nil {
-		return nil, err
-	}
+	scanner := bufio.NewScanner(f)
 
-	p := textparse.New(b)
-	i := 0
 	var mets []labels.Labels
 	hashes := map[uint64]struct{}{}
+	i := 0
 
-	for p.Next() && i < n {
+	for scanner.Scan() && i < n {
 		m := make(labels.Labels, 0, 10)
-		p.Metric((*promlabels.Labels)(unsafe.Pointer(&m)))
 
+		r := strings.NewReplacer("\"", "", "{", "", "}", "")
+		s := r.Replace(scanner.Text())
+
+		labelChunks := strings.Split(s, ",")
+		for _, labelChunk := range labelChunks {
+			split := strings.Split(labelChunk, ":")
+			fmt.Println("split: ", split)
+			m = append(m, labels.Label{Name: split[0], Value: split[1]})
+		}
+		// Order of the k/v labels matters, don't assume we'll always receive them already sorted.
+		sort.Sort(m)
 		h := m.Hash()
 		if _, ok := hashes[h]; ok {
 			continue
@@ -76,7 +81,7 @@ func readPrometheusLabels(fn string, n int) ([]labels.Labels, error) {
 		hashes[h] = struct{}{}
 		i++
 	}
-	if err := p.Err(); err != nil {
+	if err != nil {
 		return nil, err
 	}
 	if i != n {

index_test.go

@@ -236,7 +236,7 @@ func TestPersistence_index_e2e(t *testing.T) {
 	testutil.Ok(t, err)
 	defer os.RemoveAll(dir)
 
-	lbls, err := readPrometheusLabels("testdata/20k.series", 20000)
+	lbls, err := readPrometheusLabels("testdata/20kseries.json", 20000)
 	testutil.Ok(t, err)
 
 	// Sort labels as the index writer expects series in sorted order.

labels/labels_test.go

@@ -14,17 +14,14 @@
 package labels
 
 import (
+	"bufio"
 	"fmt"
-	"io/ioutil"
 	"math/rand"
 	"os"
 	"sort"
+	"strings"
 	"testing"
-	"unsafe"
 
-	"github.com/pkg/errors"
-	promlabels "github.com/prometheus/prometheus/pkg/labels"
-	"github.com/prometheus/prometheus/pkg/textparse"
 	"github.com/prometheus/tsdb/testutil"
 )
 
@@ -129,20 +126,25 @@ func readPrometheusLabels(fn string, n int) ([]Labels, error) {
 	}
 	defer f.Close()
 
-	b, err := ioutil.ReadAll(f)
-	if err != nil {
-		return nil, err
-	}
+	scanner := bufio.NewScanner(f)
 
-	p := textparse.New(b)
-	i := 0
 	var mets []Labels
 	hashes := map[uint64]struct{}{}
+	i := 0
 
-	for p.Next() && i < n {
+	for scanner.Scan() && i < n {
 		m := make(Labels, 0, 10)
-		p.Metric((*promlabels.Labels)(unsafe.Pointer(&m)))
 
+		r := strings.NewReplacer("\"", "", "{", "", "}", "")
+		s := r.Replace(scanner.Text())
+
+		labelChunks := strings.Split(s, ",")
+		for _, labelChunk := range labelChunks {
+			split := strings.Split(labelChunk, ":")
+			m = append(m, Label{Name: split[0], Value: split[1]})
+		}
+		// Order of the k/v labels matters, don't assume we'll always receive them already sorted.
+		sort.Sort(m)
 		h := m.Hash()
 		if _, ok := hashes[h]; ok {
 			continue
@@ -151,12 +153,6 @@ func readPrometheusLabels(fn string, n int) ([]Labels, error) {
 		hashes[h] = struct{}{}
 		i++
 	}
-	if err := p.Err(); err != nil {
-		return nil, err
-	}
-	if i != n {
-		return mets, errors.Errorf("requested %d metrics but found %d", n, i)
-	}
 	return mets, nil
 }
 

wal_test.go

@@ -71,7 +71,7 @@ func TestSegmentWAL_Truncate(t *testing.T) {
 		numMetrics = 20000
 		batch      = 100
 	)
-	series, err := readPrometheusLabels("testdata/20k.series", numMetrics)
+	series, err := readPrometheusLabels("testdata/20kseries.json", numMetrics)
 	testutil.Ok(t, err)
 
 	dir, err := ioutil.TempDir("", "test_wal_log_truncate")
@@ -150,7 +150,7 @@ func TestSegmentWAL_Log_Restore(t *testing.T) {
 	)
 	// Generate testing data. It does not make semantical sense but
 	// for the purpose of this test.
-	series, err := readPrometheusLabels("testdata/20k.series", numMetrics)
+	series, err := readPrometheusLabels("testdata/20kseries.json", numMetrics)
 	testutil.Ok(t, err)
 
 	dir, err := ioutil.TempDir("", "test_wal_log_restore")