2 files changed, 73 insertions, 24 deletions
diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/logger.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/logger.go
new file mode 100644
index 000000000..4cf0b589d
--- /dev/null
+++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/logger.go
@@ -0,0 +1,13 @@
+package memlimit
+
+import (
+	"context"
+	"log/slog"
+)
+
+type noopLogger struct{}
+
+func (noopLogger) Enabled(context.Context, slog.Level) bool  { return false }
+func (noopLogger) Handle(context.Context, slog.Record) error { return nil }
+func (d noopLogger) WithAttrs([]slog.Attr) slog.Handler      { return d }
+func (d noopLogger) WithGroup(string) slog.Handler           { return d }
diff --git a/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go b/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go
index 0d7a9853c..f7aaf071c 100644
--- a/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go
+++ b/vendor/github.com/KimMachineGun/automemlimit/memlimit/memlimit.go
@@ -3,8 +3,7 @@ package memlimit
 import (
 	"errors"
 	"fmt"
-	"io"
-	"log"
+	"log/slog"
 	"math"
 	"os"
 	"runtime/debug"
@@ -12,8 +11,9 @@ import (
 )
 
 const (
-	envGOMEMLIMIT         = "GOMEMLIMIT"
-	envAUTOMEMLIMIT       = "AUTOMEMLIMIT"
+	envGOMEMLIMIT   = "GOMEMLIMIT"
+	envAUTOMEMLIMIT = "AUTOMEMLIMIT"
+	// Deprecated: use memlimit.WithLogger instead
 	envAUTOMEMLIMIT_DEBUG = "AUTOMEMLIMIT_DEBUG"
 
 	defaultAUTOMEMLIMIT = 0.9
@@ -25,7 +25,7 @@ var (
 )
 
 type config struct {
-	logger   *log.Logger
+	logger   *slog.Logger
 	ratio    float64
 	provider Provider
 }
@@ -42,6 +42,25 @@ func WithRatio(ratio float64) Option {
 	}
 }
 
+// WithProvider configures the provider.
+//
+// Default: FromCgroup
+func WithProvider(provider Provider) Option {
+	return func(cfg *config) {
+		cfg.provider = provider
+	}
+}
+
+// WithLogger configures the logger.
+// It automatically attaches the "package" attribute to the logs.
+//
+// Default: slog.New(noopLogger{})
+func WithLogger(logger *slog.Logger) Option {
+	return func(cfg *config) {
+		cfg.logger = memlimitLogger(logger)
+	}
+}
+
 // WithEnv configures whether to use environment variables.
 //
 // Default: false
@@ -51,13 +70,11 @@ func WithEnv() Option {
 	return func(cfg *config) {}
 }
 
-// WithProvider configures the provider.
-//
-// Default: FromCgroup
-func WithProvider(provider Provider) Option {
-	return func(cfg *config) {
-		cfg.provider = provider
+func memlimitLogger(logger *slog.Logger) *slog.Logger {
+	if logger == nil {
+		return slog.New(noopLogger{})
 	}
+	return logger.With(slog.String("package", "github.com/KimMachineGun/automemlimit/memlimit"))
 }
 
 // SetGoMemLimitWithOpts sets GOMEMLIMIT with options and environment variables.
@@ -74,54 +91,69 @@ func WithProvider(provider Provider) Option {
 // Options:
 //   - WithRatio
 //   - WithProvider
+//   - WithLogger
 func SetGoMemLimitWithOpts(opts ...Option) (_ int64, _err error) {
+	// init config
 	cfg := &config{
-		logger:   log.New(io.Discard, "", log.LstdFlags),
+		logger:   slog.New(noopLogger{}),
 		ratio:    defaultAUTOMEMLIMIT,
 		provider: FromCgroup,
 	}
-	if os.Getenv(envAUTOMEMLIMIT_DEBUG) == "true" {
-		cfg.logger = log.Default()
+	// TODO: remove this
+	if debug, ok := os.LookupEnv(envAUTOMEMLIMIT_DEBUG); ok {
+		defaultLogger := memlimitLogger(slog.Default())
+		defaultLogger.Warn("AUTOMEMLIMIT_DEBUG is deprecated, use memlimit.WithLogger instead")
+		if debug == "true" {
+			cfg.logger = defaultLogger
+		}
 	}
 	for _, opt := range opts {
 		opt(cfg)
 	}
+
+	// log error if any on return
 	defer func() {
 		if _err != nil {
-			cfg.logger.Println(_err)
+			cfg.logger.Error("failed to set GOMEMLIMIT", slog.Any("error", _err))
 		}
 	}()
 
+	// parse experiments
 	exps, err := parseExperiments()
 	if err != nil {
 		return 0, fmt.Errorf("failed to parse experiments: %w", err)
 	}
 	if exps.System {
-		cfg.logger.Println("system experiment is enabled: using system memory limit as a fallback")
+		cfg.logger.Info("system experiment is enabled: using system memory limit as a fallback")
 		cfg.provider = ApplyFallback(cfg.provider, FromSystem)
 	}
 
+	// capture the current GOMEMLIMIT for rollback in case of panic
 	snapshot := debug.SetMemoryLimit(-1)
 	defer func() {
-		err := recover()
-		if err != nil {
+		panicErr := recover()
+		if panicErr != nil {
 			if _err != nil {
-				cfg.logger.Println(_err)
+				cfg.logger.Error("failed to set GOMEMLIMIT", slog.Any("error", _err))
 			}
-			_err = fmt.Errorf("panic during setting the Go's memory limit, rolling back to previous value %d: %v", snapshot, err)
+			_err = fmt.Errorf("panic during setting the Go's memory limit, rolling back to previous limit %d: %v",
+				snapshot, panicErr,
+			)
 			debug.SetMemoryLimit(snapshot)
 		}
 	}()
 
+	// check if GOMEMLIMIT is already set
 	if val, ok := os.LookupEnv(envGOMEMLIMIT); ok {
-		cfg.logger.Printf("GOMEMLIMIT is set already, skipping: %s\n", val)
+		cfg.logger.Info("GOMEMLIMIT is already set, skipping", slog.String(envGOMEMLIMIT, val))
 		return 0, nil
 	}
 
+	// parse AUTOMEMLIMIT
 	ratio := cfg.ratio
 	if val, ok := os.LookupEnv(envAUTOMEMLIMIT); ok {
 		if val == "off" {
-			cfg.logger.Printf("AUTOMEMLIMIT is set to off, skipping\n")
+			cfg.logger.Info("AUTOMEMLIMIT is set to off, skipping")
 			return 0, nil
 		}
 		_ratio, err := strconv.ParseFloat(val, 64)
@@ -131,20 +163,24 @@ func SetGoMemLimitWithOpts(opts ...Option) (_ int64, _err error) {
 		ratio = _ratio
 	}
 
+	// set GOMEMLIMIT
 	limit, err := setGoMemLimit(ApplyRatio(cfg.provider, ratio))
 	if err != nil {
 		if errors.Is(err, ErrNoLimit) {
-			cfg.logger.Printf("memory is not limited, skipping: %v\n", err)
+			cfg.logger.Info("memory is not limited, skipping: %v", err)
 			return 0, nil
 		}
 		return 0, fmt.Errorf("failed to set GOMEMLIMIT: %w", err)
 	}
 
-	cfg.logger.Printf("GOMEMLIMIT=%d\n", limit)
+	cfg.logger.Info("GOMEMLIMIT is updated", slog.Int64(envGOMEMLIMIT, limit))
 
 	return limit, nil
 }
 
+// SetGoMemLimitWithEnv sets GOMEMLIMIT with the value from the environment variables.
+// Since WithEnv is deprecated, this function is equivalent to SetGoMemLimitWithOpts().
+// Deprecated: use SetGoMemLimitWithOpts instead.
 func SetGoMemLimitWithEnv() {
 	_, _ = SetGoMemLimitWithOpts()
 }