From 696804452e459982e9cf861fdd316642b3e8b973 Mon Sep 17 00:00:00 2001
From: Paul Pan <i@0x7f.app>
Date: Sat, 6 Jan 2024 19:23:12 +0800
Subject: [PATCH] feat: use injector shutdown rather than manually call Close()

---
 cmd/woj/woj.go                |  5 +++++
 internal/app/server/server.go | 16 +---------------
 internal/repo/db/pg.go        |  5 +++++
 3 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/cmd/woj/woj.go b/cmd/woj/woj.go
index c4d7100..96e3877 100644
--- a/cmd/woj/woj.go
+++ b/cmd/woj/woj.go
@@ -105,6 +105,11 @@ func wrap(f func(i *do.Injector) error) func(*cli.Context) error {
 		}()
 
 		injector := prepareServices(c)
+		defer func() {
+			if err := injector.Shutdown(); err != nil {
+				slog.Printf("shutdown injector failed: %v", err)
+			}
+		}()
 
 		logger := do.MustInvoke[log.Service](injector)
 		defer func() { _ = logger.GetRawLogger().Sync() }()
diff --git a/internal/app/server/server.go b/internal/app/server/server.go
index 163823c..c54f182 100644
--- a/internal/app/server/server.go
+++ b/internal/app/server/server.go
@@ -8,7 +8,6 @@ import (
 	"git.0x7f.app/WOJ/woj-server/internal/misc/config"
 	"git.0x7f.app/WOJ/woj-server/internal/misc/log"
 	"git.0x7f.app/WOJ/woj-server/internal/model"
-	"git.0x7f.app/WOJ/woj-server/internal/repo/db"
 	"git.0x7f.app/WOJ/woj-server/internal/web/router"
 	"git.0x7f.app/WOJ/woj-server/pkg/utils"
 	"git.0x7f.app/WOJ/woj-server/pkg/zapasynq"
@@ -24,15 +23,8 @@ import (
 )
 
 func RunServerMigrate(i *do.Injector) error {
-	slog := do.MustInvoke[log.Service](i).GetLogger("app.server")
 
-	// Migrate and shutdown database
-	err := do.MustInvoke[db.Service](i).Close()
-	if err != nil {
-		slog.Warn("Database Close Failed", zap.Error(err))
-	}
-
-	return err
+	return nil
 }
 
 func RunServer(i *do.Injector) error {
@@ -98,11 +90,5 @@ func RunServer(i *do.Injector) error {
 	// Graceful Shutdown Queue
 	queueSrv.Shutdown()
 
-	// Graceful Shutdown Database
-	err = do.MustInvoke[db.Service](i).Close()
-	if err != nil {
-		slog.Warn("Database Close Failed", zap.Error(err))
-	}
-
 	return err
 }
diff --git a/internal/repo/db/pg.go b/internal/repo/db/pg.go
index 3d0b39d..77aba43 100644
--- a/internal/repo/db/pg.go
+++ b/internal/repo/db/pg.go
@@ -24,6 +24,7 @@ type Service interface {
 	Get() *gorm.DB
 	Close() error
 	HealthCheck() error
+	Shutdown() error
 }
 
 func NewService(i *do.Injector) (Service, error) {
@@ -60,6 +61,10 @@ func (s *service) HealthCheck() error {
 	return s.err
 }
 
+func (s *service) Shutdown() error {
+	return s.Close()
+}
+
 func (s *service) setup(conf *model.Config) {
 	s.log.Info("Connecting to database...")