Add Stop() for uber/fx hook

server/grpc/grpc.go

@@ -3,7 +3,7 @@ package grpc
 import (
 	"net"
 
-	"github.com/asim/emque/proto"
+	mq "github.com/asim/emque/proto"
 	"github.com/asim/emque/server"
 	"github.com/asim/emque/server/util"
 	"google.golang.org/grpc"
@@ -12,6 +12,7 @@ import (
 
 type grpcServer struct {
 	options *server.Options
+	srv     *grpc.Server
 }
 
 func (g *grpcServer) Run() error {
@@ -50,6 +51,7 @@ func (g *grpcServer) Run() error {
 
 	// new grpc server
 	srv := grpc.NewServer(opts...)
+	g.srv = srv
 
 	// register MQ server
 	mq.RegisterMQServer(srv, new(handler))
@@ -58,6 +60,12 @@ func (g *grpcServer) Run() error {
 	return srv.Serve(l)
 }
 
+func (g *grpcServer) Stop() error {
+	g.srv.GracefulStop()
+	return nil
+
+}
+
 func New(opts ...server.Option) *grpcServer {
 	options := new(server.Options)
 	for _, o := range opts {

server/http/http.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"crypto/tls"
 	"net/http"
 	"os"
@@ -12,6 +13,7 @@ import (
 
 type httpServer struct {
 	options *server.Options
+	srv     *http.Server
 }
 
 func (h *httpServer) Run() error {
@@ -58,9 +60,15 @@ func (h *httpServer) Run() error {
 		TLSConfig: config,
 	}
 
+	h.srv = srv
+
 	return srv.Serve(l)
 }
 
+func (h *httpServer) Stop() error {
+	return h.srv.Shutdown(context.TODO())
+}
+
 func New(opts ...server.Option) *httpServer {
 	options := new(server.Options)
 	for _, o := range opts {

server/server.go

@@ -2,4 +2,5 @@ package server
 
 type Server interface {
 	Run() error
+	Stop() error
 }