feat(logger): Wrapper to extract and log context (#4863)

Add a contextHandler to wrap around our slog.Handler so that key-value
stored in context will be automatically added to the log as attributes.

Services will have to explicitly register what context keys they want to
show up in logs.
To register a context key to be logged (after initializing the logger):
```
logger.RegisterContextKey(ctxKey, logKey)
```

Author: Ly-Joey

go/cmd/gitter/gitter.go

@@ -33,7 +33,7 @@ import (
 type contextKey string
 
 const (
-	urlKey contextKey = "url"
+	urlKey contextKey = "repoURL"
 )
 
 const getGitEndpoint = "/getgit"
@@ -54,7 +54,7 @@ const shutdownTimeout = 10 * time.Second
 // runCmd executes a command with context cancellation handled by sending SIGINT.
 // It logs cancellation errors separately as requested.
 func runCmd(ctx context.Context, dir string, env []string, name string, args ...string) error {
-	logger.DebugContext(ctx, "Running command", slog.String("cmd", name), slog.String("url", ctx.Value(urlKey).(string)), slog.Any("args", args))
+	logger.DebugContext(ctx, "Running command", slog.String("cmd", name), slog.Any("args", args))
 	cmd := exec.CommandContext(ctx, name, args...)
 	if dir != "" {
 		cmd.Dir = dir
@@ -64,7 +64,7 @@ func runCmd(ctx context.Context, dir string, env []string, name string, args ...
 	}
 	// Use SIGINT instead of SIGKILL for graceful shutdown of subprocesses
 	cmd.Cancel = func() error {
-		logger.DebugContext(ctx, "SIGINT sent to command", slog.String("cmd", name), slog.String("url", ctx.Value(urlKey).(string)), slog.Any("args", args))
+		logger.DebugContext(ctx, "SIGINT sent to command", slog.String("cmd", name), slog.Any("args", args))
 		return cmd.Process.Signal(syscall.SIGINT)
 	}
 	// Ensure it eventually dies if it ignores SIGINT
@@ -74,13 +74,13 @@ func runCmd(ctx context.Context, dir string, env []string, name string, args ...
 	if err != nil {
 		if ctx.Err() != nil {
 			// Log separately if cancelled
-			logger.WarnContext(ctx, "Command cancelled", slog.String("cmd", name), slog.String("url", ctx.Value(urlKey).(string)), slog.Any("err", ctx.Err()))
+			logger.WarnContext(ctx, "Command cancelled", slog.String("cmd", name), slog.Any("err", ctx.Err()))
 			return fmt.Errorf("command %s cancelled: %w", name, ctx.Err())
 		}
 
 		return fmt.Errorf("command %s failed: %w, output: %s", name, err, out)
 	}
-	logger.DebugContext(ctx, "Command completed successfully", slog.String("cmd", name), slog.String("url", ctx.Value(urlKey).(string)), slog.String("out", string(out)))
+	logger.DebugContext(ctx, "Command completed successfully", slog.String("cmd", name), slog.String("out", string(out)))
 
 	return nil
 }
@@ -137,7 +137,7 @@ func fetchBlob(ctx context.Context, url string, forceUpdate bool) ([]byte, error
 
 	// Check if we need to fetch
 	if forceUpdate || !ok || time.Since(accessTime) > fetchTimeout {
-		logger.InfoContext(ctx, "Fetching git blob", slog.String("url", ctx.Value(urlKey).(string)), slog.Duration("sinceAccessTime", time.Since(accessTime)))
+		logger.InfoContext(ctx, "Fetching git blob", slog.Duration("sinceAccessTime", time.Since(accessTime)))
 		if _, err := os.Stat(path.Join(repoPath, ".git")); os.IsNotExist(err) {
 			// Clone
 			err := runCmd(ctx, "", []string{"GIT_TERMINAL_PROMPT=0"}, "git", "clone", "--", url, repoPath)
@@ -156,7 +156,7 @@ func fetchBlob(ctx context.Context, url string, forceUpdate bool) ([]byte, error
 			if err != nil && isIndexLockError(err) {
 				// index.lock exists, likely a previous git reset got terminated and wasn't cleaned up properly.
 				// We can remove the file and retry the command
-				logger.WarnContext(ctx, "index.lock exists, attempting to remove and retry", slog.String("url", ctx.Value(urlKey).(string)))
+				logger.WarnContext(ctx, "index.lock exists, attempting to remove and retry")
 				indexLockPath := filepath.Join(repoPath, ".git", "index.lock")
 				if err := os.Remove(indexLockPath); err != nil {
 					return nil, fmt.Errorf("failed to remove index.lock in %s: %w", repoPath, err)
@@ -169,7 +169,7 @@ func fetchBlob(ctx context.Context, url string, forceUpdate bool) ([]byte, error
 			}
 		}
 
-		logger.InfoContext(ctx, "Archiving git blob", slog.String("url", ctx.Value(urlKey).(string)))
+		logger.InfoContext(ctx, "Archiving git blob")
 		// Archive
 		// tar --zstd -cf <archivePath> -C "<gitStorePath>/<repoDirName>" .
 		// using -C to archive the relative path so it unzips nicely
@@ -201,6 +201,7 @@ func fetchBlob(ctx context.Context, url string, forceUpdate bool) ([]byte, error
 
 func main() {
 	logger.InitGlobalLogger()
+	logger.RegisterContextKey(urlKey, "repoURL")
 	defer logger.Close()
 
 	port := flag.Int("port", 8888, "Listen port")
@@ -283,7 +284,7 @@ func gitHandler(w http.ResponseWriter, r *http.Request) {
 
 	ctx := context.WithValue(r.Context(), urlKey, url)
 
-	logger.InfoContext(ctx, "Received request", slog.String("url", url), slog.Bool("forceUpdate", forceUpdate), slog.String("remoteAddr", r.RemoteAddr))
+	logger.InfoContext(ctx, "Received request", slog.Bool("forceUpdate", forceUpdate), slog.String("remoteAddr", r.RemoteAddr))
 	// If request came from a local ip, don't do the check
 	if !isLocalRequest(r) {
 		// Check if url starts with protocols: http(s)://, git://, ssh://, (s)ftp://
@@ -308,7 +309,7 @@ func gitHandler(w http.ResponseWriter, r *http.Request) {
 	})
 
 	if err != nil {
-		logger.ErrorContext(ctx, "Error fetching/archiving blob", slog.String("url", ctx.Value(urlKey).(string)), slog.Any("error", err))
+		logger.ErrorContext(ctx, "Error fetching blob", slog.Any("error", err))
 		if isAuthError(err) {
 			http.Error(w, fmt.Sprintf("Error fetching blob: %v", err), http.StatusForbidden)
 			return
@@ -322,11 +323,11 @@ func gitHandler(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Disposition", "attachment; filename=\"git-blob.zst\"")
 	w.WriteHeader(http.StatusOK)
 	if _, err := io.Copy(w, bytes.NewReader(fileData.([]byte))); err != nil {
-		logger.ErrorContext(ctx, "Error copying file", slog.String("url", url), slog.Any("error", err))
+		logger.ErrorContext(ctx, "Error copying file", slog.Any("error", err))
 		http.Error(w, "Error copying file", http.StatusInternalServerError)
 
 		return
 	}
 
-	logger.InfoContext(ctx, "Request completed successfully", slog.String("url", ctx.Value(urlKey).(string)))
+	logger.InfoContext(ctx, "Request completed successfully")
 }

go/logger/contexthandler.go

@@ -0,0 +1,46 @@
+package logger
+
+import (
+	"context"
+	"log/slog"
+	"sync"
+)
+
+var (
+	ctxKeysMu sync.RWMutex
+	ctxKeys   = make(map[any]string)
+)
+
+// RegisterContextKey registers a context key that should be automatically
+// extracted and logged if present in the context.
+func RegisterContextKey(ctxKey any, logKey string) {
+	ctxKeysMu.Lock()
+	defer ctxKeysMu.Unlock()
+	ctxKeys[ctxKey] = logKey
+}
+
+// contextHandler is a wrapper around slog.Handler that automatically
+// extracts registered context keys and logs their values as part of the attributes.
+type contextHandler struct {
+	slog.Handler
+}
+
+func (h *contextHandler) Handle(ctx context.Context, r slog.Record) error {
+	ctxKeysMu.RLock()
+	for ctxKey, logKey := range ctxKeys {
+		if val := ctx.Value(ctxKey); val != nil {
+			r.AddAttrs(slog.Any(logKey, val))
+		}
+	}
+	ctxKeysMu.RUnlock()
+
+	return h.Handler.Handle(ctx, r)
+}
+
+func (h *contextHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
+	return &contextHandler{h.Handler.WithAttrs(attrs)}
+}
+
+func (h *contextHandler) WithGroup(name string) slog.Handler {
+	return &contextHandler{h.Handler.WithGroup(name)}
+}

go/logger/init.go

@@ -20,7 +20,7 @@ import (
 )
 
 var (
-	slogLogger = slog.New(newLocalHandler(os.Stdout))
+	slogLogger = slog.New(&contextHandler{newLocalHandler(os.Stdout)})
 	once       sync.Once
 	tp         *sdktrace.TracerProvider
 	projectID  string
@@ -49,7 +49,7 @@ func InitGlobalLogger() {
 			initTracing(context.Background(), projectID, serviceName)
 		}
 		handler := slog.NewJSONHandler(os.Stdout, cloudHandlerOptions())
-		slogLogger = slog.New(handler)
+		slogLogger = slog.New(&contextHandler{handler})
 	})
 }