feat: add support for log rotation configs

KusionStack · Jan 13, 2025 · 52338ca · 52338ca
1 parent 5c4dd6c
commit 52338ca
Show file tree

Hide file tree

Showing 2 changed files with 115 additions and 7 deletions.
diff --git a/pkg/log/hclog.go b/pkg/log/hclog.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"path/filepath"
+	"strconv"
 
 	"github.com/hashicorp/go-hclog"
 	"google.golang.org/grpc/metadata"
@@ -13,8 +14,11 @@ import (
 )
 
 const (
-	KusionModuleName = "kusion_module_name"
-	KusionTraceID    = "kusion_trace_id"
+	KusionModuleName          = "kusion_module_name"
+	KusionTraceID             = "kusion_trace_id"
+	KusionModuleLogMaxSize    = "kusion_module_log_max_size"
+	KusionModuleLogMaxBackups = "kusion_module_log_max_backups"
+	KusionModuleLogMaxAge     = "kusion_module_log_max_age"
 )
 
 func GetModuleLogger(ctx context.Context) hclog.Logger {
@@ -26,12 +30,15 @@ func GetModuleLogger(ctx context.Context) hclog.Logger {
 
 	kusionDataDir, _ := kfile.KusionDataFolder()
 	logFile := filepath.Join(kusionDataDir, Folder, "modules", moduleName, fmt.Sprintf("%s.log", moduleName))
+	maxSize, maxBackups, maxAge := getRotationConfigs(ctx)
+
 	lumberjackLogger := &lumberjack.Logger{
-		Filename:  logFile,
-		MaxSize:   10,
-		Compress:  false,
-		LocalTime: true,
-		MaxAge:    28,
+		Filename:   logFile,
+		MaxSize:    maxSize,
+		MaxBackups: maxBackups,
+		MaxAge:     maxAge,
+		Compress:   false,
+		LocalTime:  true,
 	}
 
 	return hclog.New(&hclog.LoggerOptions{
@@ -56,3 +63,42 @@ func getModuleName(ctx context.Context) string {
 	}
 	return names[0]
 }
+
+func getRotationConfigs(ctx context.Context) (logMaxSize, logMaxBackups, logMaxAge int) {
+	var err error
+
+	logMaxSizeStrs := metadata.ValueFromIncomingContext(ctx, KusionModuleLogMaxSize)
+	if len(logMaxSizeStrs) == 0 {
+		// Set the default max size of kusion module logger to 10MB.
+		logMaxSize = 10
+	} else {
+		logMaxSize, err = strconv.Atoi(logMaxSizeStrs[0])
+		if err != nil {
+			logMaxSize = 10
+		}
+	}
+
+	logMaxBackupStrs := metadata.ValueFromIncomingContext(ctx, KusionModuleLogMaxBackups)
+	if len(logMaxBackupStrs) == 0 {
+		// Set the default max backups of kusion module logger to 10.
+		logMaxBackups = 10
+	} else {
+		logMaxBackups, err = strconv.Atoi(logMaxBackupStrs[0])
+		if err != nil {
+			logMaxBackups = 10
+		}
+	}
+
+	logMaxAgeStrs := metadata.ValueFromIncomingContext(ctx, KusionModuleLogMaxAge)
+	if len(logMaxAgeStrs) == 0 {
+		// Set the default max age of kusion module logger to 28 days.
+		logMaxAge = 28
+	} else {
+		logMaxAge, err = strconv.Atoi(logMaxAgeStrs[0])
+		if err != nil {
+			logMaxAge = 28
+		}
+	}
+
+	return
+}
diff --git a/pkg/log/hclog_test.go b/pkg/log/hclog_test.go
@@ -76,3 +76,65 @@ func TestGetModuleName(t *testing.T) {
 		})
 	}
 }
+
+func TestGetRotationConfigs(t *testing.T) {
+	tests := []struct {
+		name             string
+		input            map[string]string
+		wantedMaxSize    int
+		wantedMaxBackups int
+		wantedMaxAge     int
+	}{
+		{
+			name:             "empty input",
+			input:            nil,
+			wantedMaxSize:    10,
+			wantedMaxBackups: 10,
+			wantedMaxAge:     28,
+		},
+		{
+			name: "empty input values",
+			input: map[string]string{
+				KusionModuleLogMaxSize:    "",
+				KusionModuleLogMaxBackups: "",
+				KusionModuleLogMaxAge:     "",
+			},
+			wantedMaxSize:    10,
+			wantedMaxBackups: 10,
+			wantedMaxAge:     28,
+		},
+		{
+			name: "invalid input values",
+			input: map[string]string{
+				KusionModuleLogMaxSize:    "?",
+				KusionModuleLogMaxBackups: "?",
+				KusionModuleLogMaxAge:     "?",
+			},
+			wantedMaxSize:    10,
+			wantedMaxBackups: 10,
+			wantedMaxAge:     28,
+		},
+		{
+			name: "customized valid input values",
+			input: map[string]string{
+				KusionModuleLogMaxSize:    "100",
+				KusionModuleLogMaxBackups: "100",
+				KusionModuleLogMaxAge:     "100",
+			},
+			wantedMaxSize:    100,
+			wantedMaxBackups: 100,
+			wantedMaxAge:     100,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := metadata.NewIncomingContext(context.Background(), metadata.New(tt.input))
+			if gotMaxSize, gotMaxBackups, gotMaxAge := getRotationConfigs(ctx); gotMaxSize != tt.wantedMaxSize ||
+				gotMaxBackups != tt.wantedMaxBackups || gotMaxAge != tt.wantedMaxAge {
+				t.Errorf("getRotationConfigs() = %d, %d, %d want %d, %d, %d", gotMaxSize, gotMaxBackups, gotMaxAge,
+					tt.wantedMaxSize, tt.wantedMaxBackups, tt.wantedMaxAge)
+			}
+		})
+	}
+}