Fix logging behavior around setting goproc limits

cmd/spicedb/main.go

@@ -3,17 +3,13 @@ package main
 import (
 	"errors"
 	"fmt"
-	"log/slog"
 	"os"
 
-	"github.com/KimMachineGun/automemlimit/memlimit"
 	mcobra "github.com/muesli/mango-cobra"
 	"github.com/muesli/roff"
 	"github.com/rs/zerolog"
-	slogzerolog "github.com/samber/slog-zerolog/v2"
 	"github.com/sercand/kuberesolver/v5"
 	"github.com/spf13/cobra"
-	"go.uber.org/automaxprocs/maxprocs"
 	"google.golang.org/grpc/balancer"
 
 	_ "google.golang.org/grpc/xds"
@@ -29,33 +25,16 @@ import (
 var errParsing = errors.New("parsing error")
 
 func main() {
+	// Set up root logger
+	// This will typically be overwritten by the logging setup for a given command.
+	log.SetGlobalLogger(zerolog.New(os.Stderr).Level(zerolog.InfoLevel))
+
 	// Enable Kubernetes gRPC resolver
 	kuberesolver.RegisterInCluster()
 
 	// Enable consistent hashring gRPC load balancer
 	balancer.Register(cmdutil.ConsistentHashringBuilder)
 
-	globalLogger := zerolog.New(os.Stderr).Level(zerolog.DebugLevel)
-	log.SetGlobalLogger(globalLogger)
-	slogger := slog.New(slogzerolog.Option{Level: slog.LevelDebug, Logger: &globalLogger}.NewZerologHandler())
-
-	undo, err := maxprocs.Set(maxprocs.Logger(globalLogger.Printf))
-	if err != nil {
-		log.Fatal().Err(err).Msg("failed to set maxprocs")
-	}
-	defer undo()
-
-	_, _ = memlimit.SetGoMemLimitWithOpts(
-		memlimit.WithRatio(0.9),
-		memlimit.WithProvider(
-			memlimit.ApplyFallback(
-				memlimit.FromCgroup,
-				memlimit.FromSystem,
-			),
-		),
-		memlimit.WithLogger(slogger),
-	)
-
 	// Create a root command
 	rootCmd := cmd.NewRootCommand("spicedb")
 	rootCmd.SetFlagErrorFunc(func(cmd *cobra.Command, err error) error {

pkg/cmd/limits/limits.go

@@ -0,0 +1,51 @@
+package limits
+
+import (
+	"log/slog"
+
+	"github.com/KimMachineGun/automemlimit/memlimit"
+	"github.com/jzelinskie/cobrautil/v2"
+	"github.com/rs/zerolog"
+	slogzerolog "github.com/samber/slog-zerolog/v2"
+	"github.com/spf13/cobra"
+	"go.uber.org/automaxprocs/maxprocs"
+
+	log "github.com/authzed/spicedb/internal/logging"
+)
+
+// SetLimitsRunE wraps the RunFunc with setup logic for memory limits and maxprocs
+// limits for the go process. Implementing it here means that we can use already-
+// established loggers and wrap only those commands that require it.
+func SetLimitsRunE() cobrautil.CobraRunFunc {
+	return func(cmd *cobra.Command, args []string) error {
+		// Need to invert the slog => zerolog map so that we can get the correct
+		// slog loglevel for memlimit logging
+		logLevelMap := make(map[zerolog.Level]slog.Level, len(slogzerolog.LogLevels))
+		for sLevel, zLevel := range slogzerolog.LogLevels {
+			logLevelMap[zLevel] = sLevel
+		}
+
+		logLevel := logLevelMap[zerolog.DefaultContextLogger.GetLevel()]
+
+		slogger := slog.New(slogzerolog.Option{Level: logLevel, Logger: zerolog.DefaultContextLogger}.NewZerologHandler())
+
+		undo, err := maxprocs.Set(maxprocs.Logger(zerolog.DefaultContextLogger.Printf))
+		if err != nil {
+			log.Fatal().Err(err).Msg("failed to set maxprocs")
+		}
+		defer undo()
+
+		_, _ = memlimit.SetGoMemLimitWithOpts(
+			memlimit.WithRatio(0.9),
+			memlimit.WithProvider(
+				memlimit.ApplyFallback(
+					memlimit.FromCgroup,
+					memlimit.FromSystem,
+				),
+			),
+			memlimit.WithLogger(slogger),
+		)
+
+		return nil
+	}
+}

pkg/cmd/server/defaults.go

@@ -35,6 +35,7 @@ import (
 	datastoremw "github.com/authzed/spicedb/internal/middleware/datastore"
 	dispatchmw "github.com/authzed/spicedb/internal/middleware/dispatcher"
 	"github.com/authzed/spicedb/internal/middleware/servicespecific"
+	"github.com/authzed/spicedb/pkg/cmd/limits"
 	"github.com/authzed/spicedb/pkg/datastore"
 	logmw "github.com/authzed/spicedb/pkg/middleware/logging"
 	"github.com/authzed/spicedb/pkg/middleware/requestid"
@@ -71,6 +72,7 @@ func DefaultPreRunE(programName string) cobrautil.CobraRunFunc {
 				logging.SetGlobalLogger(logger)
 			}),
 		).RunE(),
+		limits.SetLimitsRunE(),
 		cobraotel.New("spicedb",
 			cobraotel.WithLogger(zerologr.New(&logging.Logger)),
 		).RunE(),