Refactor OnChange to be set after initialization through a seperate

method.

Even with the channel providing guardrails on initialization, the
previous attempt would still cause spurious panic's with newPM
being nil by the time we attempt to call the method on it.

I didn't take a look at the generated assembly but I assume that even
with the channel block there was still a race as to when the address of
newPM was copied from the outer scope's stack into the closure's stack.

This could likely be avoided by taking the address of the pointer and
passing that instead, but started to feel too clever to me.

Author: majiru

service/matching/matching_engine.go

@@ -430,16 +430,10 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 	tqConfig.loadCause = loadCause
 	logger, throttledLogger, metricsHandler := e.loggerAndMetricsForPartition(namespaceEntry, partition, tqConfig)
 	onFatalErr := func(cause unloadCause) { newPM.unloadFromEngine(cause) }
-	done := make(chan struct{})
-	onUserDataChanged := func(from, to *persistencespb.VersionedTaskQueueUserData) {
-		<-done
-		newPM.userDataChanged(from, to)
-	}
 	userDataManager := newUserDataManager(
 		e.taskManager,
 		e.matchingRawClient,
 		onFatalErr,
-		onUserDataChanged,
 		partition,
 		tqConfig,
 		logger,
@@ -456,10 +450,10 @@ func (e *matchingEngineImpl) getTaskQueuePartitionManager(
 		metricsHandler,
 		userDataManager,
 	)
-	close(done)
 	if err != nil {
 		return nil, false, err
 	}
+	userDataManager.SetOnChange(newPM.userDataChanged)
 
 	// If it gets here, write lock and check again in case a task queue is created between the two locks
 	e.partitionsLock.Lock()

service/matching/physical_task_queue_manager_test.go

@@ -85,7 +85,7 @@ func (s *PhysicalTaskQueueManagerTestSuite) SetupTest() {
 	prtn := s.physicalTaskQueueKey.Partition()
 	tqConfig := newTaskQueueConfig(prtn.TaskQueue(), engine.config, nsName)
 	onFatalErr := func(unloadCause) { s.T().Fatal("user data manager called onFatalErr") }
-	udMgr := newUserDataManager(engine.taskManager, engine.matchingRawClient, onFatalErr, nil, prtn, tqConfig, engine.logger, engine.namespaceRegistry)
+	udMgr := newUserDataManager(engine.taskManager, engine.matchingRawClient, onFatalErr, prtn, tqConfig, engine.logger, engine.namespaceRegistry)
 
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Millisecond)
 	prtnMgr, err := newTaskQueuePartitionManager(ctx, engine, ns, prtn, tqConfig, engine.logger, nil, metrics.NoopMetricsHandler, udMgr)

service/matching/task_queue_partition_manager_test.go

@@ -607,6 +607,10 @@ type mockUserDataManager struct {
 	data *persistencespb.VersionedTaskQueueUserData
 }
 
+func (m *mockUserDataManager) SetOnChange(fn UserDataChangedFunc) {
+	// noop
+}
+
 func (m *mockUserDataManager) Start() {
 	// noop
 }

service/matching/user_data_manager.go

@@ -50,6 +50,7 @@ type (
 		// Handles the maybe-long-poll GetUserData RPC.
 		HandleGetUserDataRequest(ctx context.Context, req *matchingservice.GetTaskQueueUserDataRequest) (*matchingservice.GetTaskQueueUserDataResponse, error)
 		CheckTaskQueueUserDataPropagation(context.Context, int64, int, int) error
+		SetOnChange(fn UserDataChangedFunc)
 	}
 
 	UserDataUpdateOptions struct {
@@ -65,6 +66,8 @@ type (
 	// Extra care should be taken to avoid mutating the current user data to avoid keeping uncommitted data in memory.
 	UserDataUpdateFunc func(*persistencespb.TaskQueueUserData) (*persistencespb.TaskQueueUserData, bool, error)
 
+	UserDataChangedFunc func(from, to *persistencespb.VersionedTaskQueueUserData)
+
 	// userDataManager is responsible for fetching and keeping user data up-to-date in-memory
 	// for a given TQ partition.
 	//
@@ -75,7 +78,7 @@ type (
 	userDataManagerImpl struct {
 		lock              sync.Mutex
 		onFatalErr        func(unloadCause)
-		onUserDataChanged func(from, to *persistencespb.VersionedTaskQueueUserData) // if set, call this in new goroutine when user data changes
+		onUserDataChanged UserDataChangedFunc // if set, call this in new goroutine when user data changes
 		partition         tqid.Partition
 		userData          *persistencespb.VersionedTaskQueueUserData
 		userDataChanged   chan struct{}
@@ -110,15 +113,13 @@ func newUserDataManager(
 	store persistence.TaskManager,
 	matchingClient matchingservice.MatchingServiceClient,
 	onFatalErr func(unloadCause),
-	onUserDataChanged func(from, to *persistencespb.VersionedTaskQueueUserData),
 	partition tqid.Partition,
 	config *taskQueueConfig,
 	logger log.Logger,
 	registry namespace.Registry,
 ) *userDataManagerImpl {
 	m := &userDataManagerImpl{
 		onFatalErr:        onFatalErr,
-		onUserDataChanged: onUserDataChanged,
 		partition:         partition,
 		userDataChanged:   make(chan struct{}),
 		config:            config,
@@ -135,6 +136,12 @@ func newUserDataManager(
 	return m
 }
 
+func (m *userDataManagerImpl) SetOnChange(fn UserDataChangedFunc) {
+	m.lock.Lock()
+	m.onUserDataChanged = fn
+	m.lock.Unlock()
+}
+
 func (m *userDataManagerImpl) Start() {
 	if m.store != nil {
 		m.goroGroup.Go(m.loadUserData)

service/matching/user_data_manager_test.go

@@ -66,7 +66,7 @@ func createUserDataManager(
 		onFatalErr = func(unloadCause) { t.Fatal("user data manager called onFatalErr") }
 	}
 
-	return newUserDataManager(tm, testOpts.matchingClientMock, onFatalErr, nil, testOpts.dbq.Partition(), newTaskQueueConfig(testOpts.dbq.Partition().TaskQueue(), testOpts.config, ns), logger, mockNamespaceCache)
+	return newUserDataManager(tm, testOpts.matchingClientMock, onFatalErr, testOpts.dbq.Partition(), newTaskQueueConfig(testOpts.dbq.Partition().TaskQueue(), testOpts.config, ns), logger, mockNamespaceCache)
 }
 
 func TestUserData_LoadOnInit(t *testing.T) {