Fix Go runtime shutdown function ctx

CHANGELOG.md

@@ -4,6 +4,8 @@ All notable changes to this project are documented below.
 The format is based on [keep a changelog](http://keepachangelog.com) and this project uses [semantic versioning](http://semver.org).
 
 ## [Unreleased]
+### Fixed
+- Correctly wire Go runtime shutdown function context.
 
 ## [3.22.0] - 2024-06-09
 ### Added

server/api_test.go

@@ -55,7 +55,7 @@ var (
 	}
 	metrics       = NewLocalMetrics(logger, logger, nil, cfg)
 	storageIdx, _ = NewLocalStorageIndex(logger, nil, &StorageConfig{DisableIndexOnly: false}, metrics)
-	_             = cfg.Validate(logger)
+	_             = ValidateConfig(logger, cfg)
 )
 
 type DummyMessageRouter struct{}

server/runtime_go.go

@@ -2674,6 +2674,7 @@ func (ri *RuntimeGoInitializer) RegisterFleetManager(fleetManager runtime.FleetM
 
 func (ri *RuntimeGoInitializer) RegisterShutdown(fn func(ctx context.Context, logger runtime.Logger, db *sql.DB, nk runtime.NakamaModule)) error {
 	ri.shutdownFunction = func(ctx context.Context) {
+		ctx = NewRuntimeGoContext(ctx, ri.node, ri.version, ri.env, RuntimeExecutionModeShutdown, nil, nil, 0, "", "", nil, "", "", "", "")
 		fn(ctx, ri.logger.WithField("mode", RuntimeExecutionModeShutdown.String()), ri.db, ri.nk)
 	}
 

server/shutdown.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"errors"
 	"go.uber.org/zap"
 	"os"
 	"time"
@@ -14,18 +15,20 @@ func HandleShutdown(ctx context.Context, logger *zap.Logger, matchRegistry Match
 	runtimeShutdownFnDone := make(chan struct{}, 1)
 
 	if graceSeconds != 0 {
-		timer = time.NewTimer(time.Duration(graceSeconds) * time.Second)
-		timerCh = timer.C
-
+		graceDuration := time.Duration(graceSeconds) * time.Second
 		if shutdownFn != nil {
 			go func() {
-				shutdownFn(ctx)
+				shCtx, _ := context.WithTimeoutCause(context.WithoutCancel(ctx), graceDuration, errors.New("grace period expired"))
+				shutdownFn(shCtx)
 				close(runtimeShutdownFnDone)
 			}()
 		} else {
 			close(runtimeShutdownFnDone)
 		}
 
+		timer = time.NewTimer(graceDuration)
+		timerCh = timer.C
+
 		logger.Info("Shutdown started - use CTRL^C to force stop server", zap.Int("grace_period_sec", graceSeconds))
 	} else {
 		// No grace period.