refactoring profile handler into entrypoints (#113)

* moved profutils profiling server start to entrypoint so it is not included in single binary Signed-off-by: Daniel Rammer <daniel@union.ai> * fixed imports Signed-off-by: Daniel Rammer <daniel@union.ai> --------- Signed-off-by: Daniel Rammer <daniel@union.ai>
flyteorg · Aug 23, 2023 · b5a958f · b5a958f
1 parent 87e16cf
commit b5a958f
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 21 deletions.
diff --git a/cmd/entrypoints/serve.go b/cmd/entrypoints/serve.go
@@ -3,12 +3,14 @@ package entrypoints
 import (
 	"context"
 
-	"github.com/flyteorg/flytestdlib/contextutils"
-	"github.com/flyteorg/flytestdlib/promutils/labeled"
-
 	"github.com/flyteorg/datacatalog/pkg/config"
 	"github.com/flyteorg/datacatalog/pkg/rpc/datacatalogservice"
+	"github.com/flyteorg/datacatalog/pkg/runtime"
+	"github.com/flyteorg/flytestdlib/contextutils"
 	"github.com/flyteorg/flytestdlib/logger"
+	"github.com/flyteorg/flytestdlib/profutils"
+	"github.com/flyteorg/flytestdlib/promutils/labeled"
+
 	"github.com/spf13/cobra"
 )
 
@@ -27,6 +29,16 @@ var serveCmd = &cobra.Command{
 			}
 		}()
 
+		// Serve profiling endpoint.
+		dataCatalogConfig := runtime.NewConfigurationProvider().ApplicationConfiguration().GetDataCatalogConfig()
+		go func() {
+			err := profutils.StartProfilingServerWithDefaultHandlers(
+				context.Background(), dataCatalogConfig.ProfilerPort, nil)
+			if err != nil {
+				logger.Panicf(context.Background(), "Failed to Start profiling and Metrics server. Error, %v", err)
+			}
+		}()
+
 		// Set Keys
 		labeled.SetMetricKeys(contextutils.AppNameKey, contextutils.ProjectKey, contextutils.DomainKey)
 

diff --git a/cmd/entrypoints/serve_dummy.go b/cmd/entrypoints/serve_dummy.go
@@ -5,6 +5,7 @@ import (
 
 	"github.com/flyteorg/datacatalog/pkg/config"
 	"github.com/flyteorg/datacatalog/pkg/rpc/datacatalogservice"
+	"github.com/flyteorg/flytestdlib/logger"
 	"github.com/spf13/cobra"
 )
 
@@ -14,6 +15,15 @@ var serveDummyCmd = &cobra.Command{
 	RunE: func(cmd *cobra.Command, args []string) error {
 		ctx := context.Background()
 		cfg := config.GetConfig()
+
+		// serve a http healthcheck endpoint
+		go func() {
+			err := datacatalogservice.ServeHTTPHealthCheck(ctx, cfg)
+			if err != nil {
+				logger.Errorf(ctx, "Unable to serve http", cfg.GetGrpcHostAddress(), err)
+			}
+		}()
+
 		return datacatalogservice.Serve(ctx, cfg)
 	},
 }

diff --git a/pkg/rpc/datacatalogservice/service.go b/pkg/rpc/datacatalogservice/service.go
@@ -21,7 +21,6 @@ import (
 	catalog "github.com/flyteorg/flyteidl/gen/pb-go/flyteidl/datacatalog"
 	"github.com/flyteorg/flytestdlib/contextutils"
 	"github.com/flyteorg/flytestdlib/logger"
-	"github.com/flyteorg/flytestdlib/profutils"
 	"github.com/flyteorg/flytestdlib/promutils"
 	"github.com/flyteorg/flytestdlib/storage"
 )
@@ -106,15 +105,6 @@ func NewDataCatalogService() *DataCatalogService {
 	repos := repositories.GetRepository(ctx, repositories.POSTGRES, *dbConfigValues, catalogScope)
 	logger.Infof(ctx, "Created DB connection.")
 
-	// Serve profiling endpoint.
-	go func() {
-		err := profutils.StartProfilingServerWithDefaultHandlers(
-			context.Background(), dataCatalogConfig.ProfilerPort, nil)
-		if err != nil {
-			logger.Panicf(context.Background(), "Failed to Start profiling and Metrics server. Error, %v", err)
-		}
-	}()
-
 	return &DataCatalogService{
 		DatasetManager:  impl.NewDatasetManager(repos, dataStorageClient, catalogScope.NewSubScope("dataset")),
 		ArtifactManager: impl.NewArtifactManager(repos, dataStorageClient, storagePrefix, catalogScope.NewSubScope("artifact")),
@@ -173,14 +163,6 @@ func ServeHTTPHealthCheck(ctx context.Context, cfg *config.Config) error {
 
 // Create and start the gRPC server and http healthcheck endpoint
 func Serve(ctx context.Context, cfg *config.Config) error {
-	// serve a http healthcheck endpoint
-	go func() {
-		err := ServeHTTPHealthCheck(ctx, cfg)
-		if err != nil {
-			logger.Errorf(ctx, "Unable to serve http", cfg.GetGrpcHostAddress(), err)
-		}
-	}()
-
 	grpcServer := newGRPCDummyServer(ctx, cfg)
 
 	grpcListener, err := net.Listen("tcp", cfg.GetGrpcHostAddress())