From d28d9649ed0d16e8db8459f09f35df5f21fb72b6 Mon Sep 17 00:00:00 2001
From: yanghao05 <yanghao05@qianxin.com>
Date: Thu, 1 Jun 2023 11:02:10 +0800
Subject: [PATCH] add close

---
 container/container.go            | 13 +++++++++++++
 go.mod                            |  1 -
 go.sum                            |  3 ---
 providers/database/redis/redis.go |  3 +++
 services/grpc.go                  |  2 ++
 services/http.go                  |  2 ++
 6 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/container/container.go b/container/container.go
index 9d407b3..0c1194f 100644
--- a/container/container.go
+++ b/container/container.go
@@ -11,6 +11,7 @@ import (
 
 var Container *dig.Container = dig.New()
 var Cancel context.CancelFunc
+var closeable []func()
 
 func init() {
 	if err := Container.Provide(func() context.Context {
@@ -20,6 +21,18 @@ func init() {
 	}); err != nil {
 		log.Fatal(err)
 	}
+
+	closeable = make([]func(), 0)
+}
+
+func AddCloseAble(c func()) {
+	closeable = append(closeable, c)
+}
+
+func Close() {
+	for _, c := range closeable {
+		c()
+	}
 }
 
 type ProviderContainer struct {
diff --git a/go.mod b/go.mod
index d5ed8a5..43d1e8b 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,6 @@ require (
 	github.com/gin-gonic/gin v1.8.2
 	github.com/glebarez/sqlite v1.5.0
 	github.com/go-gormigrate/gormigrate/v2 v2.0.2
-	github.com/go-redis/redis/v8 v8.11.5
 	github.com/gofrs/uuid v4.0.0+incompatible
 	github.com/golang-jwt/jwt/v4 v4.4.3
 	github.com/imroc/req/v3 v3.34.0
diff --git a/go.sum b/go.sum
index 4cb793a..7c4c0ea 100644
--- a/go.sum
+++ b/go.sum
@@ -159,8 +159,6 @@ github.com/go-playground/universal-translator v0.18.1 h1:Bcnm0ZwsGyWbCzImXv+pAJn
 github.com/go-playground/universal-translator v0.18.1/go.mod h1:xekY+UJKNuX9WP91TpwSH2VMlDf28Uj24BCp08ZFTUY=
 github.com/go-playground/validator/v10 v10.11.2 h1:q3SHpufmypg+erIExEKUmsgmhDTyhcJ38oeKGACXohU=
 github.com/go-playground/validator/v10 v10.11.2/go.mod h1:NieE624vt4SCTJtD87arVLvdmjPAeV8BQlHtMnw9D7s=
-github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI=
-github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
 github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE=
 github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
@@ -424,7 +422,6 @@ github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822/go.mod h1:+n7T8m
 github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
 github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
-github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
 github.com/onsi/ginkgo/v2 v2.9.2 h1:BA2GMJOtfGAfagzYtrAlufIP0lq6QERkFmHLMLPwFSU=
 github.com/onsi/ginkgo/v2 v2.9.2/go.mod h1:WHcJJG2dIlcCqVfBAwUCrJxSPFb6v4azBwgxeMeDuts=
 github.com/onsi/gomega v1.27.4 h1:Z2AnStgsdSayCMDiCU42qIz+HLqEPcgiOCXjAU/w+8E=
diff --git a/providers/database/redis/redis.go b/providers/database/redis/redis.go
index e2c6f4b..c5ea791 100644
--- a/providers/database/redis/redis.go
+++ b/providers/database/redis/redis.go
@@ -24,6 +24,9 @@ func Provide(opts ...opt.Option) error {
 			return nil, errors.Wrap(err, "failed to ping")
 		}
 
+		container.AddCloseAble(func() {
+			client.Close()
+		})
 		return client, nil
 	}, o.DiOptions()...)
 }
diff --git a/services/grpc.go b/services/grpc.go
index 2b18051..20f5ae1 100644
--- a/services/grpc.go
+++ b/services/grpc.go
@@ -14,6 +14,8 @@ type GrpcService struct {
 }
 
 func ServeGrpc() error {
+	defer container.Close()
+
 	return container.Container.Invoke(func(grpc GrpcService) error {
 		for _, svc := range grpc.Services {
 			grpc.Server.RegisterService(svc.Name(), svc.Register)
diff --git a/services/http.go b/services/http.go
index dac2c60..a79247a 100644
--- a/services/http.go
+++ b/services/http.go
@@ -14,6 +14,8 @@ type Http struct {
 }
 
 func ServeHttp() error {
+	defer container.Close()
+
 	return container.Container.Invoke(func(http Http) error {
 		for _, route := range http.Routes {
 			route.Register()