Add wrapper for context logging (and reg)

Author: Ly-Joey

go/cmd/gitter/gitter.go

@@ -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 @@
+// contextHandler is a wrapper around slog.Handler that automatically
+// extracts registered context keys and logs their values as part of the attributes.
+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
+}
+
+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

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