Include undo and apply functional options pattern for setting logger

.golangci.yml

@@ -122,6 +122,9 @@ issues:
     - path: gomaxecs.go
       linters:
         - gochecknoinits # enable init function for setting GOMAXPROCS.
+    - path: maxprocs/maxprocs_test.go
+      linters:
+        - paralleltest # disable paralleltest for testing GOMAXPROCS env variable.
 
 linters-settings:
   depguard:

gomaxecs.go

@@ -29,5 +29,5 @@ import (
 )
 
 func init() {
-	maxprocs.Set(log.Default())
+	_, _ = maxprocs.Set(maxprocs.WithLogger(log.Printf))
 }

internal/config/config.go

@@ -33,10 +33,10 @@ const (
 	httpTimeout = 5
 )
 
-func New() Config {
-	uri := metadataURI()
+func New(opts ...Option) Config {
+	uri := GetECSMetadataURI()
 
-	return Config{
+	cfg := Config{
 		TaskMetadataURI:      uri + taskPath,
 		ContainerMetadataURI: uri,
 		Client: Client{
@@ -50,20 +50,30 @@ func New() Config {
 			ResponseHeaderTimeout: time.Second,
 		},
 	}
+
+	for _, opt := range opts {
+		opt(&cfg)
+	}
+
+	return cfg
 }
 
-func metadataURI() string {
+// GetECSMetadataURI returns the ECS metadata URI.
+func GetECSMetadataURI() string {
 	uri := os.Getenv(metaURIEnv)
 	return strings.TrimRight(uri, "/")
 }
 
-// Config represents the packagge configuration.
+// Config represents the package configuration.
 type Config struct {
 	ContainerMetadataURI string
 	TaskMetadataURI      string
 	Client               Client
+	log                  logger
 }
 
+type logger func(format string, args ...any)
+
 // Client represents the HTTP client configuration.
 type Client struct {
 	HTTPTimeout           time.Duration
@@ -75,3 +85,19 @@ type Client struct {
 	TLSHandshakeTimeout   time.Duration
 	ResponseHeaderTimeout time.Duration
 }
+
+func (c Config) Log(format string, args ...any) {
+	if c.log != nil {
+		c.log(format, args...)
+	}
+}
+
+// WithLogger sets the logger for the config.
+func WithLogger(logger logger) Option {
+	return func(cfg *Config) {
+		cfg.log = logger
+	}
+}
+
+// Option represents a configuration option for the config.
+type Option func(*Config)

internal/config/config_test.go

@@ -1,6 +1,8 @@
 package config_test
 
 import (
+	"bytes"
+	"log"
 	"testing"
 	"time"
 
@@ -9,7 +11,7 @@ import (
 	"github.com/rdforte/gomaxecs/internal/config"
 )
 
-func TestConfig_LoadConfiguration(t *testing.T) {
+func TestConfig_New_LoadConfiguration(t *testing.T) {
 	metaURIEnv := "ECS_CONTAINER_METADATA_URI_V4"
 	uri := "mock-ecs-metadata-uri/"
 	t.Setenv(metaURIEnv, uri)
@@ -34,3 +36,48 @@ func TestConfig_LoadConfiguration(t *testing.T) {
 
 	assert.Equal(t, wantCfg, cfg)
 }
+
+func TestConfig_New_AppliesOptions(t *testing.T) {
+	t.Parallel()
+
+	opt1 := mockOption{}
+	opt2 := mockOption{}
+
+	config.New(opt1.Apply, opt2.Apply)
+
+	assert.True(t, opt1.isApplied)
+	assert.True(t, opt2.isApplied)
+}
+
+func TestConfig_WithLogger_LogsMessage(t *testing.T) {
+	t.Parallel()
+
+	buf := new(bytes.Buffer)
+	logger := log.New(buf, "", 0)
+
+	cfg := config.New(config.WithLogger(logger.Printf))
+
+	cfg.Log("test log: %s, %s", "arg1", "arg2")
+
+	wantLog := "test log: arg1, arg2\n"
+	assert.Equal(t, wantLog, buf.String())
+}
+
+func TestConfig_GetECSMetadataURI_RetrievesMetadataURIFromEnv(t *testing.T) {
+	metaURIEnv := "ECS_CONTAINER_METADATA_URI_V4"
+	uri := "mock-ecs-metadata-uri/"
+	t.Setenv(metaURIEnv, uri)
+
+	got := config.GetECSMetadataURI()
+
+	want := "mock-ecs-metadata-uri"
+	assert.Equal(t, want, got)
+}
+
+type mockOption struct {
+	isApplied bool
+}
+
+func (m *mockOption) Apply(_ *config.Config) {
+	m.isApplied = true
+}

internal/task/meta.go

@@ -29,20 +29,20 @@ import (
 	"github.com/rdforte/gomaxecs/internal/client"
 )
 
-// TaskMeta represents the ECS Task Metadata.
+// taskMeta represents the ECS Task Metadata.
 type taskMeta struct {
 	Containers []container `json:"Containers"`
 	Limits     limit       `json:"Limits"` // this is optional in the response
 }
 
-// Container represents the ECS Container Metadata.
+// container represents the ECS Container Metadata.
 type container struct {
 	//nolint:tagliatelle // ECS Agent inconsistency. All fields adhere to goPascal but this one.
 	DockerID string `json:"DockerId"`
 	Limits   limit  `json:"Limits"`
 }
 
-// Limit contains the CPU limit.
+// limit contains the CPU limit.
 type limit struct {
 	CPU float64 `json:"CPU"`
 }

internal/task/task.go

@@ -33,6 +33,7 @@ import (
 
 const (
 	cpuUnits = 10
+	minCPU   = 1
 )
 
 var errNoCPULimit = errors.New("no CPU limit found for task or container")
@@ -84,15 +85,10 @@ func (t *Task) GetMaxProcs(ctx context.Context) (int, error) {
 	}
 
 	if containerCPULimit == 0 {
-		minThreads := 1
-		return max(int(task.Limits.CPU), minThreads), nil
+		return max(int(task.Limits.CPU), minCPU), nil
 	}
 
-	cpu := int(containerCPULimit) >> cpuUnits
-	// Set a minimum of 1 for containers with less than 1 vCPU
-	if cpu == 0 {
-		cpu = 1
-	}
+	cpu := max(int(containerCPULimit)>>cpuUnits, minCPU)
 
 	taskCPULimit := int(task.Limits.CPU)
 	if taskCPULimit > 0 {

internal/task/task_test.go

@@ -158,6 +158,8 @@ func TestTask_GetMaxProcs_GetsCPUUsingContainerLimit(t *testing.T) {
 			taskCPU:      16,
 			testServer:   testServerContainerLimit,
 		},
+		// For tasks that are hosted on Amazon EC2 instances, the CPU limit is optional.
+		// https://docs.aws.amazon.com/AmazonECS/latest/developerguide/task_definition_parameters.html#task_size
 		{
 			name:         "should get cpu of 16 when task CPU limit is 0 and container CPU limit is 16384 vCPU",
 			wantCPU:      16,

maxprocs/maxprocs.go

@@ -2,24 +2,70 @@ package maxprocs
 
 import (
 	"context"
-	"log"
+	"fmt"
+	"os"
 	"runtime"
 
 	"github.com/rdforte/gomaxecs/internal/config"
-	"github.com/rdforte/gomaxecs/internal/task"
+	ecstask "github.com/rdforte/gomaxecs/internal/task"
 )
 
+const maxProcsKey = "GOMAXPROCS"
+
 // Set sets GOMAXPROCS based on the CPU limit of the container and the task.
-func Set(log *log.Logger) {
-	cfg := config.New()
-	t := task.New(cfg)
+// returns a function to reset GOMAXPROCS to its previous value and an error if one occurred.
+// If the GOMAXPROCS environment variable is set, it will honor that value.
+func Set(opts ...config.Option) (func(), error) {
+	cfg := config.New(opts...)
+	task := ecstask.New(cfg)
+
+	undoNoop := func() {
+		cfg.Log("maxprocs: No GOMAXPROCS change to reset")
+	}
+
+	if procs, ok := shouldHonorGOMAXPROCSEnv(); ok {
+		cfg.Log("maxprocs: Honoring GOMAXPROCS=%q as set in environment", procs)
+		return undoNoop, nil
+	}
 
-	procs, err := t.GetMaxProcs(context.Background())
+	prevProcs := prevMaxProcs()
+	undo := func() {
+		cfg.Log("maxprocs: Resetting GOMAXPROCS to %v", prevProcs)
+		setMaxProcs(prevProcs)
+	}
+
+	procs, err := task.GetMaxProcs(context.Background())
 	if err != nil {
-		log.Println("failed to set GOMAXPROCS:", err)
-		return
+		cfg.Log("maxprocs: Failed to set GOMAXPROCS:", err)
+		return undo, fmt.Errorf("failed to set GOMAXPROCS: %w", err)
 	}
 
+	setMaxProcs(procs)
+	cfg.Log("maxprocs: Updated GOMAXPROCS=%v", procs)
+
+	return undo, nil
+}
+
+// shouldHonorGOMAXPROCSEnv returns the GOMAXPROCS environment variable if present
+// and a boolean indicating if it should be honored.
+func shouldHonorGOMAXPROCSEnv() (string, bool) {
+	return os.LookupEnv(maxProcsKey)
+}
+
+func prevMaxProcs() int {
+	return runtime.GOMAXPROCS(0)
+}
+
+func setMaxProcs(procs int) {
 	runtime.GOMAXPROCS(procs)
-	log.Println("GOMAXPROCS set to:", procs)
+}
+
+// WithLogger sets the logger. By default, no logger is set.
+func WithLogger(printf func(format string, args ...any)) config.Option {
+	return config.WithLogger(printf)
+}
+
+// IsECS returns true if detected ECS environment.
+func IsECS() bool {
+	return len(config.GetECSMetadataURI()) > 0
 }