refactor: complete graceful shutdown implementation and add service runner (#493)

This commit is contained in:
Alessandro (Ale) Segala
2025-05-03 14:25:22 -07:00
committed by GitHub
parent 6bd6cefaa6
commit bda178c2bb
13 changed files with 454 additions and 130 deletions

View File

@@ -9,27 +9,28 @@ import (
"time"
"github.com/gin-gonic/gin"
"github.com/pocket-id/pocket-id/backend/internal/common"
"github.com/pocket-id/pocket-id/backend/internal/controller"
"github.com/pocket-id/pocket-id/backend/internal/job"
"github.com/pocket-id/pocket-id/backend/internal/middleware"
"github.com/pocket-id/pocket-id/backend/internal/service"
"github.com/pocket-id/pocket-id/backend/internal/utils/systemd"
"golang.org/x/time/rate"
"gorm.io/gorm"
"github.com/pocket-id/pocket-id/backend/internal/common"
"github.com/pocket-id/pocket-id/backend/internal/controller"
"github.com/pocket-id/pocket-id/backend/internal/middleware"
"github.com/pocket-id/pocket-id/backend/internal/utils"
"github.com/pocket-id/pocket-id/backend/internal/utils/systemd"
)
// This is used to register additional controllers for tests
var registerTestControllers []func(apiGroup *gin.RouterGroup, db *gorm.DB, appConfigService *service.AppConfigService, jwtService *service.JwtService)
var registerTestControllers []func(apiGroup *gin.RouterGroup, db *gorm.DB, svc *services)
func initRouter(ctx context.Context, db *gorm.DB, appConfigService *service.AppConfigService) {
err := initRouterInternal(ctx, db, appConfigService)
func initRouter(db *gorm.DB, svc *services) utils.Service {
runner, err := initRouterInternal(db, svc)
if err != nil {
log.Fatalf("failed to init router: %v", err)
}
return runner
}
func initRouterInternal(ctx context.Context, db *gorm.DB, appConfigService *service.AppConfigService) error {
func initRouterInternal(db *gorm.DB, svc *services) (utils.Service, error) {
// Set the appropriate Gin mode based on the environment
switch common.EnvConfig.AppEnv {
case "production":
@@ -43,23 +44,6 @@ func initRouterInternal(ctx context.Context, db *gorm.DB, appConfigService *serv
r := gin.Default()
r.Use(gin.Logger())
// Initialize services
emailService, err := service.NewEmailService(appConfigService, db)
if err != nil {
return fmt.Errorf("unable to create email service: %w", err)
}
geoLiteService := service.NewGeoLiteService(ctx)
auditLogService := service.NewAuditLogService(db, appConfigService, emailService, geoLiteService)
jwtService := service.NewJwtService(appConfigService)
webauthnService := service.NewWebAuthnService(db, jwtService, auditLogService, appConfigService)
userService := service.NewUserService(db, jwtService, auditLogService, emailService, appConfigService)
customClaimService := service.NewCustomClaimService(db)
oidcService := service.NewOidcService(db, jwtService, appConfigService, auditLogService, customClaimService)
userGroupService := service.NewUserGroupService(db, appConfigService)
ldapService := service.NewLdapService(db, appConfigService, userService, userGroupService)
apiKeyService := service.NewApiKeyService(db, emailService)
rateLimitMiddleware := middleware.NewRateLimitMiddleware()
// Setup global middleware
@@ -67,56 +51,31 @@ func initRouterInternal(ctx context.Context, db *gorm.DB, appConfigService *serv
r.Use(middleware.NewErrorHandlerMiddleware().Add())
r.Use(rateLimitMiddleware.Add(rate.Every(time.Second), 60))
scheduler, err := job.NewScheduler()
if err != nil {
return fmt.Errorf("failed to create job scheduler: %w", err)
}
err = scheduler.RegisterLdapJobs(ctx, ldapService, appConfigService)
if err != nil {
return fmt.Errorf("failed to register LDAP jobs in scheduler: %w", err)
}
err = scheduler.RegisterDbCleanupJobs(ctx, db)
if err != nil {
return fmt.Errorf("failed to register DB cleanup jobs in scheduler: %w", err)
}
err = scheduler.RegisterFileCleanupJobs(ctx, db)
if err != nil {
return fmt.Errorf("failed to register file cleanup jobs in scheduler: %w", err)
}
err = scheduler.RegisterApiKeyExpiryJob(ctx, apiKeyService, appConfigService)
if err != nil {
return fmt.Errorf("failed to register API key expiration jobs in scheduler: %w", err)
}
// Run the scheduler in a background goroutine, until the context is canceled
go scheduler.Run(ctx)
// Initialize middleware for specific routes
authMiddleware := middleware.NewAuthMiddleware(apiKeyService, userService, jwtService)
authMiddleware := middleware.NewAuthMiddleware(svc.apiKeyService, svc.userService, svc.jwtService)
fileSizeLimitMiddleware := middleware.NewFileSizeLimitMiddleware()
// Set up API routes
apiGroup := r.Group("/api")
controller.NewApiKeyController(apiGroup, authMiddleware, apiKeyService)
controller.NewWebauthnController(apiGroup, authMiddleware, middleware.NewRateLimitMiddleware(), webauthnService, appConfigService)
controller.NewOidcController(apiGroup, authMiddleware, fileSizeLimitMiddleware, oidcService, jwtService)
controller.NewUserController(apiGroup, authMiddleware, middleware.NewRateLimitMiddleware(), userService, appConfigService)
controller.NewAppConfigController(apiGroup, authMiddleware, appConfigService, emailService, ldapService)
controller.NewAuditLogController(apiGroup, auditLogService, authMiddleware)
controller.NewUserGroupController(apiGroup, authMiddleware, userGroupService)
controller.NewCustomClaimController(apiGroup, authMiddleware, customClaimService)
controller.NewApiKeyController(apiGroup, authMiddleware, svc.apiKeyService)
controller.NewWebauthnController(apiGroup, authMiddleware, middleware.NewRateLimitMiddleware(), svc.webauthnService, svc.appConfigService)
controller.NewOidcController(apiGroup, authMiddleware, fileSizeLimitMiddleware, svc.oidcService, svc.jwtService)
controller.NewUserController(apiGroup, authMiddleware, middleware.NewRateLimitMiddleware(), svc.userService, svc.appConfigService)
controller.NewAppConfigController(apiGroup, authMiddleware, svc.appConfigService, svc.emailService, svc.ldapService)
controller.NewAuditLogController(apiGroup, svc.auditLogService, authMiddleware)
controller.NewUserGroupController(apiGroup, authMiddleware, svc.userGroupService)
controller.NewCustomClaimController(apiGroup, authMiddleware, svc.customClaimService)
// Add test controller in non-production environments
if common.EnvConfig.AppEnv != "production" {
for _, f := range registerTestControllers {
f(apiGroup, db, appConfigService, jwtService)
f(apiGroup, db, svc)
}
}
// Set up base routes
baseGroup := r.Group("/")
controller.NewWellKnownController(baseGroup, jwtService)
controller.NewWellKnownController(baseGroup, svc.jwtService)
// Set up the server
srv := &http.Server{
@@ -129,41 +88,46 @@ func initRouterInternal(ctx context.Context, db *gorm.DB, appConfigService *serv
// Set up the listener
listener, err := net.Listen("tcp", srv.Addr)
if err != nil {
return fmt.Errorf("failed to create TCP listener: %w", err)
return nil, fmt.Errorf("failed to create TCP listener: %w", err)
}
log.Printf("Server listening on %s", srv.Addr)
// Service runner function
runFn := func(ctx context.Context) error {
log.Printf("Server listening on %s", srv.Addr)
// Notify systemd that we are ready
err = systemd.SdNotifyReady()
if err != nil {
log.Printf("[WARN] Unable to notify systemd that the service is ready: %v", err)
// continue to serve anyway since it's not that important
}
// Start the server in a background goroutine
go func() {
defer listener.Close()
// Start the server in a background goroutine
go func() {
defer listener.Close()
// Next call blocks until the server is shut down
srvErr := srv.Serve(listener)
if srvErr != http.ErrServerClosed {
log.Fatalf("Error starting app server: %v", srvErr)
}
}()
// Next call blocks until the server is shut down
srvErr := srv.Serve(listener)
if srvErr != http.ErrServerClosed {
log.Fatalf("Error starting app server: %v", srvErr)
// Notify systemd that we are ready
err = systemd.SdNotifyReady()
if err != nil {
// Log the error only
log.Printf("[WARN] Unable to notify systemd that the service is ready: %v", err)
}
}()
// Block until the context is canceled
<-ctx.Done()
// Block until the context is canceled
<-ctx.Done()
// Handle graceful shutdown
// Note we use the background context here as ctx has been canceled already
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
shutdownErr := srv.Shutdown(shutdownCtx) //nolint:contextcheck
shutdownCancel()
if shutdownErr != nil {
// Log the error only (could be context canceled)
log.Printf("[WARN] App server shutdown error: %v", shutdownErr)
// Handle graceful shutdown
// Note we use the background context here as ctx has been canceled already
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
shutdownErr := srv.Shutdown(shutdownCtx) //nolint:contextcheck
shutdownCancel()
if shutdownErr != nil {
// Log the error only (could be context canceled)
log.Printf("[WARN] App server shutdown error: %v", shutdownErr)
}
return nil
}
return nil
return runFn, nil
}