Migrate logging to slog and more structured logs

Author: J12934

balancer/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"log"
+	"log/slog"
 	"net/http"
 	"os"
 
@@ -52,7 +53,7 @@ func StartBalancerServer(b *bundle.Bundle) {
 	router := http.NewServeMux()
 	routes.AddRoutes(router, b)
 
-	b.Log.Println("Starting MultiJuicer balancer on :8080")
+	b.Log.Info("Starting MultiJuicer balancer on :8080")
 	server := &http.Server{
 		Addr:    ":8080",
 		Handler: router,
@@ -63,14 +64,14 @@ func StartBalancerServer(b *bundle.Bundle) {
 	}
 }
 
-func StartLLMGatewayServer(gateway *llmgateway.Gateway, logger *log.Logger) {
+func StartLLMGatewayServer(gateway *llmgateway.Gateway, logger *slog.Logger) {
 	router := http.NewServeMux()
 	router.Handle("/", gateway)
 	server := &http.Server{
 		Addr:    ":8082",
 		Handler: router,
 	}
-	logger.Println("Starting LLM gateway on :8082")
+	logger.Info("Starting LLM gateway on :8082")
 	if err := server.ListenAndServe(); err != nil {
 		log.Fatalf("Failed to start LLM gateway server: %v", err)
 	}

balancer/pkg/bundle/bundle.go

@@ -5,8 +5,9 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"log"
+	"log/slog"
 	"os"
+	"strings"
 	"time"
 
 	"github.com/juice-shop/multi-juicer/balancer/pkg/passcode"
@@ -28,7 +29,7 @@ type Bundle struct {
 	BcryptRounds           int
 	StaticAssetsDirectory  string `json:"staticAssetsDirectory"`
 	Config                 *Config
-	Log                    *log.Logger
+	Log                    *slog.Logger
 
 	// LongPollDefaultWaitTimeout amount of time that HTTP Long Polling Endpoints wait for new data to arrive before returning a empty no changes response
 	LongPollDefaultWaitTimeout time.Duration
@@ -184,6 +185,21 @@ type NotificationService interface {
 	SetEndDate(ctx context.Context, endDate *time.Time) error
 }
 
+// ParseLogLevel converts a log level string to a slog.Level.
+// Valid values: "debug", "info", "warn"/"warning", "error". Defaults to info.
+func ParseLogLevel(level string) slog.Level {
+	switch strings.ToLower(level) {
+	case "debug":
+		return slog.LevelDebug
+	case "warn", "warning":
+		return slog.LevelWarn
+	case "error":
+		return slog.LevelError
+	default:
+		return slog.LevelInfo
+	}
+}
+
 func getJuiceShopUrlForTeam(team string, bundle *Bundle) string {
 	return fmt.Sprintf("http://juiceshop-%s.%s.svc.cluster.local:3000", team, bundle.RuntimeEnvironment.Namespace)
 }
@@ -251,7 +267,7 @@ func New() *Bundle {
 		GeneratePasscode:           passcode.GetPasscodeGeneratorWithPasscodeLength(config.TeamPasscodeLength),
 		GetJuiceShopUrlForTeam:     getJuiceShopUrlForTeam,
 		BcryptRounds:               bcrypt.DefaultCost,
-		Log:                        log.New(os.Stdout, "", log.LstdFlags),
+		Log:                        slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: ParseLogLevel(os.Getenv("LOG_LEVEL"))})),
 		LongPollDefaultWaitTimeout: 25 * time.Second,
 		Config:                     config,
 		JuiceShopChallenges:        challenges,

balancer/pkg/llmgateway/gateway.go

@@ -4,7 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"io"
-	"log"
+	"log/slog"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
@@ -29,11 +29,11 @@ type Gateway struct {
 	upstreamURL *url.URL
 	apiKey      string
 	usage       *UsageTracker
-	logger      *log.Logger
+	logger      *slog.Logger
 }
 
 // NewGateway creates a new LLM gateway.
-func NewGateway(signingKey string, upstreamURL string, apiKey string, usage *UsageTracker, logger *log.Logger) (*Gateway, error) {
+func NewGateway(signingKey string, upstreamURL string, apiKey string, usage *UsageTracker, logger *slog.Logger) (*Gateway, error) {
 	u, err := url.Parse(upstreamURL)
 	if err != nil {
 		return nil, err
@@ -65,7 +65,7 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	// Check if this is a chat completions request (for usage tracking)
 	isChatCompletion := strings.Contains(r.URL.Path, "/chat/completions")
-	g.logger.Printf("LLM gateway: request from team '%s': %s %s (isChatCompletion=%v)", team, r.Method, r.URL.Path, isChatCompletion)
+	g.logger.Debug("LLM gateway: request", "team", team, "method", r.Method, "path", r.URL.Path, "isChatCompletion", isChatCompletion)
 
 	// Create reverse proxy
 	proxy := &httputil.ReverseProxy{
@@ -84,7 +84,7 @@ func (g *Gateway) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 
 	proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
-		g.logger.Printf("LLM gateway proxy error for team '%s': %v", team, err)
+		g.logger.Error("LLM gateway proxy error", "team", team, "error", err)
 		http.Error(w, `{"error":"upstream LLM API error"}`, http.StatusBadGateway)
 	}
 
@@ -102,7 +102,7 @@ func (g *Gateway) extractUsage(resp *http.Response, team string) error {
 	body, err := io.ReadAll(resp.Body)
 	resp.Body.Close()
 	if err != nil {
-		g.logger.Printf("LLM gateway: failed to read response body for team '%s': %v", team, err)
+		g.logger.Error("LLM gateway: failed to read response body", "team", team, "error", err)
 		resp.Body = io.NopCloser(bytes.NewReader(body))
 		return nil
 	}
@@ -124,7 +124,7 @@ func (g *Gateway) extractUsageFromJSON(body []byte, team string) {
 		return
 	}
 	if result.Usage != nil {
-		g.logger.Printf("LLM gateway: usage for team '%s': input_tokens=%d, output_tokens=%d", team, result.Usage.InputTokens, result.Usage.OutputTokens)
+		g.logger.Debug("LLM gateway: usage", "team", team, "input_tokens", result.Usage.InputTokens, "output_tokens", result.Usage.OutputTokens)
 		g.usage.Add(team, result.Usage.InputTokens, result.Usage.OutputTokens)
 	}
 }
@@ -148,10 +148,10 @@ func (g *Gateway) extractUsageFromSSE(body []byte, team string) {
 			continue
 		}
 		if chunk.Usage != nil {
-			g.logger.Printf("LLM gateway: SSE usage for team '%s': input_tokens=%d, output_tokens=%d", team, chunk.Usage.InputTokens, chunk.Usage.OutputTokens)
+			g.logger.Debug("LLM gateway: SSE usage", "team", team, "input_tokens", chunk.Usage.InputTokens, "output_tokens", chunk.Usage.OutputTokens)
 			g.usage.Add(team, chunk.Usage.InputTokens, chunk.Usage.OutputTokens)
 			return
 		}
 	}
-	g.logger.Printf("LLM gateway: no usage data found in SSE stream for team '%s'", team)
+	g.logger.Debug("LLM gateway: no usage data found in SSE stream", "team", team)
 }

balancer/pkg/llmgateway/gateway_test.go

@@ -3,7 +3,7 @@ package llmgateway
 import (
 	"encoding/json"
 	"io"
-	"log"
+	"log/slog"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -22,7 +22,7 @@ func signToken(team string) string {
 
 func TestGateway_MissingAuth(t *testing.T) {
 	usage := NewUsageTracker()
-	gw, _ := NewGateway(testSigningKey, "http://localhost:11434", "real-key", usage, log.New(os.Stdout, "", 0))
+	gw, _ := NewGateway(testSigningKey, "http://localhost:11434", "real-key", usage, slog.New(slog.NewTextHandler(os.Stdout, nil)))
 
 	req := httptest.NewRequest("POST", "/v1/chat/completions", nil)
 	w := httptest.NewRecorder()
@@ -35,7 +35,7 @@ func TestGateway_MissingAuth(t *testing.T) {
 
 func TestGateway_InvalidToken(t *testing.T) {
 	usage := NewUsageTracker()
-	gw, _ := NewGateway(testSigningKey, "http://localhost:11434", "real-key", usage, log.New(os.Stdout, "", 0))
+	gw, _ := NewGateway(testSigningKey, "http://localhost:11434", "real-key", usage, slog.New(slog.NewTextHandler(os.Stdout, nil)))
 
 	req := httptest.NewRequest("POST", "/v1/chat/completions", nil)
 	req.Header.Set("Authorization", "Bearer bad-token")
@@ -63,7 +63,7 @@ func TestGateway_ProxiesWithRealKey(t *testing.T) {
 	defer upstream.Close()
 
 	usage := NewUsageTracker()
-	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, log.New(os.Stdout, "", 0))
+	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, slog.New(slog.NewTextHandler(os.Stdout, nil)))
 
 	req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"test","messages":[]}`))
 	req.Header.Set("Authorization", "Bearer "+signToken("team-a"))
@@ -99,7 +99,7 @@ func TestGateway_NonCompletionEndpoint_NoUsageTracking(t *testing.T) {
 	defer upstream.Close()
 
 	usage := NewUsageTracker()
-	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, log.New(os.Stdout, "", 0))
+	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, slog.New(slog.NewTextHandler(os.Stdout, nil)))
 
 	req := httptest.NewRequest("GET", "/v1/models", nil)
 	req.Header.Set("Authorization", "Bearer "+signToken("team-a"))
@@ -128,7 +128,7 @@ func TestGateway_ResponseBodyPassedThrough(t *testing.T) {
 	defer upstream.Close()
 
 	usage := NewUsageTracker()
-	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, log.New(os.Stdout, "", 0))
+	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, slog.New(slog.NewTextHandler(os.Stdout, nil)))
 
 	req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"test","messages":[]}`))
 	req.Header.Set("Authorization", "Bearer "+signToken("team-a"))
@@ -158,7 +158,7 @@ data: [DONE]
 	defer upstream.Close()
 
 	usage := NewUsageTracker()
-	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, log.New(os.Stdout, "", 0))
+	gw, _ := NewGateway(testSigningKey, upstream.URL, "real-api-key", usage, slog.New(slog.NewTextHandler(os.Stdout, nil)))
 
 	req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"test","messages":[]}`))
 	req.Header.Set("Authorization", "Bearer "+signToken("team-a"))

balancer/pkg/llmgateway/usage.go

@@ -3,7 +3,7 @@ package llmgateway
 import (
 	"context"
 	"fmt"
-	"log"
+	"log/slog"
 	"strconv"
 	"sync"
 	"time"
@@ -52,15 +52,15 @@ func (t *UsageTracker) Add(team string, inputTokens, outputTokens int64) {
 }
 
 // FlushToAnnotations writes accumulated usage to deployment annotations and resets the counters.
-func (t *UsageTracker) FlushToAnnotations(ctx context.Context, clientset kubernetes.Interface, namespace string, logger *log.Logger) {
+func (t *UsageTracker) FlushToAnnotations(ctx context.Context, clientset kubernetes.Interface, namespace string, logger *slog.Logger) {
 	t.mu.Lock()
 	pending := t.usage
 	t.usage = make(map[string]*TeamUsage)
 	t.mu.Unlock()
 
 	for team, usage := range pending {
 		if err := t.updateTeamAnnotations(ctx, clientset, namespace, team, usage, logger); err != nil {
-			logger.Printf("Failed to flush LLM usage for team '%s': %v", team, err)
+			logger.Error("Failed to flush LLM usage", "team", team, "error", err)
 			// Put the usage back so it's not lost
 			t.Add(team, usage.InputTokens, usage.OutputTokens)
 		}
@@ -69,7 +69,7 @@ func (t *UsageTracker) FlushToAnnotations(ctx context.Context, clientset kuberne
 
 // updateTeamAnnotations uses optimistic concurrency (read resourceVersion, retry on conflict)
 // to safely increment token counters even when multiple balancer replicas are running.
-func (t *UsageTracker) updateTeamAnnotations(ctx context.Context, clientset kubernetes.Interface, namespace, team string, delta *TeamUsage, logger *log.Logger) error {
+func (t *UsageTracker) updateTeamAnnotations(ctx context.Context, clientset kubernetes.Interface, namespace, team string, delta *TeamUsage, logger *slog.Logger) error {
 	deploymentName := fmt.Sprintf("juiceshop-%s", team)
 
 	for attempt := range maxRetries {
@@ -109,14 +109,14 @@ func (t *UsageTracker) updateTeamAnnotations(ctx context.Context, clientset kube
 		if !errors.IsConflict(err) {
 			return fmt.Errorf("failed to update deployment: %w", err)
 		}
-		logger.Printf("LLM usage update conflict for team '%s' (attempt %d/%d), retrying", team, attempt+1, maxRetries)
+		logger.Warn("LLM usage update conflict, retrying", "team", team, "attempt", attempt+1, "maxRetries", maxRetries)
 	}
 
 	return fmt.Errorf("failed to update deployment after %d retries due to conflicts", maxRetries)
 }
 
 // StartFlusher periodically flushes accumulated usage to deployment annotations.
-func (t *UsageTracker) StartFlusher(ctx context.Context, clientset kubernetes.Interface, namespace string, logger *log.Logger) {
+func (t *UsageTracker) StartFlusher(ctx context.Context, clientset kubernetes.Interface, namespace string, logger *slog.Logger) {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
 	for {

balancer/pkg/llmgateway/usage_test.go

@@ -2,7 +2,7 @@ package llmgateway
 
 import (
 	"context"
-	"log"
+	"log/slog"
 	"os"
 	"testing"
 
@@ -45,7 +45,7 @@ func TestUsageTracker_FlushToAnnotations(t *testing.T) {
 	tracker := NewUsageTracker()
 	tracker.Add("team-a", 10, 20)
 
-	logger := log.New(os.Stdout, "", 0)
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	tracker.FlushToAnnotations(context.Background(), clientset, "default", logger)
 
 	// Verify annotations were updated
@@ -76,7 +76,7 @@ func TestUsageTracker_FlushRetainsOnError(t *testing.T) {
 	tracker := NewUsageTracker()
 	tracker.Add("nonexistent", 10, 20)
 
-	logger := log.New(os.Stdout, "", 0)
+	logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 	tracker.FlushToAnnotations(context.Background(), clientset, "default", logger)
 
 	// Usage should be retained

balancer/pkg/notification/notification.go

@@ -78,7 +78,7 @@ func (s *NotificationService) StartNotificationWatcher(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
-			s.bundle.Log.Printf("MultiJuicer context canceled. Exiting notification watcher.")
+			s.bundle.Log.Info("MultiJuicer context canceled. Exiting notification watcher.")
 			return
 		default:
 			s.watchConfigMap(ctx)
@@ -99,10 +99,10 @@ func (s *NotificationService) watchConfigMap(ctx context.Context) {
 	)
 	if err != nil {
 		if errors.IsNotFound(err) {
-			s.bundle.Log.Printf("Notification ConfigMap not found. Treating as no notification.")
+			s.bundle.Log.Debug("Notification ConfigMap not found. Treating as no notification.")
 			s.parseAndUpdateNotification(nil)
 		} else {
-			s.bundle.Log.Printf("Failed to get notification ConfigMap: %v", err)
+			s.bundle.Log.Error("Failed to get notification ConfigMap", "error", err)
 			return
 		}
 	} else {
@@ -117,18 +117,18 @@ func (s *NotificationService) watchConfigMap(ctx context.Context) {
 		},
 	)
 	if err != nil {
-		s.bundle.Log.Printf("Failed to start watch for notification ConfigMap: %v", err)
+		s.bundle.Log.Error("Failed to start watch for notification ConfigMap", "error", err)
 		return
 	}
 	defer watcher.Stop()
 
-	s.bundle.Log.Printf("Started watching notification ConfigMap")
+	s.bundle.Log.Debug("Started watching notification ConfigMap")
 
 	for {
 		select {
 		case event, ok := <-watcher.ResultChan():
 			if !ok {
-				s.bundle.Log.Printf("Notification ConfigMap watcher closed. Reconnecting...")
+				s.bundle.Log.Warn("Notification ConfigMap watcher closed. Reconnecting...")
 				return
 			}
 
@@ -137,11 +137,11 @@ func (s *NotificationService) watchConfigMap(ctx context.Context) {
 				configMap := event.Object.(*corev1.ConfigMap)
 				s.parseAndUpdateNotification(configMap)
 			case watch.Deleted:
-				s.bundle.Log.Printf("Notification ConfigMap deleted")
+				s.bundle.Log.Info("Notification ConfigMap deleted")
 				s.parseAndUpdateNotification(nil)
 			}
 		case <-ctx.Done():
-			s.bundle.Log.Printf("Context canceled. Exiting notification watcher.")
+			s.bundle.Log.Info("Context canceled. Exiting notification watcher.")
 			return
 		}
 	}
@@ -159,23 +159,23 @@ func (s *NotificationService) parseAndUpdateNotification(cm *corev1.ConfigMap) {
 
 	jsonData, ok := cm.Data["notification.json"]
 	if !ok {
-		s.bundle.Log.Printf("Notification ConfigMap missing 'notification.json' key")
+		s.bundle.Log.Warn("Notification ConfigMap missing 'notification.json' key")
 		s.currentNotification = nil
 		s.lastUpdate = timeutil.TruncateToMillisecond(time.Now())
 		return
 	}
 
 	var notification bundle.Notification
 	if err := json.Unmarshal([]byte(jsonData), &notification); err != nil {
-		s.bundle.Log.Printf("Failed to parse notification JSON: %v", err)
+		s.bundle.Log.Error("Failed to parse notification JSON", "error", err)
 		s.currentNotification = nil
 		s.lastUpdate = timeutil.TruncateToMillisecond(time.Now())
 		return
 	}
 
 	s.currentNotification = &notification
 	s.lastUpdate = timeutil.TruncateToMillisecond(time.Now())
-	s.bundle.Log.Printf("Updated notification: enabled=%v, message=%q", notification.Enabled, notification.Message)
+	s.bundle.Log.Info("Updated notification", "enabled", notification.Enabled, "message", notification.Message)
 }
 
 // getOrCreateConfigMap retrieves the existing notification ConfigMap or returns a new empty one.