Calculate head chunk size based on actual disk usage (prometheus#8139)

Signed-off-by: Julien Pivotto <[email protected]>

tsdb/chunks/head_chunks.go

@@ -105,10 +105,6 @@ type ChunkDiskMapper struct {
 	// from which chunks are served till they are flushed and are ready for m-mapping.
 	chunkBuffer *chunkBuffer
 
-	// The total size of bytes in the closed files.
-	// Needed to calculate the total size of all segments on disk.
-	size atomic.Int64
-
 	// If 'true', it indicated that the maxt of all the on-disk files were set
 	// after iterating through all the chunks in those files.
 	fileMaxtSet bool
@@ -181,8 +177,6 @@ func (cdm *ChunkDiskMapper) openMMapFiles() (returnErr error) {
 		chkFileIndices = append(chkFileIndices, seq)
 	}
 
-	cdm.size.Store(int64(0))
-
 	// Check for gaps in the files.
 	sort.Ints(chkFileIndices)
 	if len(chkFileIndices) == 0 {
@@ -209,8 +203,6 @@ func (cdm *ChunkDiskMapper) openMMapFiles() (returnErr error) {
 		if v := int(b.byteSlice.Range(MagicChunksSize, MagicChunksSize+ChunksFormatVersionSize)[0]); v != chunksFormatV1 {
 			return errors.Errorf("%s: invalid chunk format version %d", files[i], v)
 		}
-
-		cdm.size.Add(int64(b.byteSlice.Len()))
 	}
 
 	return nil
@@ -371,7 +363,6 @@ func (cdm *ChunkDiskMapper) cut() (returnErr error) {
 		}
 	}()
 
-	cdm.size.Add(cdm.curFileSize())
 	cdm.curFileNumBytes.Store(int64(n))
 
 	if cdm.curFile != nil {
@@ -727,7 +718,6 @@ func (cdm *ChunkDiskMapper) deleteFiles(removedFiles []int) error {
 			cdm.readPathMtx.Unlock()
 			return err
 		}
-		cdm.size.Sub(int64(cdm.mmappedChunkFiles[seq].byteSlice.Len()))
 		delete(cdm.mmappedChunkFiles, seq)
 		delete(cdm.closers, seq)
 	}
@@ -766,8 +756,8 @@ func (cdm *ChunkDiskMapper) DeleteCorrupted(originalErr error) error {
 }
 
 // Size returns the size of the chunk files.
-func (cdm *ChunkDiskMapper) Size() int64 {
-	return cdm.size.Load() + cdm.curFileSize()
+func (cdm *ChunkDiskMapper) Size() (int64, error) {
+	return fileutil.DirSize(cdm.dir.Name())
 }
 
 func (cdm *ChunkDiskMapper) curFileSize() int64 {

tsdb/db_test.go

@@ -1237,7 +1237,7 @@ func TestSizeRetention(t *testing.T) {
 	// Add some data to the WAL.
 	headApp := db.Head().Appender(context.Background())
 	for _, m := range headBlocks {
-		series := genSeries(100, 10, m.MinTime, m.MaxTime)
+		series := genSeries(100, 10, m.MinTime, m.MaxTime+1)
 		for _, s := range series {
 			it := s.Iterator()
 			for it.Next() {
@@ -1256,8 +1256,12 @@ func TestSizeRetention(t *testing.T) {
 	blockSize := int64(prom_testutil.ToFloat64(db.metrics.blocksBytes)) // Use the actual internal metrics.
 	walSize, err := db.Head().wal.Size()
 	require.NoError(t, err)
-	// Expected size should take into account block size + WAL size
-	expSize := blockSize + walSize
+	cdmSize, err := db.Head().chunkDiskMapper.Size()
+	require.NoError(t, err)
+	require.NotZero(t, cdmSize)
+	// Expected size should take into account block size + WAL size + Head
+	// chunks size
+	expSize := blockSize + walSize + cdmSize
 	actSize, err := fileutil.DirSize(db.Dir())
 	require.NoError(t, err)
 	require.Equal(t, expSize, actSize, "registered size doesn't match actual disk size")
@@ -1270,7 +1274,20 @@ func TestSizeRetention(t *testing.T) {
 	blockSize = int64(prom_testutil.ToFloat64(db.metrics.blocksBytes)) // Use the actual internal metrics.
 	walSize, err = db.Head().wal.Size()
 	require.NoError(t, err)
-	expSize = blockSize + walSize
+	cdmSize, err = db.Head().chunkDiskMapper.Size()
+	require.NoError(t, err)
+	require.NotZero(t, cdmSize)
+	expSize = blockSize + walSize + cdmSize
+	actSize, err = fileutil.DirSize(db.Dir())
+	require.NoError(t, err)
+	require.Equal(t, expSize, actSize, "registered size doesn't match actual disk size")
+
+	// Truncate Chunk Disk Mapper and compare sizes.
+	require.NoError(t, db.Head().chunkDiskMapper.Truncate(900))
+	cdmSize, err = db.Head().chunkDiskMapper.Size()
+	require.NoError(t, err)
+	require.NotZero(t, cdmSize)
+	expSize = blockSize + walSize + cdmSize
 	actSize, err = fileutil.DirSize(db.Dir())
 	require.NoError(t, err)
 	require.Equal(t, expSize, actSize, "registered size doesn't match actual disk size")
@@ -1287,8 +1304,11 @@ func TestSizeRetention(t *testing.T) {
 	blockSize = int64(prom_testutil.ToFloat64(db.metrics.blocksBytes))
 	walSize, err = db.Head().wal.Size()
 	require.NoError(t, err)
+	cdmSize, err = db.Head().chunkDiskMapper.Size()
+	require.NoError(t, err)
+	require.NotZero(t, cdmSize)
 	// Expected size should take into account block size + WAL size
-	expSize = blockSize + walSize
+	expSize = blockSize + walSize + cdmSize
 	actRetentionCount := int(prom_testutil.ToFloat64(db.metrics.sizeRetentionCount))
 	actSize, err = fileutil.DirSize(db.Dir())
 	require.NoError(t, err)

tsdb/head.go

@@ -2373,7 +2373,8 @@ func (h *Head) Size() int64 {
 	if h.wal != nil {
 		walSize, _ = h.wal.Size()
 	}
-	return walSize + h.chunkDiskMapper.Size()
+	cdmSize, _ := h.chunkDiskMapper.Size()
+	return walSize + cdmSize
 }
 
 func (h *RangeHead) Size() int64 {