manifest: improve API around TableStats

Using `TableMetadata.TableStats` is very fragile; any code that uses
it must first check that the stats are valid. Code that doesn't do so
relies on subtle higher level synchronization assumptions.

We improve this with a simple change: we unexport the stats and add an
getter which only returns the `*TableStats` if they have been
populated, and a setter with can be called at most once.

All access through these methods is safe (without assumptions about
any external locks), as long as we populate the stats only once.

Author: RaduBerinde

compaction.go

@@ -2810,10 +2810,12 @@ func (d *DB) runCopyCompaction(
 		SmallestSeqNum:           inputMeta.SmallestSeqNum,
 		LargestSeqNum:            inputMeta.LargestSeqNum,
 		LargestSeqNumAbsolute:    inputMeta.LargestSeqNumAbsolute,
-		Stats:                    inputMeta.Stats,
 		Virtual:                  inputMeta.Virtual,
 		SyntheticPrefixAndSuffix: inputMeta.SyntheticPrefixAndSuffix,
 	}
+	if inputStats, ok := inputMeta.Stats(); ok {
+		newMeta.PopulateStats(inputStats)
+	}
 	if inputMeta.HasPointKeys {
 		newMeta.ExtendPointKeyBounds(c.comparer.Compare,
 			inputMeta.PointKeyBounds.Smallest(),

compaction_picker.go

@@ -701,7 +701,10 @@ func (c *candidateLevelInfo) shouldCompact() bool {
 }
 
 func tableTombstoneCompensation(t *manifest.TableMetadata) uint64 {
-	return t.Stats.PointDeletionsBytesEstimate + t.Stats.RangeDeletionsBytesEstimate
+	if stats, ok := t.Stats(); ok {
+		return stats.PointDeletionsBytesEstimate + stats.RangeDeletionsBytesEstimate
+	}
+	return 0
 }
 
 // tableCompensatedSize returns t's size, including an estimate of the physical
@@ -1199,7 +1202,9 @@ func pickCompactionSeedFile(
 			// out of the overlapping bytes helps prioritize these compactions
 			// that are cheaper than their file sizes suggest.
 			if outputLevel == numLevels-1 && outputFile.LargestSeqNum < earliestSnapshotSeqNum {
-				overlappingBytes -= outputFile.Stats.RangeDeletionsBytesEstimate
+				if stats, ok := outputFile.Stats(); ok {
+					overlappingBytes -= stats.RangeDeletionsBytesEstimate
+				}
 			}
 
 			// If the file in the next level extends beyond f's largest key,
@@ -1550,7 +1555,8 @@ var elisionOnlyAnnotator = &manifest.Annotator[manifest.TableMetadata]{
 			if f.IsCompacting() {
 				return false, true
 			}
-			if !f.StatsValid() {
+			stats, statsValid := f.Stats()
+			if !statsValid {
 				return false, false
 			}
 			// Bottommost files are large and not worthwhile to compact just
@@ -1567,7 +1573,7 @@ var elisionOnlyAnnotator = &manifest.Annotator[manifest.TableMetadata]{
 			// `NumEntries` and `RangeDeletionsBytesEstimate` are both zero) are excluded
 			// from elision-only compactions.
 			// TODO(travers): Consider an alternative heuristic for elision of range-keys.
-			return f.Stats.RangeDeletionsBytesEstimate*10 >= f.Size || f.Stats.NumDeletions*10 > f.Stats.NumEntries, true
+			return stats.RangeDeletionsBytesEstimate*10 >= f.Size || stats.NumDeletions*10 > stats.NumEntries, true
 		},
 		Compare: func(f1 *manifest.TableMetadata, f2 *manifest.TableMetadata) bool {
 			return f1.LargestSeqNum < f2.LargestSeqNum
@@ -1733,6 +1739,7 @@ func (p *compactionPickerByScore) pickTombstoneDensityCompaction(
 	}
 
 	var candidate *manifest.TableMetadata
+	var candidateTombstoneDenseBlocksRatio float64
 	var level int
 	// If a candidate file has a very high overlapping ratio, point tombstones
 	// in it are likely sparse in keyspace even if the sstable itself is tombstone
@@ -1747,18 +1754,20 @@ func (p *compactionPickerByScore) pickTombstoneDensityCompaction(
 	for l := numLevels - 2; l >= 0; l-- {
 		iter := p.vers.Levels[l].Iter()
 		for f := iter.First(); f != nil; f = iter.Next() {
-			if f.IsCompacting() || !f.StatsValid() || f.Size == 0 {
+			if f.IsCompacting() || f.Size == 0 {
 				continue
 			}
-			if f.Stats.TombstoneDenseBlocksRatio < p.opts.Experimental.TombstoneDenseCompactionThreshold {
+			stats, statsValid := f.Stats()
+			if !statsValid || stats.TombstoneDenseBlocksRatio < p.opts.Experimental.TombstoneDenseCompactionThreshold {
 				continue
 			}
 			overlaps := p.vers.Overlaps(lastNonEmptyLevel, f.UserKeyBounds())
 			if float64(overlaps.AggregateSizeSum())/float64(f.Size) > maxOverlappingRatio {
 				continue
 			}
-			if candidate == nil || candidate.Stats.TombstoneDenseBlocksRatio < f.Stats.TombstoneDenseBlocksRatio {
+			if candidate == nil || candidateTombstoneDenseBlocksRatio < stats.TombstoneDenseBlocksRatio {
 				candidate = f
+				candidateTombstoneDenseBlocksRatio = stats.TombstoneDenseBlocksRatio
 				level = l
 			}
 		}

compaction_picker_test.go

@@ -72,7 +72,6 @@ func loadVersion(
 				}
 			}
 
-			var lastFile *manifest.TableMetadata
 			for i := uint64(1); sizes[level] < int64(size); i++ {
 				var key InternalKey
 				if level == 0 {
@@ -81,19 +80,15 @@ func loadVersion(
 				} else {
 					key = base.MakeInternalKey([]byte(fmt.Sprintf("%04d", i)), base.SeqNum(i), InternalKeyKindSet)
 				}
-				m := (&manifest.TableMetadata{
+				m := &manifest.TableMetadata{
 					TableNum:              base.TableNum(uint64(level)*100_000 + i),
 					SmallestSeqNum:        key.SeqNum(),
 					LargestSeqNum:         key.SeqNum(),
 					LargestSeqNumAbsolute: key.SeqNum(),
 					Size:                  1,
-					Stats: manifest.TableStats{
-						RangeDeletionsBytesEstimate: 0,
-					},
-				}).ExtendPointKeyBounds(opts.Comparer.Compare, key, key)
+				}
+				m.ExtendPointKeyBounds(opts.Comparer.Compare, key, key)
 				m.InitPhysicalBacking()
-				m.StatsMarkValid()
-				lastFile = m
 				if size >= 100 {
 					// If the requested size of the level is very large only add a single
 					// file in order to avoid massive blow-up in the number of files in
@@ -108,13 +103,14 @@ func loadVersion(
 						m.ExtendPointKeyBounds(opts.Comparer.Compare, key, endKey)
 					}
 				}
+				if compensation > 0 {
+					m.PopulateStats(&manifest.TableStats{
+						RangeDeletionsBytesEstimate: compensation,
+					})
+				}
 				files[level] = append(files[level], m)
 				sizes[level] += int64(m.Size)
 			}
-			// Let all the compensation be due to the last file.
-			if lastFile != nil && compensation > 0 {
-				lastFile.Stats.RangeDeletionsBytesEstimate = compensation
-			}
 		}
 	}
 
@@ -139,10 +135,11 @@ func parseTableMeta(t *testing.T, s string, opts *Options) (*manifest.TableMetad
 	if len(parts) != 2 {
 		return nil, errors.Errorf("malformed table spec: %s. usage: <optional-file-num>:start.SET.1-end.SET.2", s)
 	}
-	m := (&manifest.TableMetadata{
+	m := &manifest.TableMetadata{
 		TableNum: tableNum,
 		Size:     1028,
-	}).ExtendPointKeyBounds(
+	}
+	m.ExtendPointKeyBounds(
 		opts.Comparer.Compare,
 		base.ParseInternalKey(strings.TrimSpace(parts[0])),
 		base.ParseInternalKey(strings.TrimSpace(parts[1])),
@@ -160,9 +157,10 @@ func parseTableMeta(t *testing.T, s string, opts *Options) (*manifest.TableMetad
 			if err != nil {
 				return nil, err
 			}
-			m.Stats.RangeDeletionsBytesEstimate = uint64(v)
-			m.Stats.NumDeletions = 1 // At least one range del responsible for the deletion bytes.
-			m.StatsMarkValid()
+			m.PopulateStats(&manifest.TableStats{
+				RangeDeletionsBytesEstimate: uint64(v),
+				NumDeletions:                1, // At least one range del responsible for the deletion bytes.
+			})
 		}
 	}
 	m.SmallestSeqNum = m.Smallest().SeqNum()
@@ -1306,8 +1304,10 @@ func TestCompactionPickerCompensatedSize(t *testing.T) {
 		t.Run("", func(t *testing.T) {
 			f := &manifest.TableMetadata{Size: tc.size}
 			f.InitPhysicalBacking()
-			f.Stats.PointDeletionsBytesEstimate = tc.pointDelEstimateBytes
-			f.Stats.RangeDeletionsBytesEstimate = tc.rangeDelEstimateBytes
+			f.PopulateStats(&manifest.TableStats{
+				PointDeletionsBytesEstimate: tc.pointDelEstimateBytes,
+				RangeDeletionsBytesEstimate: tc.rangeDelEstimateBytes,
+			})
 			gotBytes := tableCompensatedSize(f)
 			require.Equal(t, tc.wantBytes, gotBytes)
 		})

compaction_test.go

@@ -3390,18 +3390,17 @@ func TestTombstoneDensityCompactionMoveOptimization(t *testing.T) {
 	meta := &manifest.TableMetadata{
 		TableNum: 1,
 		Size:     1024,
-		Stats: manifest.TableStats{
-			NumEntries:                10,
-			NumDeletions:              8,
-			TombstoneDenseBlocksRatio: 0.9, // Above threshold
-		},
 	}
+	meta.PopulateStats(&manifest.TableStats{
+		NumEntries:                10,
+		NumDeletions:              8,
+		TombstoneDenseBlocksRatio: 0.9, // Above threshold
+	})
 	meta.ExtendPointKeyBounds(opts.Comparer.Compare,
 		base.ParseInternalKey("a.SET.1"),
 		base.ParseInternalKey("z.SET.2"),
 	)
 	meta.InitPhysicalBacking()
-	meta.StatsMarkValid()
 
 	// Set up the version: L4 has the file, L5 and L6 are empty.
 	var files [numLevels][]*manifest.TableMetadata
@@ -3490,18 +3489,17 @@ func TestTombstoneDensityCompactionMoveOptimization_NoMoveWithOverlap(t *testing
 	metaL4 := &manifest.TableMetadata{
 		TableNum: 1,
 		Size:     1024,
-		Stats: manifest.TableStats{
-			NumEntries:                10,
-			NumDeletions:              8,
-			TombstoneDenseBlocksRatio: 0.9, // Above threshold
-		},
 	}
+	metaL4.PopulateStats(&manifest.TableStats{
+		NumEntries:                10,
+		NumDeletions:              8,
+		TombstoneDenseBlocksRatio: 0.9, // Above threshold
+	})
 	metaL4.ExtendPointKeyBounds(opts.Comparer.Compare,
 		base.ParseInternalKey("a.SET.1"),
 		base.ParseInternalKey("z.SET.2"),
 	)
 	metaL4.InitPhysicalBacking()
-	metaL4.StatsMarkValid()
 
 	// Create an overlapping file in L5.
 	metaL5 := &manifest.TableMetadata{
@@ -3513,7 +3511,7 @@ func TestTombstoneDensityCompactionMoveOptimization_NoMoveWithOverlap(t *testing
 		base.ParseInternalKey("z.SET.3"),
 	)
 	metaL5.InitPhysicalBacking()
-	metaL5.StatsMarkValid()
+	metaL5.PopulateStats(&manifest.TableStats{})
 
 	// Set up the version: L4 has metaL4, L5 has metaL5.
 	var files [numLevels][]*manifest.TableMetadata
@@ -3571,18 +3569,17 @@ func TestTombstoneDensityCompactionMoveOptimization_GrandparentOverlapTooLarge(t
 	metaL4 := &manifest.TableMetadata{
 		TableNum: 1,
 		Size:     1024,
-		Stats: manifest.TableStats{
-			NumEntries:                10,
-			NumDeletions:              8,
-			TombstoneDenseBlocksRatio: 0.9,
-		},
 	}
+	metaL4.PopulateStats(&manifest.TableStats{
+		NumEntries:                10,
+		NumDeletions:              8,
+		TombstoneDenseBlocksRatio: 0.9,
+	})
 	metaL4.ExtendPointKeyBounds(opts.Comparer.Compare,
 		base.ParseInternalKey("a.SET.1"),
 		base.ParseInternalKey("z.SET.2"),
 	)
 	metaL4.InitPhysicalBacking()
-	metaL4.StatsMarkValid()
 
 	// Large overlapping file in L6 (grandparent level).
 	metaL6 := &manifest.TableMetadata{
@@ -3594,7 +3591,7 @@ func TestTombstoneDensityCompactionMoveOptimization_GrandparentOverlapTooLarge(t
 		base.ParseInternalKey("z.SET.3"),
 	)
 	metaL6.InitPhysicalBacking()
-	metaL6.StatsMarkValid()
+	metaL6.PopulateStats(&manifest.TableStats{})
 
 	var files [numLevels][]*manifest.TableMetadata
 	files[inputLevel] = []*manifest.TableMetadata{metaL4}
@@ -3634,18 +3631,17 @@ func TestTombstoneDensityCompactionMoveOptimization_BelowDensityThreshold(t *tes
 	meta := &manifest.TableMetadata{
 		TableNum: 1,
 		Size:     1024,
-		Stats: manifest.TableStats{
-			NumEntries:                10,
-			NumDeletions:              5,
-			TombstoneDenseBlocksRatio: 0.5, // Below threshold
-		},
 	}
+	meta.PopulateStats(&manifest.TableStats{
+		NumEntries:                10,
+		NumDeletions:              5,
+		TombstoneDenseBlocksRatio: 0.5, // Below threshold
+	})
 	meta.ExtendPointKeyBounds(opts.Comparer.Compare,
 		base.ParseInternalKey("a.SET.1"),
 		base.ParseInternalKey("z.SET.2"),
 	)
 	meta.InitPhysicalBacking()
-	meta.StatsMarkValid()
 
 	var files [numLevels][]*manifest.TableMetadata
 	files[inputLevel] = []*manifest.TableMetadata{meta}

data_test.go

@@ -1233,28 +1233,33 @@ func runTableStatsCmd(td *datadriven.TestData, d *DB) string {
 				continue
 			}
 
-			if !f.StatsValid() {
+			stats, statsValid := f.Stats()
+			if !statsValid {
 				d.waitTableStats()
+				stats, statsValid = f.Stats()
+				if !statsValid {
+					panic("stats not valid after waitTableStats")
+				}
 			}
 
 			var b bytes.Buffer
-			fmt.Fprintf(&b, "num-entries: %d\n", f.Stats.NumEntries)
-			fmt.Fprintf(&b, "num-deletions: %d\n", f.Stats.NumDeletions)
-			fmt.Fprintf(&b, "num-range-key-sets: %d\n", f.Stats.NumRangeKeySets)
-			fmt.Fprintf(&b, "point-deletions-bytes-estimate: %d\n", f.Stats.PointDeletionsBytesEstimate)
-			fmt.Fprintf(&b, "range-deletions-bytes-estimate: %d\n", f.Stats.RangeDeletionsBytesEstimate)
+			fmt.Fprintf(&b, "num-entries: %d\n", stats.NumEntries)
+			fmt.Fprintf(&b, "num-deletions: %d\n", stats.NumDeletions)
+			fmt.Fprintf(&b, "num-range-key-sets: %d\n", stats.NumRangeKeySets)
+			fmt.Fprintf(&b, "point-deletions-bytes-estimate: %d\n", stats.PointDeletionsBytesEstimate)
+			fmt.Fprintf(&b, "range-deletions-bytes-estimate: %d\n", stats.RangeDeletionsBytesEstimate)
 
 			// If requested, override the tombstone density ratio for testing
 			if forceTombstoneDensityRatio >= 0 {
-				f.Stats.TombstoneDenseBlocksRatio = forceTombstoneDensityRatio
+				stats.TombstoneDenseBlocksRatio = forceTombstoneDensityRatio
 			}
 			// Only include the tombstone-dense-blocks-ratio if it was forced or is non-zero
-			if forceTombstoneDensityRatio >= 0 || f.Stats.TombstoneDenseBlocksRatio > 0 {
-				fmt.Fprintf(&b, "tombstone-dense-blocks-ratio: %0.1f\n", f.Stats.TombstoneDenseBlocksRatio)
+			if forceTombstoneDensityRatio >= 0 || stats.TombstoneDenseBlocksRatio > 0 {
+				fmt.Fprintf(&b, "tombstone-dense-blocks-ratio: %0.1f\n", stats.TombstoneDenseBlocksRatio)
 			}
 
-			if !f.Stats.CompressionStats.IsEmpty() {
-				fmt.Fprintf(&b, "compression: %s\n", f.Stats.CompressionStats.String())
+			if !stats.CompressionStats.IsEmpty() {
+				fmt.Fprintf(&b, "compression: %s\n", stats.CompressionStats.String())
 			}
 
 			return b.String()

db.go

@@ -2350,13 +2350,11 @@ func (d *DB) SSTables(opts ...SSTablesOption) ([][]SSTableInfo, error) {
 					continue
 				}
 			}
-			var tableStats manifest.TableStats
-			if m.StatsValid() {
-				tableStats = m.Stats
-			}
 			destTables[j] = SSTableInfo{
-				TableInfo:  m.TableInfo(),
-				TableStats: tableStats,
+				TableInfo: m.TableInfo(),
+			}
+			if stats, ok := m.Stats(); ok {
+				destTables[j].TableStats = *stats
 			}
 			if opt.withProperties {
 				p, err := d.fileCache.getTableProperties(

file_cache_test.go

@@ -358,7 +358,7 @@ func TestVirtualReadsWiring(t *testing.T) {
 		base.MakeInternalKey([]byte{'a'}, seqNumA, InternalKeyKindSet))
 	v1.ExtendPointKeyBounds(DefaultComparer.Compare, v1.PointKeyBounds.Smallest(), v1.PointKeyBounds.Largest())
 	v1.AttachVirtualBacking(parentFile.TableBacking)
-	v1.Stats.NumEntries = 1
+	v1.PopulateStats(&manifest.TableStats{NumEntries: 1})
 
 	v2 := &manifest.TableMetadata{
 		TableNum:              f2,
@@ -378,7 +378,7 @@ func TestVirtualReadsWiring(t *testing.T) {
 		base.MakeInternalKey([]byte{'k'}, seqNumRangeUnset, InternalKeyKindRangeKeyUnset))
 	v2.ExtendPointKeyBounds(DefaultComparer.Compare, v2.PointKeyBounds.Smallest(), v2.PointKeyBounds.Largest())
 	v2.AttachVirtualBacking(parentFile.TableBacking)
-	v2.Stats.NumEntries = 6
+	v2.PopulateStats(&manifest.TableStats{NumEntries: 6})
 
 	v1.PointKeyBounds.SetInternalKeyBounds(v1.Smallest(), v1.Largest())