[dbnode] Get rid of excessive locking when adding new namespaces

src/dbnode/storage/bootstrap.go

@@ -79,6 +79,7 @@ type bootstrapManager struct {
 	processProvider             bootstrap.ProcessProvider
 	state                       BootstrapState
 	hasPending                  bool
+	pendingOnCompleteFns        []BootstrapCompleteFn
 	sleepFn                     sleepFn
 	nowFn                       clock.NowFn
 	lastBootstrapCompletionTime xtime.UnixNano
@@ -116,23 +117,29 @@ func (m *bootstrapManager) LastBootstrapCompletionTime() (xtime.UnixNano, bool)
 	return bsTime, bsTime > 0
 }
 
-func (m *bootstrapManager) BootstrapEnqueue() *BootstrapAsyncResult {
-	bootstrapAsyncResult := newBootstrapAsyncResult()
-	go func(r *BootstrapAsyncResult) {
-		if result, err := m.startBootstrap(r); err != nil && !result.AlreadyBootstrapping {
+func (m *bootstrapManager) BootstrapEnqueue(
+	opts BootstrapEnqueueOptions,
+) {
+	go func() {
+		result, err := m.startBootstrap(opts.OnCompleteFn)
+		if err != nil && !result.AlreadyBootstrapping {
 			m.instrumentation.emitAndLogInvariantViolation(err, "error bootstrapping")
 		}
-	}(bootstrapAsyncResult)
-	return bootstrapAsyncResult
+	}()
 }
 
 func (m *bootstrapManager) Bootstrap() (BootstrapResult, error) {
-	bootstrapAsyncResult := newBootstrapAsyncResult()
-	return m.startBootstrap(bootstrapAsyncResult)
+	return m.startBootstrap(nil)
 }
 
-func (m *bootstrapManager) startBootstrap(asyncResult *BootstrapAsyncResult) (BootstrapResult, error) {
+func (m *bootstrapManager) startBootstrap(
+	onCompleteFn BootstrapCompleteFn,
+) (BootstrapResult, error) {
 	m.Lock()
+	if onCompleteFn != nil {
+		// Append completion fn if specified.
+		m.pendingOnCompleteFns = append(m.pendingOnCompleteFns, onCompleteFn)
+	}
 	switch m.state {
 	case Bootstrapping:
 		// NB(r): Already bootstrapping, now a consequent bootstrap
@@ -144,26 +151,19 @@ func (m *bootstrapManager) startBootstrap(asyncResult *BootstrapAsyncResult) (Bo
 		m.hasPending = true
 		m.Unlock()
 		result := BootstrapResult{AlreadyBootstrapping: true}
-		asyncResult.bootstrapResult = result
-		asyncResult.bootstrapStarted.Done()
-		asyncResult.bootstrapCompleted.Done()
 		return result, errBootstrapEnqueued
 	default:
 		m.state = Bootstrapping
 	}
 	m.Unlock()
 	// NB(xichen): disable filesystem manager before we bootstrap to minimize
 	// the impact of file operations on bootstrapping performance
+	m.instrumentation.log.Info("disable fileOps and wait")
 	m.mediator.DisableFileOpsAndWait()
 	defer m.mediator.EnableFileOps()
+	m.instrumentation.log.Info("fileOps disabled")
 
 	var result BootstrapResult
-	asyncResult.bootstrapStarted.Done()
-	defer func() {
-		asyncResult.bootstrapResult = result
-		asyncResult.bootstrapCompleted.Done()
-	}()
-
 	// Keep performing bootstraps until none pending and no error returned.
 	for i := 0; true; i++ {
 		// NB(r): Decouple implementation of bootstrap so can override in tests.
@@ -209,7 +209,19 @@ func (m *bootstrapManager) startBootstrap(asyncResult *BootstrapAsyncResult) (Bo
 	m.Lock()
 	m.lastBootstrapCompletionTime = xtime.ToUnixNano(m.nowFn())
 	m.state = Bootstrapped
+	// NB(r): Clear out the pending completion functions and execute them if
+	// needed.
+	pendingOnCompleteFns := m.pendingOnCompleteFns
+	m.pendingOnCompleteFns = nil
 	m.Unlock()
+
+	if len(pendingOnCompleteFns) > 0 {
+		// Execute any on complete functions that were queued.
+		for _, fn := range pendingOnCompleteFns {
+			fn(result)
+		}
+	}
+
 	return result, nil
 }
 

src/dbnode/storage/bootstrap_test.go

@@ -98,9 +98,15 @@ func testDatabaseBootstrapWithBootstrapError(t *testing.T, async bool) {
 
 	var result BootstrapResult
 	if async {
-		asyncResult := bsm.BootstrapEnqueue()
-		asyncResult.WaitForStart()
-		result = asyncResult.Result()
+		var wg sync.WaitGroup
+		wg.Add(1)
+		bsm.BootstrapEnqueue(BootstrapEnqueueOptions{
+			OnCompleteFn: func(r BootstrapResult) {
+				result = r
+				wg.Done()
+			},
+		})
+		wg.Wait()
 	} else {
 		result, err = bsm.Bootstrap()
 		require.NoError(t, err)

src/dbnode/storage/database.go

@@ -292,20 +292,12 @@ func (d *db) UpdateOwnedNamespaces(newNamespaces namespace.Map) error {
 		}
 	}
 
-	d.Lock()
-	defer d.Unlock()
-
+	d.RLock()
 	removes, adds, updates := d.namespaceDeltaWithLock(newNamespaces)
+	noBootstrapNeeded := d.bootstraps == 0
+	d.RUnlock()
 	if err := d.logNamespaceUpdate(removes, adds, updates); err != nil {
-		enrichedErr := fmt.Errorf("unable to log namespace updates: %v", err)
-		d.log.Error(enrichedErr.Error())
-		return enrichedErr
-	}
-
-	// add any namespaces marked for addition
-	if err := d.addNamespacesWithLock(adds); err != nil {
-		enrichedErr := fmt.Errorf("unable to add namespaces: %v", err)
-		d.log.Error(enrichedErr.Error())
+		d.log.Error("unable to log namespace updates", zap.Error(err))
 		return err
 	}
 
@@ -317,14 +309,61 @@ func (d *db) UpdateOwnedNamespaces(newNamespaces namespace.Map) error {
 			"restart the process if you want changes to take effect")
 	}
 
-	// enqueue bootstraps if new namespaces
 	if len(adds) > 0 {
-		d.queueBootstrapWithLock()
-	}
+		if noBootstrapNeeded {
+			d.Lock()
+			defer d.Unlock()
+			return d.addNamespacesWithLock(adds)
+		}
+
+		if mediator := d.mediator; mediator == nil || !mediator.IsOpen() {
+			d.log.Info("mediator is closed, adding namespaces")
+			d.Lock()
+			if err := d.addNamespacesWithLock(adds); err != nil {
+				d.Unlock()
+				d.log.Error("unable to add namespaces", zap.Error(err))
+				return err
+			}
+			d.Unlock()
 
+			d.enqueueBootstrap()
+			return nil
+		}
+
+		// NB: mediator is opened, we need to disable fileOps and wait for all the background processes to complete
+		// so that we could update namespaces safely. Otherwise, there is a high chance in getting
+		// invariant violation panic because cold/warm flush will receive new namespaces
+		// in the middle of their operations.
+		d.disableFileOpsAndWait()
+		// add any namespaces marked for addition
+		d.Lock()
+		if err := d.addNamespacesWithLock(adds); err != nil {
+			d.Unlock()
+			d.log.Error("unable to add namespaces", zap.Error(err))
+			d.enableFileOps()
+			return err
+		}
+		// enqueue bootstrap and enable file ops when bootstrap is completed
+		d.enqueueBootstrapWithLock(d.enableFileOps)
+		d.Unlock()
+	}
 	return nil
 }
 
+func (d *db) disableFileOpsAndWait() {
+	if mediator := d.mediator; mediator != nil && mediator.IsOpen() {
+		d.log.Info("waiting for file ops to be disabled")
+		mediator.DisableFileOpsAndWait()
+	}
+}
+
+func (d *db) enableFileOps() {
+	if mediator := d.mediator; mediator != nil && mediator.IsOpen() {
+		d.log.Info("enabling file ops")
+		mediator.EnableFileOps()
+	}
+}
+
 func (d *db) namespaceDeltaWithLock(newNamespaces namespace.Map) ([]ident.ID, []namespace.Metadata, []namespace.Metadata) {
 	var (
 		existing = d.namespaces
@@ -459,10 +498,20 @@ func (d *db) AssignShardSet(shardSet sharding.ShardSet) {
 	if receivedNewShards {
 		d.lastReceivedNewShards = d.nowFn()
 	}
+	if d.bootstraps == 0 {
+		// no need to enqueue bootstrap, just assign shards and return
+		d.assignShardsWithLock(shardSet)
+		d.Unlock()
+		return
+	}
 	d.Unlock()
 
-	if !d.mediator.IsOpen() {
-		d.assignShardSet(shardSet)
+	if mediator := d.mediator; mediator == nil || !mediator.IsOpen() {
+		d.Lock()
+		d.assignShardsWithLock(shardSet)
+		d.Unlock()
+
+		d.enqueueBootstrap()
 		return
 	}
 
@@ -480,28 +529,32 @@ func (d *db) AssignShardSet(shardSet sharding.ShardSet) {
 }
 
 func (d *db) assignShardSet(shardSet sharding.ShardSet) {
+	d.RLock()
+	receivedNewShards := d.hasReceivedNewShardsWithLock(shardSet)
+	d.RUnlock()
+
+	if receivedNewShards {
+		d.disableFileOpsAndWait()
+
+		d.Lock()
+		d.assignShardsWithLock(shardSet)
+		d.enqueueBootstrapWithLock(d.enableFileOps)
+		d.Unlock()
+		return
+	}
+
 	d.Lock()
-	defer d.Unlock()
+	d.assignShardsWithLock(shardSet)
+	d.Unlock()
+}
 
+func (d *db) assignShardsWithLock(shardSet sharding.ShardSet) {
 	d.log.Info("assigning shards", zap.Uint32s("shards", shardSet.AllIDs()))
-
-	receivedNewShards := d.hasReceivedNewShardsWithLock(shardSet)
 	d.shardSet = shardSet
-
 	for _, elem := range d.namespaces.Iter() {
 		ns := elem.Value()
 		ns.AssignShardSet(shardSet)
 	}
-
-	if receivedNewShards {
-		// Only trigger a bootstrap if the node received new shards otherwise
-		// the nodes will perform lots of small bootstraps (that accomplish nothing)
-		// during topology changes as other nodes mark their shards as available.
-		//
-		// These small bootstraps can significantly delay topology changes as they prevent
-		// the nodes from marking themselves as bootstrapped and durable, for example.
-		d.queueBootstrapWithLock()
-	}
 }
 
 func (d *db) hasReceivedNewShardsWithLock(incoming sharding.ShardSet) bool {
@@ -533,7 +586,12 @@ func (d *db) ShardSet() sharding.ShardSet {
 	return shardSet
 }
 
-func (d *db) queueBootstrapWithLock() {
+func (d *db) enqueueBootstrap() {
+	d.log.Info("enqueuing bootstrap")
+	d.mediator.BootstrapEnqueue(BootstrapEnqueueOptions{})
+}
+
+func (d *db) enqueueBootstrapWithLock(onCompleteFn func()) {
 	// Only perform a bootstrap if at least one bootstrap has already occurred. This enables
 	// the ability to open the clustered database and assign shardsets to the non-clustered
 	// database when it receives an initial topology (as well as topology changes) without
@@ -542,11 +600,16 @@ func (d *db) queueBootstrapWithLock() {
 	// the non-clustered database bootstrapped by assigning it shardsets which will trigger new
 	// bootstraps since d.bootstraps > 0 will be true.
 	if d.bootstraps > 0 {
-		bootstrapAsyncResult := d.mediator.BootstrapEnqueue()
-		// NB(linasn): We need to wait for the bootstrap to start and set it's state to Bootstrapping in order
-		// to safely run fileOps in mediator later.
-		bootstrapAsyncResult.WaitForStart()
+		d.log.Info("enqueuing bootstrap with onComplete function")
+		d.mediator.BootstrapEnqueue(BootstrapEnqueueOptions{
+			OnCompleteFn: func(_ BootstrapResult) {
+				onCompleteFn()
+			},
+		})
+		return
 	}
+
+	onCompleteFn()
 }
 
 func (d *db) Namespace(id ident.ID) (Namespace, bool) {