fix: Server shutdown grace period

configs/service.yaml

@@ -3,6 +3,7 @@ open_saves_cloud: "gcp"
 open_saves_bucket: ""
 open_saves_project: ""
 log_level: "info"
+shutdown_grace_period: "5s"
 cache_default_ttl: "5m"
 
 redis_address: "localhost:6379"
@@ -13,4 +14,4 @@ redis_max_retries: 3
 redis_min_retry_backoff: "8ms"
 redis_max_retry_backoff: "512ms"
 
-blob_max_inline_size: 65536
+blob_max_inline_size: 65536

internal/app/server/open_saves_test.go

@@ -52,10 +52,11 @@ import (
 )
 
 const (
-	testProject    = "triton-for-games-dev"
-	testBucket     = "gs://triton-integration"
-	testPort       = "8000"
-	testBufferSize = 1024 * 1024
+	testProject             = "triton-for-games-dev"
+	testBucket              = "gs://triton-integration"
+	testPort                = "8000"
+	testBufferSize          = 1024 * 1024
+	testShutdownGracePeriod = "1s"
 	// The threshold of comparing times.
 	// Since the server will actually access the backend datastore,
 	// we need enough time to prevent flaky tests.
@@ -75,6 +76,7 @@ func getServiceConfig() (*config.ServiceConfig, error) {
 	viper.Set(config.OpenSavesBucket, testBucket)
 	viper.Set(config.OpenSavesProject, testProject)
 	viper.Set(config.OpenSavesPort, testPort)
+	viper.Set(config.ShutdownGracePeriod, testShutdownGracePeriod)
 	sc, err := config.Load(configPath)
 	serviceConfig = sc
 	return sc, err
@@ -124,6 +126,8 @@ func TestOpenSaves_HealthCheck(t *testing.T) {
 	if want := healthgrpc.HealthCheckResponse_NOT_SERVING; got.Status != want {
 		t.Errorf("hc.Check got: %v, want: %v", got.Status, want)
 	}
+	// Sleep long enough to free up the connection
+	time.Sleep(serviceConfig.ShutdownGracePeriod)
 }
 
 func TestOpenSaves_RunServer(t *testing.T) {
@@ -148,6 +152,8 @@ func TestOpenSaves_RunServer(t *testing.T) {
 			t.Errorf("got err calling server.Run: %v", err)
 		}
 	})
+	// Sleep long enough to free up the connection
+	time.Sleep(serviceConfig.ShutdownGracePeriod)
 }
 
 func getOpenSavesServer(ctx context.Context, t *testing.T, cloud string) (*openSavesServer, *bufconn.Listener) {

internal/app/server/server.go

@@ -20,6 +20,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"github.com/googleforgames/open-saves/internal/pkg/config"
 
@@ -65,8 +66,13 @@ func Run(ctx context.Context, network string, cfg *config.ServiceConfig) error {
 		case <-ctx.Done():
 			log.Infoln("context cancelled")
 		}
-		log.Infoln("stopping open saves server gracefully")
+		log.Infoln("set health check to not serving")
 		healthcheck.SetServingStatus(serviceName, healthgrpc.HealthCheckResponse_NOT_SERVING)
+
+		log.Infoln("starting server shutdown grace period")
+		time.Sleep(cfg.ShutdownGracePeriod)
+
+		log.Infoln("stopping open saves server gracefully")
 		s.GracefulStop()
 		log.Infoln("stopped open saves server gracefully")
 	}()

internal/pkg/config/loader.go

@@ -92,10 +92,11 @@ func Load(path string) (*ServiceConfig, error) {
 	}
 
 	serverConfig := ServerConfig{
-		Address: fmt.Sprintf(":%d", viper.GetUint(OpenSavesPort)),
-		Cloud:   viper.GetString(OpenSavesCloud),
-		Bucket:  viper.GetString(OpenSavesBucket),
-		Project: viper.GetString(OpenSavesProject),
+		Address:             fmt.Sprintf(":%d", viper.GetUint(OpenSavesPort)),
+		Cloud:               viper.GetString(OpenSavesCloud),
+		Bucket:              viper.GetString(OpenSavesBucket),
+		Project:             viper.GetString(OpenSavesProject),
+		ShutdownGracePeriod: viper.GetDuration(ShutdownGracePeriod),
 	}
 
 	// Cloud Run environment populates the PORT env var, so check for it here.

internal/pkg/config/model.go

@@ -17,11 +17,12 @@ package config
 import "time"
 
 const (
-	OpenSavesPort    = "open_saves_port"
-	OpenSavesCloud   = "open_saves_cloud"
-	OpenSavesBucket  = "open_saves_bucket"
-	OpenSavesProject = "open_saves_project"
-	LogLevel         = "log_level"
+	OpenSavesPort       = "open_saves_port"
+	OpenSavesCloud      = "open_saves_cloud"
+	OpenSavesBucket     = "open_saves_bucket"
+	OpenSavesProject    = "open_saves_project"
+	LogLevel            = "log_level"
+	ShutdownGracePeriod = "shutdown_grace_period"
 
 	CacheDefaultTTL = "cache_default_ttl"
 
@@ -45,10 +46,11 @@ type ServiceConfig struct {
 
 // ServerConfig defines common fields needed to start Open Saves.
 type ServerConfig struct {
-	Address string
-	Cloud   string
-	Bucket  string
-	Project string
+	Address             string
+	Cloud               string
+	Bucket              string
+	Project             string
+	ShutdownGracePeriod time.Duration
 }
 
 // CacheConfig has configurations for caching control (not Redis specific).