Make compact lifecycle more flexible to be overridden for sharded compaction

pkg/block/block.go

@@ -37,6 +37,8 @@ const (
 	IndexHeaderFilename = "index-header"
 	// ChunksDirname is the known dir name for chunks with compressed samples.
 	ChunksDirname = "chunks"
+	// PartitionInfoFilename is JSON filename for partition information.
+	PartitionInfoFilename = "partition-info.json"
 
 	// DebugMetas is a directory for debug meta files that happen in the past. Useful for debugging.
 	DebugMetas = "debug/metas"
@@ -153,6 +155,14 @@ func upload(ctx context.Context, logger log.Logger, bkt objstore.Bucket, bdir st
 		return cleanUp(logger, bkt, id, errors.Wrap(err, "upload index"))
 	}
 
+	// level 1 blocks should not have partition info file
+	if meta.Compaction.Level > 1 {
+		if err := objstore.UploadFile(ctx, logger, bkt, filepath.Join(bdir, PartitionInfoFilename), path.Join(id.String(), PartitionInfoFilename)); err != nil {
+			// Don't call cleanUp here. Partition info file acts in a similar way as meta file.
+			return errors.Wrap(err, "upload partition info")
+		}
+	}
+
 	// Meta.json always need to be uploaded as a last item. This will allow to assume block directories without meta file to be pending uploads.
 	if err := bkt.Upload(ctx, path.Join(id.String(), MetaFilename), strings.NewReader(metaEncoded.String())); err != nil {
 		// Don't call cleanUp here. Despite getting error, meta.json may have been uploaded in certain cases,

pkg/block/fetcher.go

@@ -576,6 +576,10 @@ func (f *LabelShardedMetaFilter) Filter(_ context.Context, metas map[ulid.ULID]*
 
 var _ MetadataFilter = &DeduplicateFilter{}
 
+type IDeduplicateFilter interface {
+	DuplicateIDs() []ulid.ULID
+}
+
 // DeduplicateFilter is a BaseFetcher filter that filters out older blocks that have exactly the same data.
 // Not go-routine safe.
 type DeduplicateFilter struct {

pkg/block/partition_info.go

@@ -0,0 +1,79 @@
+package block
+
+import (
+	"encoding/json"
+	"io"
+	"os"
+	"path/filepath"
+
+	"github.com/go-kit/log"
+	"github.com/prometheus/prometheus/tsdb/fileutil"
+
+	"github.com/thanos-io/thanos/pkg/runutil"
+)
+
+type PartitionInfo struct {
+	PartitionedGroupID uint32 `json:"partitionedGroupID"`
+	PartitionCount     int    `json:"partitionCount"`
+	PartitionID        int    `json:"partitionID"`
+}
+
+// WriteToDir writes the encoded partition info into <dir>/partition-info.json.
+func (p PartitionInfo) WriteToDir(logger log.Logger, dir string) error {
+	// Make any changes to the file appear atomic.
+	path := filepath.Join(dir, PartitionInfoFilename)
+	tmp := path + ".tmp"
+
+	f, err := os.Create(tmp)
+	if err != nil {
+		return err
+	}
+
+	if err := p.Write(f); err != nil {
+		runutil.CloseWithLogOnErr(logger, f, "close partition info")
+		return err
+	}
+	if err := f.Close(); err != nil {
+		return err
+	}
+	return renameFile(logger, tmp, path)
+}
+
+// Write writes the given encoded partition info to writer.
+func (p PartitionInfo) Write(w io.Writer) error {
+	enc := json.NewEncoder(w)
+	enc.SetIndent("", "\t")
+	return enc.Encode(&p)
+}
+
+func renameFile(logger log.Logger, from, to string) error {
+	if err := os.RemoveAll(to); err != nil {
+		return err
+	}
+	if err := os.Rename(from, to); err != nil {
+		return err
+	}
+
+	// Directory was renamed; sync parent dir to persist rename.
+	pdir, err := fileutil.OpenDir(filepath.Dir(to))
+	if err != nil {
+		return err
+	}
+
+	if err = fileutil.Fdatasync(pdir); err != nil {
+		runutil.CloseWithLogOnErr(logger, pdir, "close dir")
+		return err
+	}
+	return pdir.Close()
+}
+
+// Read the block partition info from the given reader.
+func ReadPartitionInfo(rc io.ReadCloser) (_ *PartitionInfo, err error) {
+	defer runutil.ExhaustCloseWithErrCapture(&err, rc, "close partition info JSON")
+
+	var p PartitionInfo
+	if err = json.NewDecoder(rc).Decode(&p); err != nil {
+		return nil, err
+	}
+	return &p, nil
+}