fix: fix admin and public servers mixed issue

instill-ai · Feb 20, 2023 · 8f8adef · 8f8adef
1 parent c721e97
commit 8f8adef
Show file tree

Hide file tree

Showing 4 changed files with 45 additions and 25 deletions.
diff --git a/cmd/main/main.go b/cmd/main/main.go
@@ -116,11 +116,14 @@ func main() {
 
 	repository := repository.NewRepository(db)
 
-	grpcS := grpc.NewServer(grpcServerOpts...)
-	reflection.Register(grpcS)
+	adminGrpcS := grpc.NewServer(grpcServerOpts...)
+	reflection.Register(adminGrpcS)
+
+	publicGrpcS := grpc.NewServer(grpcServerOpts...)
+	reflection.Register(publicGrpcS)
 
 	mgmtPB.RegisterMgmtAdminServiceServer(
-		grpcS,
+		adminGrpcS,
 		handler.NewAdminHandler(service.NewService(repository)))
 
 	// Usage collection
@@ -137,10 +140,25 @@ func main() {
 	}
 
 	mgmtPB.RegisterMgmtPublicServiceServer(
-		grpcS,
+		publicGrpcS,
 		handler.NewPublicHandler(service.NewService(repository), usg))
 
-	gwS := runtime.NewServeMux(
+	adminServeMux := runtime.NewServeMux(
+		runtime.WithForwardResponseOption(handler.HttpResponseModifier),
+		runtime.WithErrorHandler(handler.ErrorHandler),
+		runtime.WithMarshalerOption(runtime.MIMEWildcard, &runtime.JSONPb{
+			MarshalOptions: protojson.MarshalOptions{
+				UseProtoNames:   true,
+				EmitUnpopulated: true,
+				UseEnumNumbers:  false,
+			},
+			UnmarshalOptions: protojson.UnmarshalOptions{
+				DiscardUnknown: true,
+			},
+		}),
+	)
+
+	publicServeMux := runtime.NewServeMux(
 		runtime.WithForwardResponseOption(handler.HttpResponseModifier),
 		runtime.WithErrorHandler(handler.ErrorHandler),
 		runtime.WithMarshalerOption(runtime.MIMEWildcard, &runtime.JSONPb{
@@ -163,51 +181,51 @@ func main() {
 		dialOpts = []grpc.DialOption{grpc.WithTransportCredentials(insecure.NewCredentials())}
 	}
 
-	if err := mgmtPB.RegisterMgmtAdminServiceHandlerFromEndpoint(ctx, gwS, fmt.Sprintf(":%v", config.Config.Server.AdminPort), dialOpts); err != nil {
+	if err := mgmtPB.RegisterMgmtAdminServiceHandlerFromEndpoint(ctx, adminServeMux, fmt.Sprintf(":%v", config.Config.Server.AdminPort), dialOpts); err != nil {
 		logger.Fatal(err.Error())
 	}
 
-	if err := mgmtPB.RegisterMgmtPublicServiceHandlerFromEndpoint(ctx, gwS, fmt.Sprintf(":%v", config.Config.Server.PublicPort), dialOpts); err != nil {
+	if err := mgmtPB.RegisterMgmtPublicServiceHandlerFromEndpoint(ctx, publicServeMux, fmt.Sprintf(":%v", config.Config.Server.PublicPort), dialOpts); err != nil {
 		logger.Fatal(err.Error())
 	}
 
-	adminHttpServer := &http.Server{
+	adminHTTPServer := &http.Server{
 		Addr:    fmt.Sprintf(":%v", config.Config.Server.AdminPort),
-		Handler: grpcHandlerFunc(grpcS, gwS, config.Config.Server.CORSOrigins),
+		Handler: grpcHandlerFunc(adminGrpcS, adminServeMux, config.Config.Server.CORSOrigins),
 	}
 
-	publicHttpServer := &http.Server{
+	publicHTTPServer := &http.Server{
 		Addr:    fmt.Sprintf(":%v", config.Config.Server.PublicPort),
-		Handler: grpcHandlerFunc(grpcS, gwS, config.Config.Server.CORSOrigins),
+		Handler: grpcHandlerFunc(publicGrpcS, publicServeMux, config.Config.Server.CORSOrigins),
 	}
 
 	// Wait for interrupt signal to gracefully shutdown the server with a timeout of 5 seconds.
 	quitSig := make(chan os.Signal, 1)
 	errSig := make(chan error)
 	if config.Config.Server.HTTPS.Cert != "" && config.Config.Server.HTTPS.Key != "" {
 		go func() {
-			if err := adminHttpServer.ListenAndServeTLS(config.Config.Server.HTTPS.Cert, config.Config.Server.HTTPS.Key); err != nil {
+			if err := adminHTTPServer.ListenAndServeTLS(config.Config.Server.HTTPS.Cert, config.Config.Server.HTTPS.Key); err != nil {
 				errSig <- err
 			}
 		}()
 		go func() {
-			if err := publicHttpServer.ListenAndServeTLS(config.Config.Server.HTTPS.Cert, config.Config.Server.HTTPS.Key); err != nil {
+			if err := publicHTTPServer.ListenAndServeTLS(config.Config.Server.HTTPS.Cert, config.Config.Server.HTTPS.Key); err != nil {
 				errSig <- err
 			}
 		}()
 	} else {
 		go func() {
-			if err := adminHttpServer.ListenAndServe(); err != nil {
+			if err := adminHTTPServer.ListenAndServe(); err != nil {
 				errSig <- err
 			}
 		}()
 		go func() {
-			if err := publicHttpServer.ListenAndServe(); err != nil {
+			if err := publicHTTPServer.ListenAndServe(); err != nil {
 				errSig <- err
 			}
 		}()
 	}
-	logger.Info("gRPC server is running.")
+	logger.Info("gRPC servers are running.")
 
 	// kill (no param) default send syscall.SIGTERM
 	// kill -2 is syscall.SIGINT
@@ -219,6 +237,7 @@ func main() {
 		logger.Error(fmt.Sprintf("Fatal error: %v\n", err))
 	case <-quitSig:
 		logger.Info("Shutting down server...")
-		grpcS.GracefulStop()
+		adminGrpcS.GracefulStop()
+		publicGrpcS.GracefulStop()
 	}
 }
diff --git a/pkg/handler/adminhandler.go b/pkg/handler/adminhandler.go
@@ -11,14 +11,13 @@ import (
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
+	fieldmask_utils "github.com/mennanov/fieldmask-utils"
+
 	"github.com/instill-ai/mgmt-backend/pkg/datamodel"
 	"github.com/instill-ai/mgmt-backend/pkg/logger"
 	"github.com/instill-ai/mgmt-backend/pkg/service"
-
 	"github.com/instill-ai/x/sterr"
 
-	fieldmask_utils "github.com/mennanov/fieldmask-utils"
-
 	mgmtPB "github.com/instill-ai/protogen-go/vdp/mgmt/v1alpha"
 	checkfield "github.com/instill-ai/x/checkfield"
 )

diff --git a/pkg/handler/interceptor.go b/pkg/handler/interceptor.go
@@ -7,7 +7,7 @@ import (
 	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
 )
 
-// RecoveryInterceptor - panic handler
+// RecoveryInterceptorOpt - panic handler
 func RecoveryInterceptorOpt() grpc_recovery.Option {
 	return grpc_recovery.WithRecoveryHandler(func(p interface{}) (err error) {
 		return status.Errorf(codes.Unknown, "panic triggered: %v", p)

diff --git a/pkg/handler/publichandler.go b/pkg/handler/publichandler.go
@@ -7,20 +7,22 @@ import (
 
 	"github.com/gofrs/uuid"
 	"github.com/iancoleman/strcase"
+	"google.golang.org/genproto/googleapis/rpc/errdetails"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	fieldmask_utils "github.com/mennanov/fieldmask-utils"
+
 	"github.com/instill-ai/mgmt-backend/config"
 	"github.com/instill-ai/mgmt-backend/pkg/datamodel"
 	"github.com/instill-ai/mgmt-backend/pkg/logger"
 	"github.com/instill-ai/mgmt-backend/pkg/service"
 	"github.com/instill-ai/mgmt-backend/pkg/usage"
 	"github.com/instill-ai/x/sterr"
-	"google.golang.org/genproto/googleapis/rpc/errdetails"
-	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/status"
 
 	healthcheckPB "github.com/instill-ai/protogen-go/vdp/healthcheck/v1alpha"
 	mgmtPB "github.com/instill-ai/protogen-go/vdp/mgmt/v1alpha"
 	checkfield "github.com/instill-ai/x/checkfield"
-	fieldmask_utils "github.com/mennanov/fieldmask-utils"
 )
 
 // TODO: Validate mask based on the field behavior. Currently, the fields are hard-coded.