db: fix L0Organizer race during logAndApply

In `logAndApply`, we `Apply` the bulk version edit without holding
`DB.mu`. `Apply` internally updates `L0Organizer` which is supposed to
be protected by `DB.mu`.

Note that the same thing was happening before the `L0Organizer`, but
because the incremental updating of sublevels doesn't work with
compactions, we were always generating new `L0Sublevels`, so the
previous version's sublevels were still usable.

We separate the updating of L0 sublevels from `Apply` and split it
into two steps so we can perform as much work as possible outside of
DB.mu.

Author: RaduBerinde

internal/keyspan/keyspanimpl/level_iter_test.go

@@ -314,8 +314,9 @@ func TestLevelIterEquivalence(t *testing.T) {
 			b.AddedTables[6] = amap
 			l0Organizer := manifest.NewL0Organizer(base.DefaultComparer, 0 /* flushSplitBytes */)
 			emptyVersion := manifest.NewInitialVersion(base.DefaultComparer)
-			v, err := b.Apply(emptyVersion, l0Organizer, 0)
+			v, err := b.Apply(emptyVersion, 0)
 			require.NoError(t, err)
+			l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(b, v), v)
 			levelIter.Init(
 				context.Background(),
 				keyspan.SpanIterOptions{}, base.DefaultComparer.Compare, tableNewIters,

internal/manifest/l0_sublevels.go

@@ -2073,6 +2073,7 @@ type L0Organizer struct {
 	cmp             base.Compare
 	formatKey       base.FormatKey
 	flushSplitBytes int64
+	generation      int64
 
 	// levelMetadata is the current L0.
 	levelMetadata LevelMetadata
@@ -2108,68 +2109,108 @@ func NewL0Organizer(comparer *base.Comparer, flushSplitBytes int64) *L0Organizer
 	return o
 }
 
-// SublevelFiles returns the sublevels as LevelSlices. The returned value (both
-// the slice and each LevelSlice) is immutable. The L0Organizer creates new
-// slices every time L0 changes.
-func (o *L0Organizer) SublevelFiles() []LevelSlice {
-	return o.l0Sublevels.Levels
-}
-
-// Update the L0 organizer with the given L0 changes.
-func (o *L0Organizer) Update(
-	addedL0Tables map[base.FileNum]*TableMetadata,
-	deletedL0Tables map[base.FileNum]*TableMetadata,
-	newLevelMeta *LevelMetadata,
-) {
+// PrepareUpdate is the first step in the two-step process to update the
+// L0Organizer. This first step performs as much work as it can without
+// modifying the L0Organizer.
+//
+// This method can be called concurrently with other methods (other than
+// PerformUpdate). It allows doing most of the update work outside an important
+// lock.
+func (o *L0Organizer) PrepareUpdate(bve *BulkVersionEdit, newVersion *Version) L0PreparedUpdate {
+	addedL0Tables := bve.AddedTables[0]
+	deletedL0Tables := bve.DeletedTables[0]
+	newLevelMeta := &newVersion.Levels[0]
 	if invariants.Enabled && invariants.Sometimes(10) {
 		// Verify that newLevelMeta = m.levelMetadata + addedL0Tables - deletedL0Tables.
 		verifyLevelMetadataTransition(&o.levelMetadata, newLevelMeta, addedL0Tables, deletedL0Tables)
 	}
-	o.levelMetadata = *newLevelMeta
+
 	if len(addedL0Tables) == 0 && len(deletedL0Tables) == 0 {
-		return
+		return L0PreparedUpdate{
+			generation:   o.generation,
+			newSublevels: o.l0Sublevels,
+		}
 	}
-	// If we only added tables, try to use addL0Files.
+
 	if len(deletedL0Tables) == 0 {
 		if files, ok := o.l0Sublevels.canUseAddL0Files(addedL0Tables, newLevelMeta); ok {
-			newSublevels := o.l0Sublevels.addL0Files(files, o.flushSplitBytes, newLevelMeta)
-			// In invariants mode, sometimes rebuild from scratch to verify that
-			// AddL0Files did the right thing. Note that NewL0Sublevels updates
-			// fields in TableMetadata like L0Index, so we don't want to do this
-			// every time.
-			if invariants.Enabled && invariants.Sometimes(10) {
-				expectedSublevels, err := newL0Sublevels(newLevelMeta, o.cmp, o.formatKey, o.flushSplitBytes)
-				if err != nil {
-					panic(fmt.Sprintf("error when regenerating sublevels: %s", err))
-				}
-				s1 := describeSublevels(o.formatKey, false /* verbose */, expectedSublevels.Levels)
-				s2 := describeSublevels(o.formatKey, false /* verbose */, newSublevels.Levels)
-				if s1 != s2 {
-					// Add verbosity.
-					s1 := describeSublevels(o.formatKey, true /* verbose */, expectedSublevels.Levels)
-					s2 := describeSublevels(o.formatKey, true /* verbose */, newSublevels.Levels)
-					panic(fmt.Sprintf("incremental L0 sublevel generation produced different output than regeneration: %s != %s", s1, s2))
-				}
+			return L0PreparedUpdate{
+				generation: o.generation,
+				addL0Files: files,
 			}
-			o.l0Sublevels = newSublevels
-			return
 		}
 	}
-	var err error
-	o.l0Sublevels, err = newL0Sublevels(newLevelMeta, o.cmp, o.formatKey, o.flushSplitBytes)
+	newSublevels, err := newL0Sublevels(newLevelMeta, o.cmp, o.formatKey, o.flushSplitBytes)
 	if err != nil {
 		panic(errors.AssertionFailedf("error generating L0Sublevels: %s", err))
 	}
+
+	return L0PreparedUpdate{
+		generation:   o.generation,
+		newSublevels: newSublevels,
+	}
+}
+
+// L0PreparedUpdate is returned by L0Organizer.PrepareUpdate(), to be passed to
+// PerformUpdate().
+type L0PreparedUpdate struct {
+	generation int64
+
+	// Exactly one of the following fields will be set.
+	addL0Files   []*TableMetadata
+	newSublevels *l0Sublevels
+}
+
+// PerformUpdate applies an update the L0 organizer which was previously
+// prepared using PrepareUpdate.
+//
+// Sets newVersion.L0SublevelFiles (which is immutable once set).
+//
+// This method cannot be called concurrently with any other methods.
+func (o *L0Organizer) PerformUpdate(prepared L0PreparedUpdate, newVersion *Version) {
+	if prepared.generation != o.generation {
+		panic("invalid L0 update generation")
+	}
+	o.levelMetadata = newVersion.Levels[0]
+	o.generation++
+	if prepared.addL0Files != nil {
+		newSublevels := o.l0Sublevels.addL0Files(prepared.addL0Files, o.flushSplitBytes, &o.levelMetadata)
+		// In invariants mode, sometimes rebuild from scratch to verify that
+		// AddL0Files did the right thing. Note that NewL0Sublevels updates
+		// fields in TableMetadata like L0Index, so we don't want to do this
+		// every time.
+		if invariants.Enabled && invariants.Sometimes(10) {
+			expectedSublevels, err := newL0Sublevels(&o.levelMetadata, o.cmp, o.formatKey, o.flushSplitBytes)
+			if err != nil {
+				panic(fmt.Sprintf("error when regenerating sublevels: %s", err))
+			}
+			s1 := describeSublevels(o.formatKey, false /* verbose */, expectedSublevels.Levels)
+			s2 := describeSublevels(o.formatKey, false /* verbose */, newSublevels.Levels)
+			if s1 != s2 {
+				// Add verbosity.
+				s1 := describeSublevels(o.formatKey, true /* verbose */, expectedSublevels.Levels)
+				s2 := describeSublevels(o.formatKey, true /* verbose */, newSublevels.Levels)
+				panic(fmt.Sprintf("incremental L0 sublevel generation produced different output than regeneration: %s != %s", s1, s2))
+			}
+		}
+		o.l0Sublevels = newSublevels
+	} else {
+		o.l0Sublevels = prepared.newSublevels
+	}
+	newVersion.L0SublevelFiles = o.l0Sublevels.Levels
 }
 
-// ResetForTesting reinitializes the L0Organizer to reflect a given L0 level.
-func (o *L0Organizer) ResetForTesting(levelMetadata *LevelMetadata) {
-	o.levelMetadata = *levelMetadata
+// ResetForTesting reinitializes the L0Organizer to reflect the given version.
+// Sets v.L0SublevelFiles.
+func (o *L0Organizer) ResetForTesting(v *Version) {
+	o.levelMetadata = v.Levels[0]
+	o.generation = 0
 	var err error
-	o.l0Sublevels, err = newL0Sublevels(levelMetadata, o.cmp, o.formatKey, o.flushSplitBytes)
+	o.l0Sublevels, err = newL0Sublevels(&v.Levels[0], o.cmp, o.formatKey, o.flushSplitBytes)
 	if err != nil {
 		panic(errors.AssertionFailedf("error generating L0Sublevels: %s", err))
 	}
+	v.L0SublevelFiles = o.l0Sublevels.Levels
 }
 
 // verifyLevelMetadataTransition verifies that newLevel matches oldLevel after

internal/manifest/l0_sublevels_test.go

@@ -54,9 +54,10 @@ func readManifest(filename string) (*Version, error) {
 		if err := bve.Accumulate(&ve); err != nil {
 			return nil, err
 		}
-		if v, err = bve.Apply(v, l0Organizer, 32000); err != nil {
+		if v, err = bve.Apply(v, 32000); err != nil {
 			return nil, err
 		}
+		l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, v), v)
 	}
 	return v, nil
 }

internal/manifest/manifest_test.go

@@ -128,7 +128,8 @@ func replayManifest(
 	}
 	l0Organizer := manifest.NewL0Organizer(cmp, opts.FlushSplitBytes)
 	emptyVersion := manifest.NewInitialVersion(cmp)
-	v, err := bve.Apply(emptyVersion, l0Organizer, opts.Experimental.ReadCompactionRate)
+	v, err := bve.Apply(emptyVersion, opts.Experimental.ReadCompactionRate)
 	require.NoError(t, err)
+	l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, v), v)
 	return v, l0Organizer
 }

internal/manifest/version.go

@@ -1214,8 +1214,7 @@ func NewVersionForTesting(
 			v.Levels[l].totalSize += f.Size
 		}
 	}
-	l0Organizer.ResetForTesting(&v.Levels[0])
-	v.L0SublevelFiles = l0Organizer.SublevelFiles()
+	l0Organizer.ResetForTesting(v)
 	return v
 }
 

internal/manifest/version_edit.go

@@ -1149,9 +1149,10 @@ func (b *BulkVersionEdit) Accumulate(ve *VersionEdit) error {
 // the levels.
 //
 // curr may be nil, which is equivalent to a pointer to a zero version.
-func (b *BulkVersionEdit) Apply(
-	curr *Version, l0Organizer *L0Organizer, readCompactionRate int64,
-) (*Version, error) {
+//
+// Not that L0SublevelFiles is not initialized in the returned version; it is
+// the caller's responsibility to set it using L0Organizer.PerformUpdate().
+func (b *BulkVersionEdit) Apply(curr *Version, readCompactionRate int64) (*Version, error) {
 	comparer := curr.cmp
 	v := &Version{
 		cmp: comparer,
@@ -1293,9 +1294,6 @@ func (b *BulkVersionEdit) Apply(
 		}
 	}
 
-	l0Organizer.Update(b.AddedTables[0], b.DeletedTables[0], &v.Levels[0])
-	v.L0SublevelFiles = l0Organizer.SublevelFiles()
-
 	// We maintain stats about active references in blob files and can infer
 	// when a blob file has become a 'zombie,' and is no longer referenced in
 	// the resulting version. However, we expect the caller to explicitly list

internal/manifest/version_edit_test.go

@@ -449,10 +449,11 @@ func TestVersionEditApply(t *testing.T) {
 				if l0Organizer == nil {
 					d.Fatalf(t, "no L0 organizer")
 				}
-				newv, err := bve.Apply(v, l0Organizer, readCompactionRate)
+				newv, err := bve.Apply(v, readCompactionRate)
 				if err != nil {
 					return err.Error()
 				}
+				l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, newv), newv)
 				if saveName != "" {
 					versions[saveName] = newv
 					l0Organizers[saveName] = l0Organizer
@@ -461,7 +462,7 @@ func TestVersionEditApply(t *testing.T) {
 				// Reinitialize the L0 organizer in case we want to use the same version
 				// again (l0Organizer now reflects newv).
 				l0Organizer = NewL0Organizer(base.DefaultComparer, flushSplitBytes)
-				l0Organizer.ResetForTesting(&v.Levels[0])
+				l0Organizer.ResetForTesting(v)
 				l0Organizers[name] = l0Organizer
 
 				return newv.DebugString()

replay/replay.go

@@ -802,10 +802,11 @@ func (r *Runner) prepareWorkloadSteps(ctx context.Context) error {
 				}
 
 				// Apply the edit.
-				v, err = bve.Apply(v, l0Organizer, r.Opts.Experimental.ReadCompactionRate)
+				v, err = bve.Apply(v, r.Opts.Experimental.ReadCompactionRate)
 				if err != nil {
 					return err
 				}
+				l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, v), v)
 				// AddedTablesByFileNum maps file number to table metadata for all added
 				// sstables from accumulated version edits so we must retain it.
 				bve = manifest.BulkVersionEdit{AddedTablesByFileNum: bve.AddedTablesByFileNum}

tool/db.go

@@ -723,10 +723,11 @@ func (d *dbT) runProperties(cmd *cobra.Command, args []string) {
 		}
 		l0Organizer := manifest.NewL0Organizer(cmp, d.opts.FlushSplitBytes)
 		emptyVersion := manifest.NewInitialVersion(cmp)
-		v, err := bve.Apply(emptyVersion, l0Organizer, d.opts.Experimental.ReadCompactionRate)
+		v, err := bve.Apply(emptyVersion, d.opts.Experimental.ReadCompactionRate)
 		if err != nil {
 			return err
 		}
+		l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, v), v)
 
 		objProvider, err := objstorageprovider.Open(objstorageprovider.DefaultSettings(d.opts.FS, dirname))
 		if err != nil {

tool/manifest.go

@@ -243,11 +243,12 @@ func (m *manifestT) runDump(cmd *cobra.Command, args []string) {
 			if comparer != nil {
 				l0Organizer := manifest.NewL0Organizer(comparer, 0 /* flushSplitBytes */)
 				emptyVersion := manifest.NewInitialVersion(comparer)
-				v, err := bve.Apply(emptyVersion, l0Organizer, m.opts.Experimental.ReadCompactionRate)
+				v, err := bve.Apply(emptyVersion, m.opts.Experimental.ReadCompactionRate)
 				if err != nil {
 					fmt.Fprintf(stdout, "%s\n", err)
 					return
 				}
+				l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, v), v)
 				m.printLevels(comparer.Compare, stdout, v)
 			}
 		}()
@@ -628,7 +629,7 @@ func (m *manifestT) runCheck(cmd *cobra.Command, args []string) {
 					l0Organizer = manifest.NewL0Organizer(cmp, 0 /* flushSplitBytes */)
 					v = manifest.NewInitialVersion(cmp)
 				}
-				newv, err := bve.Apply(v, l0Organizer, m.opts.Experimental.ReadCompactionRate)
+				newv, err := bve.Apply(v, m.opts.Experimental.ReadCompactionRate)
 				if err != nil {
 					fmt.Fprintf(stdout, "%s: offset: %d err: %s\n",
 						arg, offset, err)
@@ -648,6 +649,7 @@ func (m *manifestT) runCheck(cmd *cobra.Command, args []string) {
 					ok = false
 					break
 				}
+				l0Organizer.PerformUpdate(l0Organizer.PrepareUpdate(&bve, newv), newv)
 				v = newv
 			}
 		}()