Merge pull request grafana-cold-storage#814 from grafana/importer_recover_from_position

Dieterbe · web-flow · commit 29bd120e59c2 · 2018-01-10T22:19:42.000+01:00
keep import position in whisper reader
diff --git a/cmd/mt-whisper-importer-reader/main.go b/cmd/mt-whisper-importer-reader/main.go
@@ -88,6 +88,11 @@ var (
 		0,
 		"Only import after the specified timestamp",
 	)
+	positionFile = flag.String(
+		"position-file",
+		"",
+		"file to store position and load position from",
+	)
 	verbose = flag.Bool(
 		"verbose",
 		false,
@@ -114,19 +119,28 @@ func main() {
 		panic(fmt.Sprintf("Error when parsing schemas file: %q", err))
 	}
 
+	var pos *posTracker
+	if len(*positionFile) > 0 {
+		pos, err = NewPositionTracker(*positionFile)
+		if err != nil {
+			log.Fatalf("Error instantiating position tracker: %s", err)
+		}
+		defer pos.Close()
+	}
+
 	fileChan := make(chan string)
 
 	wg := &sync.WaitGroup{}
 	wg.Add(*threads)
 	for i := 0; i < *threads; i++ {
-		go processFromChan(fileChan, wg)
+		go processFromChan(pos, fileChan, wg)
 	}
 
-	getFileListIntoChan(fileChan)
+	getFileListIntoChan(pos, fileChan)
 	wg.Wait()
 }
 
-func processFromChan(files chan string, wg *sync.WaitGroup) {
+func processFromChan(pos *posTracker, files chan string, wg *sync.WaitGroup) {
 	tr := &http.Transport{
 		TLSClientConfig: &tls.Config{InsecureSkipVerify: *insecureSSL},
 	}
@@ -194,6 +208,9 @@ func processFromChan(files chan string, wg *sync.WaitGroup) {
 			resp.Body.Close()
 		}
 
+		if pos != nil {
+			pos.Done(file)
+		}
 		processed := atomic.AddUint32(&processedCount, 1)
 		if processed%100 == 0 {
 			skipped := atomic.LoadUint32(&skippedCount)
@@ -362,7 +379,7 @@ func encodedChunksFromPoints(points []whisper.Point, intervalIn, chunkSpan uint3
 }
 
 // scan a directory and feed the list of whisper files relative to base into the given channel
-func getFileListIntoChan(fileChan chan string) {
+func getFileListIntoChan(pos *posTracker, fileChan chan string) {
 	filepath.Walk(
 		*whisperDirectory,
 		func(path string, info os.FileInfo, err error) error {
@@ -375,9 +392,15 @@ func getFileListIntoChan(fileChan chan string) {
 				atomic.AddUint32(&skippedCount, 1)
 				return nil
 			}
-			if len(path) >= 4 && path[len(path)-4:] == ".wsp" {
-				fileChan <- path
+			if len(path) < 4 || path[len(path)-4:] != ".wsp" {
+				return nil
 			}
+			if pos != nil && pos.IsDone(path) {
+				log.Debugf("Skipping file %s because it was listed as already done", path)
+				return nil
+			}
+
+			fileChan <- path
 			return nil
 		},
 	)
diff --git a/cmd/mt-whisper-importer-reader/pos_tracker.go b/cmd/mt-whisper-importer-reader/pos_tracker.go
@@ -0,0 +1,79 @@
+package main
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+	"sync"
+)
+
+type posTracker struct {
+	sync.Mutex
+	file         string
+	fd           *os.File
+	completedMap sync.Map
+	wg           sync.WaitGroup
+}
+
+func NewPositionTracker(file string) (*posTracker, error) {
+	p := &posTracker{file: file}
+
+	fd, err := os.Open(file)
+	if err != nil {
+		if !os.IsNotExist(err) {
+			return nil, err
+		}
+	} else {
+		reader := bufio.NewReader(fd)
+		var path string
+		for {
+			line, isPrefix, err := reader.ReadLine()
+			if err != nil {
+				if err == io.EOF {
+					break
+				}
+				return nil, err
+			}
+
+			path += string(line)
+			if isPrefix {
+				continue
+			} else {
+				p.completedMap.Store(path, struct{}{})
+				path = ""
+			}
+		}
+
+		fd.Close()
+	}
+
+	p.fd, err = os.OpenFile(file, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		return nil, err
+	}
+
+	return p, nil
+}
+
+func (p *posTracker) IsDone(path string) bool {
+	_, ok := p.completedMap.Load(path)
+	return ok
+}
+
+func (p *posTracker) Done(path string) {
+	p.completedMap.Store(path, struct{}{})
+	p.wg.Add(1)
+	go func() {
+		p.Lock()
+		defer p.Unlock()
+		p.fd.WriteString(fmt.Sprintf("%s\n", path))
+		p.fd.Sync()
+		p.wg.Done()
+	}()
+}
+
+func (p *posTracker) Close() {
+	p.wg.Wait()
+	p.fd.Close()
+}
diff --git a/cmd/mt-whisper-importer-reader/pos_tracker_test.go b/cmd/mt-whisper-importer-reader/pos_tracker_test.go
@@ -0,0 +1,59 @@
+package main
+
+import (
+	"os"
+	"testing"
+)
+
+func TestPositionTracker(t *testing.T) {
+	filePath := "/tmp/positionTrackerTest"
+	clearFile := func() { os.Remove(filePath) }
+	clearFile()
+	defer clearFile()
+
+	testValue1 := "file1"
+	testValue2 := "file2"
+	testValue3 := "file3"
+
+	p1, err := NewPositionTracker(filePath)
+	if err != nil {
+		t.Fatalf("Error instantiating position tracker: %s", err)
+	}
+	p1.Done(testValue1)
+	if !p1.IsDone(testValue1) {
+		t.Fatalf("Expected %s to be done, but it was not", testValue1)
+	}
+	if p1.IsDone(testValue2) {
+		t.Fatalf("Expected %s to not be done, but it was", testValue2)
+	}
+	p1.Done(testValue2)
+	if !p1.IsDone(testValue2) {
+		t.Fatalf("Expected %s to be done, but it was not", testValue2)
+	}
+	p1.Close()
+
+	// read the file into new instance of position tracker
+	p2, err := NewPositionTracker(filePath)
+	if err != nil {
+		t.Fatalf("Error instantiating position tracker: %s", err)
+	}
+	if !p2.IsDone(testValue1) || !p2.IsDone(testValue2) {
+		t.Fatalf("Expected %s and %s to be done, but it was not", testValue1, testValue2)
+	}
+
+	if p2.IsDone(testValue3) {
+		t.Fatalf("Expected %s to not be done, but it was", testValue3)
+	}
+
+	p2.Done(testValue3)
+	p2.Close()
+
+	// read the file into new instance of position tracker
+	p3, err := NewPositionTracker(filePath)
+	if err != nil {
+		t.Fatalf("Error instantiating position tracker: %s", err)
+	}
+	if !p3.IsDone(testValue1) || !p3.IsDone(testValue2) || !p3.IsDone(testValue3) {
+		t.Fatalf("Expected %s, %s and %s to be done, but it was not", testValue1, testValue2, testValue3)
+	}
+}
diff --git a/docs/tools.md b/docs/tools.md
@@ -471,6 +471,8 @@ Usage of ./mt-whisper-importer-reader:
     	Prefix to prepend before every metric name, should include the '.' if necessary
   -orgid int
     	Organization ID the data belongs to  (default 1)
+  -position-file string
+    	file to store position and load position from
   -threads int
     	Number of workers threads to process and convert .wsp files (default 10)
   -verbose
