receive/multitsdb: do not delete not uploaded blocks

If a block hasn't been uploaded yet then tell the TSDB layer not to
delete them. This prevents a nasty race where the TSDB layer can delete
a block before the shipper gets to it. I saw this happen with a very
small block.

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>

pkg/receive/multitsdb.go

@@ -15,6 +15,7 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
+	"github.com/oklog/ulid"
 	"github.com/pkg/errors"
 	"go.uber.org/atomic"
 	"golang.org/x/exp/slices"
@@ -154,7 +155,34 @@ type tenant struct {
 	exemplarsTSDB *exemplars.TSDB
 	ship          *shipper.Shipper
 
-	mtx *sync.RWMutex
+	mtx  *sync.RWMutex
+	tsdb *tsdb.DB
+
+	// For tests.
+	blocksToDeleteFn func(db *tsdb.DB) tsdb.BlocksToDeleteFunc
+}
+
+func (t *tenant) blocksToDelete(blocks []*tsdb.Block) map[ulid.ULID]struct{} {
+	t.mtx.RLock()
+	defer t.mtx.RUnlock()
+
+	if t.tsdb == nil {
+		return nil
+	}
+
+	deletable := t.blocksToDeleteFn(t.tsdb)(blocks)
+	if t.ship == nil {
+		return deletable
+	}
+
+	uploaded := t.ship.UploadedBlocks()
+	for deletableID := range deletable {
+		if _, ok := uploaded[deletableID]; !ok {
+			delete(deletable, deletableID)
+		}
+	}
+
+	return deletable
 }
 
 func newTenant() *tenant {
@@ -202,14 +230,15 @@ func (t *tenant) shipper() *shipper.Shipper {
 func (t *tenant) set(storeTSDB *store.TSDBStore, tenantTSDB *tsdb.DB, ship *shipper.Shipper, exemplarsTSDB *exemplars.TSDB) {
 	t.readyS.Set(tenantTSDB)
 	t.mtx.Lock()
-	t.setComponents(storeTSDB, ship, exemplarsTSDB)
+	t.setComponents(storeTSDB, ship, exemplarsTSDB, tenantTSDB)
 	t.mtx.Unlock()
 }
 
-func (t *tenant) setComponents(storeTSDB *store.TSDBStore, ship *shipper.Shipper, exemplarsTSDB *exemplars.TSDB) {
+func (t *tenant) setComponents(storeTSDB *store.TSDBStore, ship *shipper.Shipper, exemplarsTSDB *exemplars.TSDB, tenantTSDB *tsdb.DB) {
 	t.storeTSDB = storeTSDB
 	t.ship = ship
 	t.exemplarsTSDB = exemplarsTSDB
+	t.tsdb = tenantTSDB
 }
 
 func (t *MultiTSDB) Open() error {
@@ -383,9 +412,17 @@ func (t *MultiTSDB) pruneTSDB(ctx context.Context, logger log.Logger, tenantInst
 	tenantTSDB.mtx.Lock()
 	defer tenantTSDB.mtx.Unlock()
 
-	// Lock the entire tenant to make sure the shipper is not running in parallel.
+	// Make sure the shipper is not running in parallel.
 	tenantInstance.mtx.Lock()
-	defer tenantInstance.mtx.Unlock()
+	shipper := tenantInstance.ship
+	tenantInstance.ship = nil
+	tenantInstance.mtx.Unlock()
+
+	defer func() {
+		tenantInstance.mtx.Lock()
+		tenantInstance.ship = shipper
+		tenantInstance.mtx.Unlock()
+	}()
 
 	sinceLastAppendMillis = time.Since(time.UnixMilli(head.MaxTime())).Milliseconds()
 	if sinceLastAppendMillis <= compactThreshold {
@@ -402,8 +439,8 @@ func (t *MultiTSDB) pruneTSDB(ctx context.Context, logger log.Logger, tenantInst
 	}
 
 	level.Info(logger).Log("msg", "Pruning tenant")
-	if tenantInstance.ship != nil {
-		uploaded, err := tenantInstance.ship.Sync(ctx)
+	if shipper != nil {
+		uploaded, err := shipper.Sync(ctx)
 		if err != nil {
 			return false, err
 		}
@@ -422,7 +459,7 @@ func (t *MultiTSDB) pruneTSDB(ctx context.Context, logger log.Logger, tenantInst
 	}
 
 	tenantInstance.readyS.set(nil)
-	tenantInstance.setComponents(nil, nil, nil)
+	tenantInstance.setComponents(nil, nil, nil, nil)
 
 	return true, nil
 }
@@ -574,6 +611,8 @@ func (t *MultiTSDB) startTSDB(logger log.Logger, tenantID string, tenant *tenant
 
 	level.Info(logger).Log("msg", "opening TSDB")
 	opts := *t.tsdbOpts
+	opts.BlocksToDelete = tenant.blocksToDelete
+	tenant.blocksToDeleteFn = tsdb.DefaultBlocksToDelete
 
 	// NOTE(GiedriusS): always set to false to properly handle OOO samples - OOO samples are written into the WBL
 	// which gets later converted into a block. Without setting this flag to false, the block would get compacted

pkg/receive/multitsdb_test.go

@@ -9,17 +9,20 @@ import (
 	"math"
 	"os"
 	"path"
+	"path/filepath"
 	"strings"
 	"testing"
 	"time"
 
 	"github.com/efficientgo/core/testutil"
 	"github.com/go-kit/log"
+	"github.com/oklog/ulid"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/exemplar"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb"
+	"github.com/stretchr/testify/require"
 	"github.com/thanos-io/objstore"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
@@ -28,6 +31,7 @@ import (
 	"github.com/thanos-io/thanos/pkg/component"
 	"github.com/thanos-io/thanos/pkg/exemplars/exemplarspb"
 	"github.com/thanos-io/thanos/pkg/runutil"
+	"github.com/thanos-io/thanos/pkg/shipper"
 	"github.com/thanos-io/thanos/pkg/store"
 	"github.com/thanos-io/thanos/pkg/store/labelpb"
 	"github.com/thanos-io/thanos/pkg/store/storepb"
@@ -829,3 +833,57 @@ func BenchmarkMultiTSDB(b *testing.B) {
 		_, _ = a.Append(0, l, int64(i), float64(i))
 	}
 }
+
+func TestMultiTSDBDoesNotDeleteNotUploadedBlocks(t *testing.T) {
+	tenant := &tenant{}
+
+	t.Run("no blocks", func(t *testing.T) {
+		require.Equal(t, (map[ulid.ULID]struct{})(nil), tenant.blocksToDelete(nil))
+	})
+
+	tenant.tsdb = &tsdb.DB{}
+
+	mockBlockIDs := []ulid.ULID{
+		ulid.MustNew(1, nil),
+		ulid.MustNew(2, nil),
+	}
+
+	t.Run("no shipper", func(t *testing.T) {
+		tenant.blocksToDeleteFn = func(db *tsdb.DB) tsdb.BlocksToDeleteFunc {
+			return func(_ []*tsdb.Block) map[ulid.ULID]struct{} {
+				return map[ulid.ULID]struct{}{
+					mockBlockIDs[0]: {},
+					mockBlockIDs[1]: {},
+				}
+			}
+		}
+
+		require.Equal(t, map[ulid.ULID]struct{}{
+			mockBlockIDs[0]: {},
+			mockBlockIDs[1]: {},
+		}, tenant.blocksToDelete(nil))
+	})
+
+	t.Run("some blocks uploaded", func(t *testing.T) {
+		tenant.blocksToDeleteFn = func(db *tsdb.DB) tsdb.BlocksToDeleteFunc {
+			return func(_ []*tsdb.Block) map[ulid.ULID]struct{} {
+				return map[ulid.ULID]struct{}{
+					mockBlockIDs[0]: {},
+					mockBlockIDs[1]: {},
+				}
+			}
+		}
+
+		td := t.TempDir()
+
+		require.NoError(t, shipper.WriteMetaFile(log.NewNopLogger(), filepath.Join(td, shipper.DefaultMetaFilename), &shipper.Meta{
+			Version:  shipper.MetaVersion1,
+			Uploaded: []ulid.ULID{mockBlockIDs[0]},
+		}))
+
+		tenant.ship = shipper.New(log.NewNopLogger(), nil, td, nil, nil, metadata.BucketUploadSource, nil, false, metadata.NoneFunc, "")
+		require.Equal(t, map[ulid.ULID]struct{}{
+			mockBlockIDs[0]: {},
+		}, tenant.blocksToDelete(nil))
+	})
+}

pkg/shipper/shipper.go

@@ -355,6 +355,21 @@ func (s *Shipper) Sync(ctx context.Context) (uploaded int, err error) {
 	return uploaded, nil
 }
 
+func (s *Shipper) UploadedBlocks() map[ulid.ULID]struct{} {
+	meta, err := ReadMetaFile(s.metadataFilePath)
+	if err != nil {
+		// NOTE(GiedriusS): Sync() will inform users about any problems.
+		return nil
+	}
+
+	ret := make(map[ulid.ULID]struct{}, len(meta.Uploaded))
+	for _, id := range meta.Uploaded {
+		ret[id] = struct{}{}
+	}
+
+	return ret
+}
+
 // sync uploads the block if not exists in remote storage.
 // TODO(khyatisoneji): Double check if block does not have deletion-mark.json for some reason, otherwise log it or return error.
 func (s *Shipper) upload(ctx context.Context, meta *metadata.Meta) error {