Only force-load child partitions after successful initialization (#8230)

## What changed?
The force-load child partitions mechanism should only happen after
successful initialization of the root.

## Why?
If the root fails to load, things can get stuck in a loop where the root
loads the children and the children cause the root to be loaded again
(from userdata polling).

## How did you test it?
- [x] built
- [ ] run locally and tested manually
- [ ] covered by existing tests
- [ ] added new unit test(s)
- [ ] added new functional test(s)

Author: dnr

service/matching/matching_engine.go

@@ -383,13 +383,26 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 	create bool,
 	loadCause loadCause,
 ) (retPM taskQueuePartitionManager, retCreated bool, retErr error) {
+	var newPM *taskQueuePartitionManagerImpl
+
 	defer func() {
 		if retErr != nil || retPM == nil {
 			return
 		}
 
 		if retErr = retPM.WaitUntilInitialized(ctx); retErr != nil {
 			e.unloadTaskQueuePartition(retPM, unloadCauseInitError)
+			return
+		}
+
+		if retCreated {
+			// Whenever a root partition is loaded, we need to force all child partitions to load.
+			// If there is a backlog of tasks on any child partitions, force loading will ensure
+			// that they can forward their tasks the poller which caused the root partition to be
+			// loaded. These partitions could be managed by this matchingEngineImpl, but are most
+			// likely not. We skip checking and just make gRPC requests to force loading them all.
+			// Note that if retCreated is true, retPM must be newPM, so we can use newPM here.
+			newPM.ForceLoadAllChildPartitions()
 		}
 	}()
 
@@ -414,7 +427,6 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 	tqConfig := newTaskQueueConfig(partition.TaskQueue(), e.config, nsName)
 	tqConfig.loadCause = loadCause
 	logger, throttledLogger, metricsHandler := e.loggerAndMetricsForPartition(nsName, partition, tqConfig)
-	var newPM *taskQueuePartitionManagerImpl
 	onFatalErr := func(cause unloadCause) { newPM.unloadFromEngine(cause) }
 	onUserDataChanged := func() { newPM.userDataChanged() }
 	userDataManager := newUserDataManager(
@@ -453,15 +465,6 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 	e.partitionsLock.Unlock()
 
 	newPM.Start()
-	if newPM.Partition().IsRoot() {
-		// Whenever a root partition is loaded we need to force all other partitions to load.
-		// If there is a backlog of tasks on any child partitions force loading will ensure that they
-		// can forward their tasks the poller which caused the root partition to be loaded.
-		// These partitions could be managed by this matchingEngineImpl, but are most likely not.
-		// We skip checking and just make gRPC requests to force loading them all.
-
-		newPM.ForceLoadAllNonRootPartitions()
-	}
 	return newPM, true, nil
 }
 

service/matching/task_queue_partition_manager.go

@@ -738,10 +738,9 @@ func (pm *taskQueuePartitionManagerImpl) callerInfoContext(ctx context.Context)
 	return headers.SetCallerInfo(ctx, headers.NewBackgroundHighCallerInfo(pm.ns.Name().String()))
 }
 
-// ForceLoadAllNonRootPartitions spins off go routines which make RPC calls to all the
-func (pm *taskQueuePartitionManagerImpl) ForceLoadAllNonRootPartitions() {
+// ForceLoadAllChildPartitions spins off go routines which make RPC calls to all the
+func (pm *taskQueuePartitionManagerImpl) ForceLoadAllChildPartitions() {
 	if !pm.partition.IsRoot() {
-		pm.logger.Info("ForceLoadAllNonRootPartitions called on non-root partition. Prevented circular keep alive (loading) of partitions.")
 		return
 	}