This is an automated email from the ASF dual-hosted git repository.
Alanxtl pushed a commit to branch develop
in repository https://gitbox.apache.org/repos/asf/dubbo-go.git
The following commit(s) were added to refs/heads/develop by this push:
new 54b65bdef refactor(extension): use registry for shutdown callbacks
(#3388)
54b65bdef is described below
commit 54b65bdef41abad81eef5ef4de768b5ac5c86dba
Author: 吴杨帆 <[email protected]>
AuthorDate: Sun Jun 14 15:40:42 2026 +0800
refactor(extension): use registry for shutdown callbacks (#3388)
* refactor(extension): use registry for shutdown callbacks
* refactor(extension): simplify shutdown callback registration
---
common/extension/graceful_shutdown.go | 30 +++++-------------------------
common/extension/graceful_shutdown_test.go | 21 +++++++++++++++++++++
common/extension/registry_type.go | 10 ++++++++++
common/extension/registry_type_test.go | 15 ++++++++++++---
4 files changed, 48 insertions(+), 28 deletions(-)
diff --git a/common/extension/graceful_shutdown.go
b/common/extension/graceful_shutdown.go
index 5468ed5de..147c30139 100644
--- a/common/extension/graceful_shutdown.go
+++ b/common/extension/graceful_shutdown.go
@@ -36,8 +36,7 @@ var (
customShutdownCallbacks = list.New()
customShutdownCallbacksLock sync.RWMutex
customShutdownCallbacksMu = &customShutdownCallbacksLock
- gracefulShutdownCallbacksMu sync.RWMutex
- gracefulShutdownCallbacks = make(map[string]GracefulShutdownCallback)
+ gracefulShutdownCallbacks =
NewRegistry[GracefulShutdownCallback]("graceful shutdown callback")
)
/**
@@ -81,41 +80,22 @@ func GetAllCustomShutdownCallbacks() *list.List {
// RegisterGracefulShutdownCallback registers a protocol-level graceful
shutdown callback.
func RegisterGracefulShutdownCallback(name string, f GracefulShutdownCallback)
{
- gracefulShutdownCallbacksMu.Lock()
- defer gracefulShutdownCallbacksMu.Unlock()
-
- if _, exists := gracefulShutdownCallbacks[name]; exists {
+ if !gracefulShutdownCallbacks.RegisterIfAbsent(name, f) {
logger.Warnf("[GracefulShutdown] graceful shutdown callback %q
already registered, duplicate registration ignored", name)
- return
}
-
- gracefulShutdownCallbacks[name] = f
}
// LookupGracefulShutdownCallback returns a protocol graceful shutdown
callback by name.
func LookupGracefulShutdownCallback(name string) (GracefulShutdownCallback,
bool) {
- gracefulShutdownCallbacksMu.RLock()
- defer gracefulShutdownCallbacksMu.RUnlock()
- f, ok := gracefulShutdownCallbacks[name]
- return f, ok
+ return gracefulShutdownCallbacks.Get(name)
}
// UnregisterGracefulShutdownCallback removes a protocol graceful shutdown
callback by name.
func UnregisterGracefulShutdownCallback(name string) {
- gracefulShutdownCallbacksMu.Lock()
- defer gracefulShutdownCallbacksMu.Unlock()
- delete(gracefulShutdownCallbacks, name)
+ gracefulShutdownCallbacks.Unregister(name)
}
// GracefulShutdownCallbacks returns a snapshot of all protocol graceful
shutdown callbacks.
func GracefulShutdownCallbacks() map[string]GracefulShutdownCallback {
- gracefulShutdownCallbacksMu.RLock()
- defer gracefulShutdownCallbacksMu.RUnlock()
-
- callbacks := make(map[string]GracefulShutdownCallback,
len(gracefulShutdownCallbacks))
- for name, callback := range gracefulShutdownCallbacks {
- callbacks[name] = callback
- }
-
- return callbacks
+ return gracefulShutdownCallbacks.Snapshot()
}
diff --git a/common/extension/graceful_shutdown_test.go
b/common/extension/graceful_shutdown_test.go
index 311530936..0fca07c11 100644
--- a/common/extension/graceful_shutdown_test.go
+++ b/common/extension/graceful_shutdown_test.go
@@ -40,6 +40,27 @@ func TestGracefulShutdownCallbacksReturnsSnapshot(t
*testing.T) {
}
}
+func TestRegisterGracefulShutdownCallbackIgnoresDuplicate(t *testing.T) {
+ t.Cleanup(func() {
+ UnregisterGracefulShutdownCallback("duplicate-test")
+ })
+
+ RegisterGracefulShutdownCallback("duplicate-test",
func(context.Context) error {
+ return context.Canceled
+ })
+ RegisterGracefulShutdownCallback("duplicate-test",
func(context.Context) error {
+ return context.DeadlineExceeded
+ })
+
+ callback, ok := LookupGracefulShutdownCallback("duplicate-test")
+ if !ok {
+ t.Fatal("expected registered callback")
+ }
+ if err := callback(context.Background()); err != context.Canceled {
+ t.Fatalf("expected first callback to remain registered, got
%v", err)
+ }
+}
+
func TestGetAllCustomShutdownCallbacksReturnsSnapshot(t *testing.T) {
customShutdownCallbacksMu.Lock()
original := customShutdownCallbacks
diff --git a/common/extension/registry_type.go
b/common/extension/registry_type.go
index 9e601998d..c4f7f7a15 100644
--- a/common/extension/registry_type.go
+++ b/common/extension/registry_type.go
@@ -41,6 +41,16 @@ func (r *Registry[T]) Register(name string, v T) {
r.items[name] = v
}
+func (r *Registry[T]) RegisterIfAbsent(name string, v T) bool {
+ r.mu.Lock()
+ defer r.mu.Unlock()
+ if _, ok := r.items[name]; ok {
+ return false
+ }
+ r.items[name] = v
+ return true
+}
+
func (r *Registry[T]) Get(name string) (T, bool) {
r.mu.RLock()
defer r.mu.RUnlock()
diff --git a/common/extension/registry_type_test.go
b/common/extension/registry_type_test.go
index 0a7aedadc..e3ca761f2 100644
--- a/common/extension/registry_type_test.go
+++ b/common/extension/registry_type_test.go
@@ -38,11 +38,21 @@ func TestRegistryBasicOps(t *testing.T) {
assert.True(t, ok)
assert.Equal(t, 1, v)
+ assert.False(t, r.RegisterIfAbsent("a", 2))
+ v, ok = r.Get("a")
+ assert.True(t, ok)
+ assert.Equal(t, 1, v)
+
+ assert.True(t, r.RegisterIfAbsent("b", 2))
+ v, ok = r.Get("b")
+ assert.True(t, ok)
+ assert.Equal(t, 2, v)
+
must := r.MustGet("a")
assert.Equal(t, 1, must)
snapshot := r.Snapshot()
- assert.Equal(t, map[string]int{"a": 1}, snapshot)
+ assert.Equal(t, map[string]int{"a": 1, "b": 2}, snapshot)
snapshot["a"] = 99
v, ok = r.Get("a")
@@ -50,8 +60,7 @@ func TestRegistryBasicOps(t *testing.T) {
assert.Equal(t, 1, v)
names := r.Names()
- assert.Len(t, names, 1)
- assert.Equal(t, "a", names[0])
+ assert.ElementsMatch(t, []string{"a", "b"}, names)
r.Unregister("a")
_, ok = r.Get("a")
