feat: graceful shutdown #1977

common/constant/key.go

@@ -109,6 +109,7 @@ const (
 	GracefulShutdownProviderFilterKey    = "pshutdown"
 	GracefulShutdownConsumerFilterKey    = "cshutdown"
 	GracefulShutdownFilterShutdownConfig = "GracefulShutdownFilterShutdownConfig"
+	GracefulShutdownClosingKey           = "closing"
 	HystrixConsumerFilterKey             = "hystrix_consumer"
 	HystrixProviderFilterKey             = "hystrix_provider"
 	MetricsFilterKey                     = "metrics"

common/extension/graceful_shutdown.go

@@ -19,35 +19,41 @@ package extension
 
 import (
 	"container/list"
+	"context"
 )
 
-var customShutdownCallbacks = list.New()
-
-/**
- * AddCustomShutdownCallback
- * you should not make any assumption about the order.
- * For example, if you have more than one callbacks, and you wish the order is:
- * callback1()
- * callback2()
- * ...
- * callbackN()
- * Then you should put then together:
- * func callback() {
- *     callback1()
- *     callback2()
- *     ...
- *     callbackN()
- * }
- * I think the order of custom callbacks should be decided by the users.
- * Even though I can design a mechanism to support the ordered custom callbacks,
- * the benefit of that mechanism is low.
- * And it may introduce much complication for another users.
- */
+// GracefulShutdownCallback is the callback for graceful shutdown
+// name: protocol name such as "grpc", "tri", "dubbo"
+// returns error if notify failed
+type GracefulShutdownCallback func(ctx context.Context) error
+
+var (
+	customShutdownCallbacks   = list.New()
+	gracefulShutdownCallbacks = make(map[string]GracefulShutdownCallback)
+)
+
+// AddCustomShutdownCallback adds custom shutdown callback
 func AddCustomShutdownCallback(callback func()) {
 	customShutdownCallbacks.PushBack(callback)
 }
 
-// GetAllCustomShutdownCallbacks gets all custom shutdown callbacks
+// GetAllCustomShutdownCallbacks returns all custom shutdown callbacks
 func GetAllCustomShutdownCallbacks() *list.List {
 	return customShutdownCallbacks
 }
+
+// SetGracefulShutdownCallback sets protocol-level graceful shutdown callback
+func SetGracefulShutdownCallback(name string, f GracefulShutdownCallback) {
+	gracefulShutdownCallbacks[name] = f
+}
+
+// GetGracefulShutdownCallback returns protocol's graceful shutdown callback
+func GetGracefulShutdownCallback(name string) (GracefulShutdownCallback, bool) {
+	f, ok := gracefulShutdownCallbacks[name]
+	return f, ok
+}
+
+// GetAllGracefulShutdownCallbacks returns all protocol's graceful shutdown callbacks
+func GetAllGracefulShutdownCallbacks() map[string]GracefulShutdownCallback {
+	return gracefulShutdownCallbacks
+}

config/graceful_shutdown.go

@@ -54,7 +54,6 @@ import (
 const defaultShutDownTime = time.Second * 60
 
 func gracefulShutdownInit() {
-	// retrieve ShutdownConfig for gracefulShutdownFilter
 	gracefulShutdownConsumerFilter, exist := extension.GetFilter(constant.GracefulShutdownConsumerFilterKey)
 	if !exist {
 		return
@@ -63,6 +62,7 @@ func gracefulShutdownInit() {
 	if !exist {
 		return
 	}
+	// retrieve ShutdownConfig for gracefulShutdownFilter
 	if filter, ok := gracefulShutdownConsumerFilter.(Setter); ok && rootConfig.Shutdown != nil {
 		filter.Set(constant.GracefulShutdownFilterShutdownConfig, GetShutDown())
 	}
@@ -84,7 +84,6 @@ func gracefulShutdownInit() {
 				os.Exit(0)
 			})
 			BeforeShutdown()
-			// those signals' original behavior is exit with dump ths stack, so we try to keep the behavior
 			for _, dumpSignal := range DumpHeapShutdownSignals {
 				if sig == dumpSignal {
 					debug.WriteHeapDump(os.Stdout.Fd())
@@ -99,14 +98,8 @@ func gracefulShutdownInit() {
 // BeforeShutdown provides processing flow before shutdown
 func BeforeShutdown() {
 	destroyAllRegistries()
-	// waiting for a short time so that the clients have enough time to get the notification that server shutdowns
-	// The value of configuration depends on how long the clients will get notification.
 	waitAndAcceptNewRequests()
-
-	// reject sending/receiving the new request, but keeping waiting for accepting requests
 	waitForSendingAndReceivingRequests()
-
-	// destroy all protocols
 	destroyProtocols()
 
 	logger.Info("Graceful shutdown --- Execute the custom callbacks.")
@@ -142,7 +135,6 @@ func destroyProtocols() {
 func destroyProviderProtocols(consumerProtocols *gxset.HashSet) {
 	logger.Info("Graceful shutdown --- First destroy provider's protocols. ")
 	for _, protocol := range rootConfig.Protocols {
-		// the protocol is the consumer's protocol too, we can not destroy it.
 		if consumerProtocols.Contains(protocol.Name) {
 			continue
 		}
@@ -216,7 +208,7 @@ func waitingConsumerProcessedTimeout(shutdownConfig *ShutdownConfig) {
 }
 
 func totalTimeout() time.Duration {
-	timeout := defaultShutDownTime
+	timeout := defaultShutDownTime // 60s
 	if rootConfig.Shutdown != nil && rootConfig.Shutdown.GetTimeout() > timeout {
 		timeout = rootConfig.Shutdown.GetTimeout()
 	}

config/root_config.go

@@ -135,7 +135,6 @@ func (rc *RootConfig) Init() error {
 			return err
 		}
 	}
-
 	if err := rc.Application.Init(); err != nil {
 		return err
 	}
@@ -175,6 +174,7 @@ func (rc *RootConfig) Init() error {
 	if err := rc.MetadataReport.Init(rc); err != nil {
 		return err
 	}
+
 	if err := rc.Otel.Init(rc.Application); err != nil {
 		return err
 	}
@@ -189,6 +189,7 @@ func (rc *RootConfig) Init() error {
 	if err := initRouterConfig(rc); err != nil {
 		return err
 	}
+
 	// provider、consumer must last init
 	if err := rc.Provider.Init(rc); err != nil {
 		return err
@@ -208,6 +209,7 @@ func (rc *RootConfig) Init() error {
 func (rc *RootConfig) Start() {
 	startOnce.Do(func() {
 		gracefulShutdownInit()
+
 		rc.Consumer.Load()
 		rc.Provider.Load()
 		if err := initMetadata(rc); err != nil {

filter/graceful_shutdown/consumer_filter.go

@@ -19,7 +19,11 @@ package graceful_shutdown
 
 import (
 	"context"
+	"errors"
+	"strconv"
+	"strings"
 	"sync"
+	"time"
 )
 
 import (
@@ -50,7 +54,8 @@ func init() {
 }
 
 type consumerGracefulShutdownFilter struct {
-	shutdownConfig *global.ShutdownConfig
+	shutdownConfig  *global.ShutdownConfig
+	closingInvokers sync.Map // map[string]time.Time (url key -> expire time)
 }
 
 func newConsumerGracefulShutdownFilter() filter.Filter {
@@ -62,15 +67,31 @@ func newConsumerGracefulShutdownFilter() filter.Filter {
 	return csf
 }
 
-// Invoke adds the requests count and block the new requests if application is closing
+// Invoke adds the requests count and checks if invoker is closing
 func (f *consumerGracefulShutdownFilter) Invoke(ctx context.Context, invoker base.Invoker, invocation base.Invocation) result.Result {
+	// check if invoker is closing
+	if f.isClosingInvoker(invoker) {
+		logger.Warnf("Graceful shutdown: skipping closing invoker: %s", invoker.GetURL().String())
+		return &result.RPCResult{Err: errors.New("provider is closing")}
+	}
 	f.shutdownConfig.ConsumerActiveCount.Inc()
 	return invoker.Invoke(ctx, invocation)
 }
 
 // OnResponse reduces the number of active processes then return the process result
 func (f *consumerGracefulShutdownFilter) OnResponse(ctx context.Context, result result.Result, invoker base.Invoker, invocation base.Invocation) result.Result {
 	f.shutdownConfig.ConsumerActiveCount.Dec()
+
+	// check closing flag in response
+	if f.isClosingResponse(result) {
+		f.markClosingInvoker(invoker)
+	}
+
+	// handle request error
+	if result.Error() != nil {
+		f.handleRequestError(invoker, result.Error())
+	}
+
 	return result
 }
 
@@ -91,3 +112,89 @@ func (f *consumerGracefulShutdownFilter) Set(name string, conf any) {
 		// do nothing
 	}
 }
+
+// isClosingInvoker checks if invoker is in closing list
+func (f *consumerGracefulShutdownFilter) isClosingInvoker(invoker base.Invoker) bool {
+	key := invoker.GetURL().String()
+	if expireTime, ok := f.closingInvokers.Load(key); ok {
+		if time.Now().Before(expireTime.(time.Time)) {
+			return true
+		}
+		f.closingInvokers.Delete(key)
+	}
+	return false
+}
+
+// isClosingResponse checks if response contains closing flag
+func (f *consumerGracefulShutdownFilter) isClosingResponse(result result.Result) bool {
+	if result != nil && result.Attachments() != nil {
+		if v, ok := result.Attachments()[constant.GracefulShutdownClosingKey]; ok {
+			if v == "true" {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// markClosingInvoker marks invoker as closing and sets available=false
+func (f *consumerGracefulShutdownFilter) markClosingInvoker(invoker base.Invoker) {
+	key := invoker.GetURL().String()
+	expireTime := time.Now().Add(f.getClosingInvokerExpireTime())
+	f.closingInvokers.Store(key, expireTime)
+
+	logger.Infof("Graceful shutdown: marked invoker as closing: %s, will expire at %v, IsAvailable=%v",
+		key, expireTime, invoker.IsAvailable())
+
+	if bi, ok := invoker.(*base.BaseInvoker); ok {
+		bi.SetAvailable(false)
+		logger.Infof("Graceful shutdown: set invoker unavailable: %s, IsAvailable now=%v",
+			key, invoker.IsAvailable())
+	}
+}
+
+func (f *consumerGracefulShutdownFilter) getClosingInvokerExpireTime() time.Duration {
+	if f.shutdownConfig != nil && f.shutdownConfig.ClosingInvokerExpireTime != "" {
+		if duration, err := time.ParseDuration(f.shutdownConfig.ClosingInvokerExpireTime); err == nil && duration > 0 {
+			return duration
+		}
+	}
+	// default 30s, also try parsing numeric string as milliseconds
+	if f.shutdownConfig != nil && f.shutdownConfig.ClosingInvokerExpireTime != "" {
+		if ms, err := strconv.ParseInt(f.shutdownConfig.ClosingInvokerExpireTime, 10, 64); err == nil && ms > 0 {
+			return time.Duration(ms) * time.Millisecond
+		}
+	}
+	return 30 * time.Second
+}
+
+// handleRequestError handles request errors and marks invoker as unavailable for connection errors
+func (f *consumerGracefulShutdownFilter) handleRequestError(invoker base.Invoker, err error) {
+	if err == nil {
+		return
+	}
+
+	// check for connection-related errors
+	errMsg := err.Error()
+	isConnectionError := strings.Contains(errMsg, "client has closed") ||
+		strings.Contains(errMsg, "connection") ||
+		strings.Contains(errMsg, "EOF") ||
+		strings.Contains(errMsg, "broken pipe") ||
+		strings.Contains(errMsg, "gRPC") && strings.Contains(errMsg, "closing") ||
+		strings.Contains(errMsg, "http2") && strings.Contains(errMsg, "close")
+
+	if isConnectionError {
+		key := invoker.GetURL().String()
+		expireTime := time.Now().Add(f.getClosingInvokerExpireTime())
+		f.closingInvokers.Store(key, expireTime)
+
+		logger.Infof("Graceful shutdown: connection error detected for invoker: %s, marking as closing, will expire at %v, IsAvailable=%v",
+			key, expireTime, invoker.IsAvailable())
+
+		if bi, ok := invoker.(*base.BaseInvoker); ok {
+			bi.SetAvailable(false)
+			logger.Infof("Graceful shutdown: set invoker unavailable due to connection error: %s, IsAvailable now=%v",
+				key, invoker.IsAvailable())
+		}
+	}
+}

filter/graceful_shutdown/provider_filter.go

@@ -43,7 +43,6 @@ var (
 )
 
 func init() {
-	// `init()` is performed before config.Load(), so shutdownConfig will be retrieved after config was loaded.
 	extension.SetFilter(constant.GracefulShutdownProviderFilterKey, func() filter.Filter {
 		return newProviderGracefulShutdownFilter()
 	})
@@ -85,6 +84,12 @@ func (f *providerGracefulShutdownFilter) Invoke(ctx context.Context, invoker bas
 // OnResponse reduces the number of active processes then return the process result
 func (f *providerGracefulShutdownFilter) OnResponse(ctx context.Context, result result.Result, invoker base.Invoker, invocation base.Invocation) result.Result {
 	f.shutdownConfig.ProviderActiveCount.Dec()
+
+	// add closing flag to response
+	if f.isClosing() {
+		result.AddAttachment(constant.GracefulShutdownClosingKey, "true")
+	}
+
 	return result
 }
 
@@ -94,15 +99,13 @@ func (f *providerGracefulShutdownFilter) Set(name string, conf any) {
 		switch ct := conf.(type) {
 		case *global.ShutdownConfig:
 			f.shutdownConfig = ct
-		// only for compatibility with old config, able to directly remove after config is deleted
 		case *config.ShutdownConfig:
 			f.shutdownConfig = compatGlobalShutdownConfig(ct)
 		default:
 			logger.Warnf("the type of config for {%s} should be *global.ShutdownConfig", constant.GracefulShutdownFilterShutdownConfig)
 		}
 		return
 	default:
-		// do nothing
 	}
 }
 
@@ -112,3 +115,10 @@ func (f *providerGracefulShutdownFilter) rejectNewRequest() bool {
 	}
 	return f.shutdownConfig.RejectRequest.Load()
 }
+
+func (f *providerGracefulShutdownFilter) isClosing() bool {
+	if f.shutdownConfig == nil {
+		return false
+	}
+	return f.shutdownConfig.Closing.Load()
+}

global/shutdown_config.go

@@ -60,6 +60,10 @@ type ShutdownConfig struct {
 
 	// provider last received request timestamp
 	ProviderLastReceivedRequestTime atomic.Time
+
+	Closing atomic.Bool
+
+	ClosingInvokerExpireTime string `default:"30s" yaml:"closing-invoker-expire-time" json:"closingInvokerExpireTime,omitempty" property:"closingInvokerExpireTime"`
 }
 
 func DefaultShutdownConfig() *ShutdownConfig {
@@ -88,12 +92,14 @@ func (c *ShutdownConfig) Clone() *ShutdownConfig {
 		RejectRequestHandler:        c.RejectRequestHandler,
 		InternalSignal:              newInternalSignal,
 		OfflineRequestWindowTimeout: c.OfflineRequestWindowTimeout,
+		ClosingInvokerExpireTime:    c.ClosingInvokerExpireTime,
 	}
 
 	newShutdownConfig.RejectRequest.Store(c.RejectRequest.Load())
 	newShutdownConfig.ConsumerActiveCount.Store(c.ConsumerActiveCount.Load())
 	newShutdownConfig.ProviderActiveCount.Store(c.ProviderActiveCount.Load())
 	newShutdownConfig.ProviderLastReceivedRequestTime.Store(c.ProviderLastReceivedRequestTime.Load())
+	newShutdownConfig.Closing.Store(c.Closing.Load())
 
 	return newShutdownConfig
 }

go.mod

@@ -10,7 +10,6 @@ require (
 	github.com/apache/dubbo-getty v1.4.10
 	github.com/apache/dubbo-go-hessian2 v1.12.5
 	github.com/apolloconfig/agollo/v4 v4.4.0
-	github.com/cenkalti/backoff/v4 v4.2.1
 	github.com/creasty/defaults v1.5.2
 	github.com/dop251/goja v0.0.0-20240220182346-e401ed450204
 	github.com/dop251/goja_nodejs v0.0.0-20211022123610-8dd9abb0616d
@@ -81,6 +80,7 @@ require (
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/bits-and-blooms/bitset v1.2.0 // indirect
 	github.com/buger/jsonparser v1.1.1 // indirect
+	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/coreos/go-semver v0.3.0 // indirect
 	github.com/coreos/go-systemd/v22 v22.3.2 // indirect