fix: use atomic.Pointer to store EventHandler in case of data race (#30205)

Resolves #30204

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

pkg/config/etcd_source.go

@@ -112,7 +112,7 @@ func (es *EtcdSource) Close() {
 }
 
 func (es *EtcdSource) SetEventHandler(eh EventHandler) {
-	es.configRefresher.eh = eh
+	es.configRefresher.SetEventHandler(eh)
 }
 
 func (es *EtcdSource) UpdateOptions(opts Options) {
@@ -124,9 +124,9 @@ func (es *EtcdSource) UpdateOptions(opts Options) {
 	es.keyPrefix = opts.EtcdInfo.KeyPrefix
 	if es.configRefresher.refreshInterval != opts.EtcdInfo.RefreshInterval {
 		es.configRefresher.stop()
-		eh := es.configRefresher.eh
+		eh := es.configRefresher.GetEventHandler()
 		es.configRefresher = newRefresher(opts.EtcdInfo.RefreshInterval, es.refreshConfigurations)
-		es.configRefresher.eh = eh
+		es.configRefresher.SetEventHandler(eh)
 		es.configRefresher.start(es.GetSourceName())
 	}
 }

pkg/config/etcd_source_test.go

@@ -0,0 +1,108 @@
+// Licensed to the LF AI & Data foundation under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package config
+
+import (
+	"context"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/suite"
+	"go.etcd.io/etcd/server/v3/embed"
+	"go.uber.org/atomic"
+
+	"github.com/milvus-io/milvus/pkg/util/etcd"
+)
+
+type EtcdSourceSuite struct {
+	suite.Suite
+
+	embedEtcdServer *embed.Etcd
+	tempDir         string
+	endpoints       []string
+}
+
+func (s *EtcdSourceSuite) SetupSuite() {
+	// init embed etcd
+	embedServer, tempDir, err := etcd.StartTestEmbedEtcdServer()
+
+	s.Require().NoError(err)
+
+	s.embedEtcdServer = embedServer
+	s.tempDir = tempDir
+	s.endpoints = etcd.GetEmbedEtcdEndpoints(embedServer)
+}
+
+func (s *EtcdSourceSuite) TearDownSuite() {
+	if s.embedEtcdServer != nil {
+		s.embedEtcdServer.Close()
+	}
+	if s.tempDir != "" {
+		os.RemoveAll(s.tempDir)
+	}
+}
+
+func (s *EtcdSourceSuite) TestNewSource() {
+	source, err := NewEtcdSource(&EtcdInfo{
+		Endpoints:       s.endpoints,
+		KeyPrefix:       "by-dev",
+		RefreshInterval: time.Second,
+	})
+	s.NoError(err)
+	s.NotNil(source)
+	source.Close()
+}
+
+func (s *EtcdSourceSuite) TestUpdateOptions() {
+	source, err := NewEtcdSource(&EtcdInfo{
+		Endpoints:       s.endpoints,
+		KeyPrefix:       "test_update_options_1",
+		RefreshInterval: time.Second,
+	})
+	s.Require().NoError(err)
+	s.Require().NotNil(source)
+	defer source.Close()
+
+	called := atomic.NewBool(false)
+
+	handler := NewHandler("test_update_options", func(evt *Event) {
+		called.Store(true)
+	})
+
+	source.SetEventHandler(handler)
+
+	source.UpdateOptions(Options{
+		EtcdInfo: &EtcdInfo{
+			Endpoints:       s.endpoints,
+			KeyPrefix:       "test_update_options_2",
+			RefreshInterval: time.Millisecond * 100,
+		},
+	})
+
+	client, err := etcd.GetRemoteEtcdClient(s.endpoints)
+	s.Require().NoError(err)
+	client.Put(context.Background(), "test_update_options_2/config/abc", "def")
+
+	s.Eventually(func() bool {
+		return called.Load()
+	}, time.Second*2, time.Millisecond*100)
+}
+
+func TestEtcdSource(t *testing.T) {
+	suite.Run(t, new(EtcdSourceSuite))
+}

pkg/config/file_source.go

@@ -93,7 +93,7 @@ func (fs *FileSource) Close() {
 func (fs *FileSource) SetEventHandler(eh EventHandler) {
 	fs.RWMutex.Lock()
 	defer fs.RWMutex.Unlock()
-	fs.configRefresher.eh = eh
+	fs.configRefresher.SetEventHandler(eh)
 }
 
 func (fs *FileSource) UpdateOptions(opts Options) {

pkg/config/refresher.go

@@ -19,6 +19,7 @@ import (
 	"sync"
 	"time"
 
+	"go.uber.org/atomic"
 	"go.uber.org/zap"
 
 	"github.com/milvus-io/milvus/pkg/log"
@@ -28,7 +29,7 @@ type refresher struct {
 	refreshInterval  time.Duration
 	intervalDone     chan struct{}
 	intervalInitOnce sync.Once
-	eh               EventHandler
+	eh               atomic.Pointer[EventHandler]
 
 	fetchFunc func() error
 	stopOnce  sync.Once
@@ -86,10 +87,24 @@ func (r *refresher) fireEvents(name string, source, target map[string]string) er
 		return err
 	}
 	// Generate OnEvent Callback based on the events created
-	if r.eh != nil {
+	ptr := r.eh.Load()
+	if ptr != nil && *ptr != nil {
 		for _, e := range events {
-			r.eh.OnEvent(e)
+			(*ptr).OnEvent(e)
 		}
 	}
 	return nil
 }
+
+func (r *refresher) SetEventHandler(eh EventHandler) {
+	r.eh.Store(&eh)
+}
+
+func (r *refresher) GetEventHandler() EventHandler {
+	var eh EventHandler
+	ptr := r.eh.Load()
+	if ptr != nil {
+		eh = *ptr
+	}
+	return eh
+}