Add profiling to iac cli

pkg/engine2/cli.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"path/filepath"
 	"runtime/pprof"
 	"strings"
 
@@ -204,6 +205,10 @@ func (em *EngineMain) ListAttributes(cmd *cobra.Command, args []string) error {
 
 func (em *EngineMain) RunEngine(cmd *cobra.Command, args []string) error {
 	if engineCfg.profileTo != "" {
+		err := os.MkdirAll(filepath.Dir(engineCfg.profileTo), 0755)
+		if err != nil {
+			return fmt.Errorf("failed to create profile directory: %w", err)
+		}
 		profileF, err := os.OpenFile(engineCfg.profileTo, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
 		if err != nil {
 			return fmt.Errorf("failed to open profile file: %w", err)

pkg/engine2/path_selection/path_expansion.go

@@ -11,6 +11,7 @@ import (
 	"github.com/klothoplatform/klotho/pkg/engine2/solution_context"
 	knowledgebase "github.com/klothoplatform/klotho/pkg/knowledge_base2"
 	"github.com/klothoplatform/klotho/pkg/set"
+	"go.uber.org/zap"
 )
 
 type ExpansionInput struct {
@@ -247,12 +248,13 @@ func handleProperties(
 func ExpandPath(
 	ctx solution_context.SolutionContext,
 	input ExpansionInput,
-	path []construct.ResourceId,
+	path construct.Path,
 	resultGraph construct.Graph,
 ) error {
 	if len(path) == 2 {
 		return nil
 	}
+	zap.S().Debugf("Expanding path %s", path)
 
 	type candidate struct {
 		id             construct.ResourceId

pkg/engine2/path_selection/path_selection.go

@@ -21,7 +21,10 @@ const FUNCTIONAL_WEIGHT = 100000
 func BuildPathSelectionGraph(
 	dep construct.SimpleEdge,
 	kb knowledgebase.TemplateKB,
-	classification string) (construct.Graph, error) {
+	classification string,
+) (construct.Graph, error) {
+	zap.S().Debugf("Building path selection graph for %s", dep)
+
 	tempGraph := construct.NewAcyclicGraph(graph.Weighted())
 
 	// Check to see if there is a direct edge which satisfies the classification and if so short circuit in building the temp graph

pkg/infra/cli.go

@@ -21,6 +21,9 @@ var generateIacCfg struct {
 	inputGraph string
 	outputDir  string
 	appName    string
+	verbose    bool
+	jsonLog    bool
+	profileTo  string
 }
 
 func (i *IacCli) AddIacCli(root *cobra.Command) error {

pkg/infra/cli2.go

@@ -3,16 +3,22 @@ package infra
 import (
 	"fmt"
 	"os"
+	"path/filepath"
+	"runtime/pprof"
 
+	"github.com/klothoplatform/klotho/pkg/closenicely"
 	"github.com/klothoplatform/klotho/pkg/config"
 	construct "github.com/klothoplatform/klotho/pkg/construct2"
 	engine "github.com/klothoplatform/klotho/pkg/engine2"
 	"github.com/klothoplatform/klotho/pkg/infra/iac3"
 	"github.com/klothoplatform/klotho/pkg/infra/kubernetes"
 	"github.com/klothoplatform/klotho/pkg/io"
 	knowledgebase "github.com/klothoplatform/klotho/pkg/knowledge_base2"
+	"github.com/klothoplatform/klotho/pkg/logging"
 	"github.com/klothoplatform/klotho/pkg/templates"
 	"github.com/spf13/cobra"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 	"gopkg.in/yaml.v3"
 )
 
@@ -27,11 +33,61 @@ func AddIacCli(root *cobra.Command) error {
 	flags.StringVarP(&generateIacCfg.inputGraph, "input-graph", "i", "", "Input graph to use")
 	flags.StringVarP(&generateIacCfg.outputDir, "output-dir", "o", "", "Output directory to use")
 	flags.StringVarP(&generateIacCfg.appName, "app-name", "a", "", "App name to use")
+	flags.BoolVarP(&generateIacCfg.verbose, "verbose", "v", false, "Verbose flag")
+	flags.BoolVar(&generateIacCfg.jsonLog, "json-log", false, "Output logs in JSON format.")
+	flags.StringVar(&generateIacCfg.profileTo, "profiling", "", "Profile to file")
 	root.AddCommand(generateCmd)
 	return nil
 }
 
+const consoleEncoderName = "iac-cli"
+
+func setupLogger() (*zap.Logger, error) {
+	var zapCfg zap.Config
+	if generateIacCfg.verbose {
+		zapCfg = zap.NewDevelopmentConfig()
+	} else {
+		zapCfg = zap.NewProductionConfig()
+	}
+	if generateIacCfg.jsonLog {
+		zapCfg.Encoding = "json"
+	} else {
+		err := zap.RegisterEncoder(consoleEncoderName, func(zcfg zapcore.EncoderConfig) (zapcore.Encoder, error) {
+			return logging.NewConsoleEncoder(generateIacCfg.verbose, nil, nil), nil
+		})
+		if err != nil {
+			return nil, err
+		}
+		zapCfg.Encoding = consoleEncoderName
+	}
+
+	return zapCfg.Build()
+}
+
 func GenerateIac(cmd *cobra.Command, args []string) error {
+	z, err := setupLogger()
+	if err != nil {
+		return err
+	}
+	defer closenicely.FuncOrDebug(z.Sync)
+	zap.ReplaceGlobals(z)
+
+	if generateIacCfg.profileTo != "" {
+		err := os.MkdirAll(filepath.Dir(generateIacCfg.profileTo), 0755)
+		if err != nil {
+			return fmt.Errorf("failed to create profile directory: %w", err)
+		}
+		profileF, err := os.OpenFile(generateIacCfg.profileTo, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
+		if err != nil {
+			return fmt.Errorf("failed to open profile file: %w", err)
+		}
+		defer func() {
+			pprof.StopCPUProfile()
+			profileF.Close()
+		}()
+		pprof.StartCPUProfile(profileF)
+	}
+
 	var files []io.File
 	if generateIacCfg.inputGraph == "" {
 		return fmt.Errorf("input graph required")

pkg/logging/console.go

@@ -109,10 +109,10 @@ func (enc *ConsoleEncoder) levelPadding() string {
 func (enc *ConsoleEncoder) EncodeEntry(ent zapcore.Entry, fieldList []zapcore.Field) (*buffer.Buffer, error) {
 	line := pool.Get()
 
-	if ent.Level >= zapcore.WarnLevel {
+	if ent.Level >= zapcore.WarnLevel && enc.HadWarnings != nil {
 		enc.HadWarnings.Store(true)
 	}
-	if ent.Level >= zapcore.ErrorLevel {
+	if ent.Level >= zapcore.ErrorLevel && enc.HadErrors != nil {
 		enc.HadErrors.Store(true)
 	}