feat: support graceful shutdown on sigint/sigterm

internal/app/server/open_saves_test.go

@@ -22,6 +22,7 @@ import (
 	"reflect"
 	"runtime"
 	"strings"
+	"syscall"
 	"testing"
 	"time"
 
@@ -62,19 +63,31 @@ const (
 	blobKind       = "blob"
 )
 
-func TestOpenSaves_HealthCheck(t *testing.T) {
+var (
+	serviceConfig *config.ServiceConfig
+)
+
+func getServiceConfig() (*config.ServiceConfig, error) {
+	if serviceConfig != nil {
+		return serviceConfig, nil
+	}
 	configPath := cmd.GetEnvVarString("OPEN_SAVES_CONFIG", "../../../configs/")
 	viper.Set(config.OpenSavesBucket, testBucket)
 	viper.Set(config.OpenSavesProject, testProject)
 	viper.Set(config.OpenSavesPort, testPort)
-	cfg, err := config.Load(configPath)
+	sc, err := config.Load(configPath)
+	serviceConfig = sc
+	return sc, err
+}
+
+func TestOpenSaves_HealthCheck(t *testing.T) {
+	serviceConfig, err := getServiceConfig()
 	if err != nil {
-		t.Fatalf("got err loading config: %v", err)
+		t.Fatalf("getServiceConfig err: %v", err)
 	}
-
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
 	go func() {
-		if err := Run(ctx, "tcp", cfg); err != nil {
+		if err := Run(ctx, "tcp", serviceConfig); err != nil {
 			log.Errorf("got err calling server.Run: %v", err)
 		}
 	}()
@@ -99,6 +112,42 @@ func TestOpenSaves_HealthCheck(t *testing.T) {
 	if want := healthgrpc.HealthCheckResponse_SERVING; got.Status != want {
 		t.Errorf("hc.Check got: %v, want: %v", got.Status, want)
 	}
+	// Cancel here to stop the server. This frees up the port for later tests.
+	cancel()
+	ctx2 := context.Background()
+	got, err = hc.Check(ctx2, &healthgrpc.HealthCheckRequest{
+		Service: serviceName,
+	})
+	if err != nil {
+		t.Errorf("healthClient.Check err: %v", err)
+	}
+	if want := healthgrpc.HealthCheckResponse_NOT_SERVING; got.Status != want {
+		t.Errorf("hc.Check got: %v, want: %v", got.Status, want)
+	}
+}
+
+func TestOpenSaves_RunServer(t *testing.T) {
+	serviceConfig, err := getServiceConfig()
+	if err != nil {
+		t.Fatalf("getServiceConfig err: %v", err)
+	}
+	ctx := context.Background()
+	t.Run("cancel_context", func(t *testing.T) {
+		ctx, cancel := context.WithTimeout(ctx, 2*time.Second)
+		defer cancel()
+		if err := Run(ctx, "tcp", serviceConfig); err != nil {
+			t.Errorf("got err calling server.Run: %v", err)
+		}
+	})
+	t.Run("sigint", func(t *testing.T) {
+		go func() {
+			time.Sleep(2 * time.Second)
+			syscall.Kill(syscall.Getpid(), syscall.SIGINT)
+		}()
+		if err := Run(ctx, "tcp", serviceConfig); err != nil {
+			t.Errorf("got err calling server.Run: %v", err)
+		}
+	})
 }
 
 func getOpenSavesServer(ctx context.Context, t *testing.T, cloud string) (*openSavesServer, *bufconn.Listener) {

internal/app/server/server.go

@@ -17,6 +17,9 @@ package server
 import (
 	"context"
 	"net"
+	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/googleforgames/open-saves/internal/pkg/config"
 
@@ -46,6 +49,9 @@ func Run(ctx context.Context, network string, cfg *config.ServiceConfig) error {
 		}
 	}()
 
+	sigs := make(chan os.Signal, 1)
+	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
+
 	s := grpc.NewServer()
 	healthcheck := health.NewServer()
 	healthcheck.SetServingStatus(serviceName, healthgrpc.HealthCheckResponse_SERVING)
@@ -57,5 +63,18 @@ func Run(ctx context.Context, network string, cfg *config.ServiceConfig) error {
 	pb.RegisterOpenSavesServer(s, server)
 	reflection.Register(s)
 
+	go func() {
+		select {
+		case s := <-sigs:
+			log.Infof("received signal: %v\n", s)
+		case <-ctx.Done():
+			log.Infoln("context cancelled")
+		}
+		log.Infoln("stopping open saves server gracefully")
+		healthcheck.SetServingStatus(serviceName, healthgrpc.HealthCheckResponse_NOT_SERVING)
+		s.GracefulStop()
+		log.Infoln("stopped open saves server gracefully")
+	}()
+
 	return s.Serve(lis)
 }