use tempdir when xfering tsdb (#1814)

* use tempdir when xfering tsdb

Signed-off-by: Thor <[email protected]>

* review fixes

Signed-off-by: Thor <[email protected]>

* TransferTSDB: added back in files xfer in addition to bytes transfer

Signed-off-by: Thor <[email protected]>

* unit tests: assert the directory was created where expected

Signed-off-by: Thor <[email protected]>

* remove empty dirs, but don't remove non-empty dirs

Signed-off-by: Thor <[email protected]>

Author: thorfour

pkg/ingester/ingester_test.go

@@ -3,8 +3,11 @@ package ingester
 import (
 	"context"
 	"fmt"
+	"io/ioutil"
 	"math"
+	"math/rand"
 	"net/http"
+	"path/filepath"
 	"sort"
 	"strconv"
 	"sync"
@@ -572,3 +575,21 @@ func BenchmarkIngesterPush(b *testing.B) {
 	}
 
 }
+
+func TestRemoveEmptyDir(t *testing.T) {
+
+	// remove dir that dne
+	require.NoError(t, removeEmptyDir(fmt.Sprintf("%v", rand.Int63())))
+
+	// remove empty dir
+	dir, err := ioutil.TempDir("", "TestRemoveEmptyDir")
+	require.NoError(t, err)
+	require.NoError(t, removeEmptyDir(dir))
+
+	// remove non-empty dir
+	dir, err = ioutil.TempDir("", "TestRemoveEmptyDir")
+	require.NoError(t, err)
+
+	ioutil.WriteFile(filepath.Join(dir, "tempfile"), []byte("hello world"), 0777)
+	require.NotNil(t, removeEmptyDir(dir))
+}

pkg/ingester/lifecycle_test.go

@@ -4,6 +4,7 @@ import (
 	"io"
 	"io/ioutil"
 	"math"
+	"os"
 	"testing"
 	"time"
 
@@ -423,6 +424,7 @@ func TestV2IngesterTransfer(t *testing.T) {
 	require.NoError(t, err)
 	dir2, err := ioutil.TempDir("", "tsdb")
 	require.NoError(t, err)
+	require.NoError(t, os.Remove(dir2)) // remove the destination dir so there isn't a move conflict
 
 	// Start the first ingester, and get it into ACTIVE state.
 	cfg1 := defaultIngesterTestConfig()
@@ -521,4 +523,10 @@ func TestV2IngesterTransfer(t *testing.T) {
 	response, err = ing2.Query(ctx, request)
 	require.NoError(t, err)
 	assert.Equal(t, expectedResponse, response)
+
+	// Assert the data is in the expected location of dir2
+	files, err := ioutil.ReadDir(dir2)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(files))
+	require.Equal(t, "1", files[0].Name())
 }

pkg/ingester/transfer.go

@@ -42,6 +42,14 @@ var (
 		Name: "cortex_ingester_received_files",
 		Help: "The total number of files received by this ingester whilst joining",
 	})
+	receivedBytes = prometheus.NewCounter(prometheus.CounterOpts{
+		Name: "cortex_ingester_received_bytes_total",
+		Help: "The total number of bytes received by this ingester whilst joining",
+	})
+	sentBytes = prometheus.NewCounter(prometheus.CounterOpts{
+		Name: "cortex_ingester_sent_bytes_total",
+		Help: "The total number of bytes sent by this ingester whilst leaving",
+	})
 
 	once *sync.Once
 )
@@ -51,7 +59,9 @@ func init() {
 	prometheus.MustRegister(sentChunks)
 	prometheus.MustRegister(receivedChunks)
 	prometheus.MustRegister(sentFiles)
+	prometheus.MustRegister(receivedBytes)
 	prometheus.MustRegister(receivedFiles)
+	prometheus.MustRegister(sentBytes)
 }
 
 // TransferChunks receives all the chunks from another ingester.
@@ -208,6 +218,20 @@ func (i *Ingester) TransferTSDB(stream client.Ingester_TransferTSDBServer) error
 	fromIngesterID := ""
 
 	xfer := func() error {
+
+		// Validate the final directory is empty, if it exists and is empty delete it so a move can succeed
+		err := removeEmptyDir(i.cfg.TSDBConfig.Dir)
+		if err != nil {
+			return errors.Wrap(err, "remove existing TSDB directory")
+		}
+
+		tmpDir, err := ioutil.TempDir("", "tsdb_xfer")
+		if err != nil {
+			return err
+		}
+		defer os.RemoveAll(tmpDir)
+
+		bytesXfer := 0
 		filesXfer := 0
 
 		files := make(map[string]*os.File)
@@ -236,15 +260,14 @@ func (i *Ingester) TransferTSDB(stream client.Ingester_TransferTSDBServer) error
 					return errors.Wrap(err, "TransferTSDB: checkFromIngesterIsInLeavingState")
 				}
 			}
-			filesXfer++
+			bytesXfer += len(f.Data)
 
-			// TODO(thor) To avoid corruption from errors, it's probably best to write to a temp dir, and then move that to the final location
 			createfile := func(f *client.TimeSeriesFile) (*os.File, error) {
-				dir := filepath.Join(i.cfg.TSDBConfig.Dir, filepath.Dir(f.Filename))
+				dir := filepath.Join(tmpDir, filepath.Dir(f.Filename))
 				if err := os.MkdirAll(dir, 0777); err != nil {
 					return nil, errors.Wrap(err, "TransferTSDB: MkdirAll")
 				}
-				file, err := os.Create(filepath.Join(i.cfg.TSDBConfig.Dir, f.Filename))
+				file, err := os.Create(filepath.Join(tmpDir, f.Filename))
 				if err != nil {
 					return nil, errors.Wrap(err, "TransferTSDB: Create")
 				}
@@ -260,7 +283,7 @@ func (i *Ingester) TransferTSDB(stream client.Ingester_TransferTSDBServer) error
 				if err != nil {
 					return err
 				}
-
+				filesXfer++
 				files[f.Filename] = file
 			} else {
 
@@ -275,10 +298,12 @@ func (i *Ingester) TransferTSDB(stream client.Ingester_TransferTSDBServer) error
 			return errors.Wrap(err, "TransferTSDB: ClaimTokensFor")
 		}
 
+		receivedBytes.Add(float64(bytesXfer))
 		receivedFiles.Add(float64(filesXfer))
-		level.Error(util.Logger).Log("msg", "Total files xfer", "from_ingester", fromIngesterID, "num", filesXfer)
+		level.Info(util.Logger).Log("msg", "Total xfer", "from_ingester", fromIngesterID, "files", filesXfer, "bytes", bytesXfer)
 
-		return nil
+		// Move the tmpdir to the final location
+		return os.Rename(tmpDir, i.cfg.TSDBConfig.Dir)
 	}
 
 	if err := i.transfer(stream.Context(), xfer); err != nil {
@@ -644,6 +669,7 @@ func batchSend(batch int, b []byte, stream client.Ingester_TransferTSDBClient, t
 		if err != nil {
 			return err
 		}
+		sentBytes.Add(float64(len(tsfile.Data)))
 	}
 
 	// Send final data
@@ -653,7 +679,19 @@ func batchSend(batch int, b []byte, stream client.Ingester_TransferTSDBClient, t
 		if err != nil {
 			return err
 		}
+		sentBytes.Add(float64(len(tsfile.Data)))
 	}
 
 	return nil
 }
+
+func removeEmptyDir(dir string) error {
+	if _, err := os.Stat(dir); err != nil {
+		if os.IsNotExist(err) {
+			return nil
+		}
+		return err
+	}
+
+	return os.Remove(dir)
+}