From 7867ba035c8382a81e94aebe3e8ef761744ea6d5 Mon Sep 17 00:00:00 2001 From: dajneem23 Date: Mon, 4 May 2026 11:27:03 +0700 Subject: [PATCH 1/2] feat(agent): enhance Run function with context support and error handling; add logging context utilities --- agent/cmd/cmd.go | 63 +++++++++++++----- agent/cmd/cmd_test.go | 45 +++++++++++-- agent/main.go | 15 ++++- utils/log/klog.go | 100 ++++++++++++++++++++++++++++ utils/log/klog_test.go | 148 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 349 insertions(+), 22 deletions(-) create mode 100644 utils/log/klog.go create mode 100644 utils/log/klog_test.go diff --git a/agent/cmd/cmd.go b/agent/cmd/cmd.go index b887d6b45..67de707b7 100644 --- a/agent/cmd/cmd.go +++ b/agent/cmd/cmd.go @@ -14,6 +14,7 @@ package cmd import ( + "context" "flag" "fmt" "net/http" @@ -112,7 +113,10 @@ func WithEffect(f func()) Option { } // Run runs the agent. -func Run(flags *Flags, opts ...Option) { +func Run(ctx context.Context, flags *Flags, opts ...Option) error { + ctx, cancel := context.WithCancel(ctx) + defer cancel() + validateRequiredPorts(flags) var overrides options @@ -153,7 +157,7 @@ func Run(flags *Flags, opts ...Option) { if stats == nil { s, closer, err := metrics.New(config.Metrics, flags.KrakenCluster) if err != nil { - log.Fatalf("Failed to init metrics: %s", err) + return fmt.Errorf("init metrics: %w", err) } stats = s defer closers.Close(closer) @@ -162,7 +166,7 @@ func Run(flags *Flags, opts ...Option) { if flags.PeerIP == "" { localIP, err := netutil.GetLocalIP() if err != nil { - log.Fatalf("Error getting local ip: %s", err) + return fmt.Errorf("get local IP: %w", err) } flags.PeerIP = localIP } @@ -174,40 +178,40 @@ func Run(flags *Flags, opts ...Option) { pctx, err := core.NewPeerContext( config.PeerIDFactory, flags.Zone, flags.KrakenCluster, flags.PeerIP, flags.PeerPort, false) if err != nil { - log.Fatalf("Failed to create peer context: %s", err) + return fmt.Errorf("create peer context: %w", err) } cads, err := store.NewCADownloadStore(config.CADownloadStore, stats) if err != nil { - log.Fatalf("Failed to create local store: %s", err) + return fmt.Errorf("create CA download store: %w", err) } netevents, err := networkevent.NewProducer(config.NetworkEvent) if err != nil { - log.Fatalf("Failed to create network event producer: %s", err) + return fmt.Errorf("create network event producer: %w", err) } trackers, err := config.Tracker.Build() if err != nil { - log.Fatalf("Error building tracker upstream: %s", err) + return fmt.Errorf("create tracker client: %w", err) } go trackers.Monitor(nil) tls, err := config.TLS.BuildClient() if err != nil { - log.Fatalf("Error building client tls config: %s", err) + return fmt.Errorf("build TLS config: %w", err) } announceClient := announceclient.New(pctx, trackers, tls) sched, err := scheduler.NewAgentScheduler( config.Scheduler, stats, pctx, cads, netevents, trackers, announceClient, tls) if err != nil { - log.Fatalf("Error creating scheduler: %s", err) + return fmt.Errorf("create scheduler: %w", err) } buildIndexes, err := config.BuildIndex.Build() if err != nil { - log.Fatalf("Error building build-index upstream: %s", err) + return fmt.Errorf("build build-index upstream: %w", err) } tagClient := tagclient.NewClusterClient(buildIndexes, tls) @@ -216,7 +220,7 @@ func Run(flags *Flags, opts ...Option) { registry, err := config.Registry.Build(config.Registry.ReadOnlyParameters(transferer, cads, stats)) if err != nil { - log.Fatalf("Failed to init registry: %s", err) + return fmt.Errorf("init registry: %w", err) } registryAddr := fmt.Sprintf("127.0.0.1:%d", flags.AgentRegistryPort) @@ -228,13 +232,14 @@ func Run(flags *Flags, opts ...Option) { } containerRuntimeFactory, err := containerruntime.NewFactory(containerRuntimeCfg, registryAddr) if err != nil { - log.Fatalf("Failed to create container runtime factory: %s", err) + return fmt.Errorf("create container runtime factory: %w", err) } agentServer := agentserver.New( config.AgentServer, stats, cads, sched, tagClient, announceClient, containerRuntimeFactory) addr := fmt.Sprintf(":%d", flags.AgentServerPort) - log.Infof("Starting agent server on %s", addr) + log.InfoS("starting agent server", "addr", addr) + errCh := make(chan error, 2) heartbeatTicker := &timeTicker{inner: time.NewTicker(10 * time.Second)} heartbeatDone := make(chan struct{}) var heartbeatStop sync.Once @@ -250,7 +255,9 @@ func Run(flags *Flags, opts ...Option) { go func() { if err := http.ListenAndServe(addr, agentServer.Handler()); err != nil { stopHeartbeat() - log.Fatal(err) + log.ErrorS(err, "agent server exited") + cancel() + errCh <- err } }() @@ -258,7 +265,9 @@ func Run(flags *Flags, opts ...Option) { go func() { if err := registry.ListenAndServe(); err != nil { stopHeartbeat() - log.Fatal(err) + log.ErrorS(err, "registry exited") + cancel() + errCh <- err } }() @@ -271,7 +280,29 @@ func Run(flags *Flags, opts ...Option) { "registry_backup": config.RegistryBackup}, nginx.WithTLS(config.TLS)); err != nil { stopHeartbeat() - log.Fatal(err) + log.ErrorS(err, "nginx exited") + cancel() + errCh <- err + } + + return waitForShutdown(ctx, errCh) +} + +// waitForShutdown blocks until ctx is cancelled or an error arrives on errCh. +// When both are ready simultaneously (i.e. a goroutine called cancel then sent +// to errCh), the error is returned instead of nil so callers see the root cause. +func waitForShutdown(ctx context.Context, errCh <-chan error) error { + select { + case <-ctx.Done(): + select { + case err := <-errCh: + return err + default: + } + log.InfoS("shutting down", "reason", ctx.Err()) + return nil + case err := <-errCh: + return err } } diff --git a/agent/cmd/cmd_test.go b/agent/cmd/cmd_test.go index 9ae7d5c39..65076c62a 100644 --- a/agent/cmd/cmd_test.go +++ b/agent/cmd/cmd_test.go @@ -1,6 +1,8 @@ package cmd import ( + "context" + "errors" "flag" "fmt" "os" @@ -112,7 +114,7 @@ func TestRunValidation(t *testing.T) { for _, test := range tests { t.Run(test.desc, func(t *testing.T) { assert.PanicsWithValue(t, test.panic, func() { - Run(&test.flags) + _ = Run(context.Background(), &test.flags) //nolint:errcheck }) }) } @@ -131,7 +133,8 @@ func TestRunUsesProvidedConfig(t *testing.T) { called := false assert.PanicsWithValue(t, sentinel, func() { - Run( + _ = Run( //nolint:errcheck + context.Background(), flags, WithConfig(Config{}), WithMetrics(tally.NewTestScope("", nil)), @@ -159,7 +162,8 @@ func TestRunPanicsWhenConfigLoadFails(t *testing.T) { expected := fmt.Sprintf("open %s: no such file or directory", missing) assert.PanicsWithError(t, expected, func() { - Run( + _ = Run( //nolint:errcheck + context.Background(), flags, WithMetrics(tally.NewTestScope("", nil)), WithLogger(zap.NewNop()), @@ -188,7 +192,8 @@ func TestRunPanicsWhenSecretsLoadFails(t *testing.T) { expected := fmt.Sprintf("open %s: no such file or directory", missingSecrets) assert.PanicsWithError(t, expected, func() { - Run( + _ = Run( //nolint:errcheck + context.Background(), flags, WithMetrics(tally.NewTestScope("", nil)), WithLogger(zap.NewNop()), @@ -288,3 +293,35 @@ func (t clockTicker) Chan() <-chan time.Time { func (t clockTicker) Stop() { t.ticker.Stop() } + +func TestWaitForShutdown_ExternalCancel(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + errCh := make(chan error, 2) + cancel() + err := waitForShutdown(ctx, errCh) + assert.NoError(t, err) +} + +func TestWaitForShutdown_ErrorReceived(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + errCh := make(chan error, 2) + sentinel := errors.New("internal failure") + errCh <- sentinel + err := waitForShutdown(ctx, errCh) + assert.Equal(t, sentinel, err) +} + +func TestWaitForShutdown_ErrorWinsWhenBothReady(t *testing.T) { + // When cancel() and errCh <- err happen together (the common internal-error + // path), the error must always be returned — never silently swallowed as nil. + sentinel := errors.New("goroutine failure") + for range 100 { + ctx, cancel := context.WithCancel(context.Background()) + errCh := make(chan error, 2) + cancel() + errCh <- sentinel + err := waitForShutdown(ctx, errCh) + require.Equal(t, sentinel, err, "error must not be swallowed when both ctx and errCh are ready") + } +} diff --git a/agent/main.go b/agent/main.go index 636be0831..e44eb88c8 100644 --- a/agent/main.go +++ b/agent/main.go @@ -14,12 +14,23 @@ package main import ( + "context" + "log" + "os" + "os/signal" + "syscall" + "github.com/uber/kraken/agent/cmd" "github.com/uber/kraken/lib/dockerregistry" ) func main() { - cmd.Run(cmd.ParseFlags(), cmd.WithEffect(func() { + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + if err := cmd.Run(ctx, cmd.ParseFlags(), cmd.WithEffect(func() { dockerregistry.RegisterKrakenStorageDriver() - })) + })); err != nil { + log.Fatal(err) + } } diff --git a/utils/log/klog.go b/utils/log/klog.go new file mode 100644 index 000000000..71ed33053 --- /dev/null +++ b/utils/log/klog.go @@ -0,0 +1,100 @@ +// Copyright (c) 2016-2019 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package log + +import ( + "sync/atomic" + + "go.uber.org/zap" +) + +// Level is a verbosity level used with V(). Higher values are more verbose. +type Level int32 + +var _verbosity int32 // atomic; default 0 = disabled + +// SetVerbosity sets the global V() verbosity threshold at runtime. +func SetVerbosity(l Level) { + atomic.StoreInt32(&_verbosity, int32(l)) +} + +// Verbose is returned by V(). Logging methods are no-ops when the requested +// level exceeds the current verbosity threshold. +type Verbose struct { + enabled bool +} + +// Enabled reports whether this verbosity level is active. +func (v Verbose) Enabled() bool { return v.enabled } + +// Info calls the global Info when enabled. +func (v Verbose) Info(args ...interface{}) { + if v.enabled { + Info(args...) + } +} + +// Infof calls the global Infof when enabled. +func (v Verbose) Infof(format string, args ...interface{}) { + if v.enabled { + Infof(format, args...) + } +} + +// InfoS calls the global InfoS when enabled. +func (v Verbose) InfoS(msg string, keysAndValues ...interface{}) { + if v.enabled { + InfoS(msg, keysAndValues...) + } +} + +// V returns a Verbose for the given level, mirroring klog.V. +// +// log.V(2).Info("detailed trace") +// log.V(4).InfoS("very verbose", "key", val) +func V(level Level) Verbose { + return Verbose{enabled: level <= Level(atomic.LoadInt32(&_verbosity))} +} + +// InfoS structured-logs to INFO with key-value pairs, mirroring klog.InfoS. +// +// log.InfoS("pod updated", "pod", name, "namespace", ns) +func InfoS(msg string, keysAndValues ...interface{}) { + zap.S().Infow(msg, keysAndValues...) +} + +// ErrorS structured-logs to ERROR with an error and key-value pairs, +// mirroring klog.ErrorS. +// +// log.ErrorS(err, "dial failed", "host", host, "attempt", n) +func ErrorS(err error, msg string, keysAndValues ...interface{}) { + zap.S().With(zap.Error(err)).Errorw(msg, keysAndValues...) +} + +// WithValues returns a child logger that always emits the given key-value pairs. +// +// reqLog := log.WithValues("traceID", id) +// reqLog.Info("handled") +func WithValues(keysAndValues ...interface{}) *zap.SugaredLogger { + return zap.S().With(keysAndValues...) +} + +// WithName returns a child logger with the given name attached. +// +// sched := log.WithName("scheduler") +// sched.Info("tick") +func WithName(name string) *zap.SugaredLogger { + return zap.S().Named(name) +} diff --git a/utils/log/klog_test.go b/utils/log/klog_test.go new file mode 100644 index 000000000..784e14d97 --- /dev/null +++ b/utils/log/klog_test.go @@ -0,0 +1,148 @@ +// Copyright (c) 2016-2019 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package log + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "go.uber.org/zap" + "go.uber.org/zap/zapcore" + "go.uber.org/zap/zaptest/observer" +) + +// obs wires the package global to an in-memory observer and returns a restore func. +func obs(t *testing.T) (*observer.ObservedLogs, func()) { + t.Helper() + core, logs := observer.New(zapcore.DebugLevel) + SetGlobalLogger(zap.New(core).Sugar()) + return logs, func() { SetGlobalLogger(zap.NewNop().Sugar()) } +} + +// --- InfoS / ErrorS --- + +func TestInfoS(t *testing.T) { + logs, restore := obs(t) + defer restore() + + InfoS("pod updated", "pod", "kubedns", "namespace", "kube-system") + + assert.Equal(t, 1, logs.Len()) + e := logs.All()[0] + assert.Equal(t, "pod updated", e.Message) + assert.Equal(t, "kubedns", e.ContextMap()["pod"]) + assert.Equal(t, "kube-system", e.ContextMap()["namespace"]) +} + +func TestErrorS(t *testing.T) { + logs, restore := obs(t) + defer restore() + + ErrorS(errors.New("refused"), "dial failed", "host", "redis:6379") + + e := logs.All()[0] + assert.Equal(t, zapcore.ErrorLevel, e.Level) + assert.Equal(t, "dial failed", e.Message) + assert.Equal(t, "redis:6379", e.ContextMap()["host"]) + _, hasErr := e.ContextMap()["error"] + assert.True(t, hasErr, "expected 'error' field") +} + +// --- WithValues / WithName --- + +func TestWithValues(t *testing.T) { + logs, restore := obs(t) + defer restore() + + WithValues("traceID", "t1").Info("hit") + + assert.Equal(t, "t1", logs.All()[0].ContextMap()["traceID"]) +} + +func TestWithName(t *testing.T) { + logs, restore := obs(t) + defer restore() + + WithName("scheduler").Info("tick") + + assert.Equal(t, "scheduler", logs.All()[0].LoggerName) +} + +// --- V / Verbose --- + +func TestV_DisabledByDefault(t *testing.T) { + logs, restore := obs(t) + defer restore() + + SetVerbosity(0) + V(1).Info("should be silent") + + assert.Equal(t, 0, logs.Len()) +} + +func TestV_Enabled(t *testing.T) { + logs, restore := obs(t) + defer restore() + + SetVerbosity(3) + defer SetVerbosity(0) + + V(2).Info("within threshold") + V(3).InfoS("at threshold", "k", "v") + V(4).Info("above threshold – silent") + + assert.Equal(t, 2, logs.Len()) +} + +func TestV_Enabled_Check(t *testing.T) { + SetVerbosity(2) + defer SetVerbosity(0) + + assert.True(t, V(1).Enabled()) + assert.True(t, V(2).Enabled()) + assert.False(t, V(3).Enabled()) +} + +// --- FromContext / NewContext --- + +func TestFromContext_FallsBackToGlobal(t *testing.T) { + ctx := context.Background() + assert.NotNil(t, FromContext(ctx), "should return global, not nil") +} + +func TestNewContext_FromContext(t *testing.T) { + logs, restore := obs(t) + defer restore() + + ctx := NewContext(context.Background(), WithValues("requestID", "req-1")) + FromContext(ctx).Infow("handled") + + assert.Equal(t, "req-1", logs.All()[0].ContextMap()["requestID"]) +} + +func TestFromContext_InheritedByChildContext(t *testing.T) { + logs, restore := obs(t) + defer restore() + + parent := NewContext(context.Background(), WithValues("traceID", "t1")) + child, cancel := context.WithCancel(parent) + defer cancel() + + FromContext(child).Info("child log") + + assert.Equal(t, "t1", logs.All()[0].ContextMap()["traceID"]) +} From 6f60d1373f7266493e78d441764e49a81bad58ff Mon Sep 17 00:00:00 2001 From: dajneem23 Date: Mon, 4 May 2026 14:21:09 +0700 Subject: [PATCH 2/2] feat(agent): use context.WithCancel to handle graceful shutdowns --- agent/cmd/cmd.go | 91 ++++++++++++++----------- agent/cmd/cmd_test.go | 69 ++++++++++++------- nginx/nginx.go | 23 ++++++- utils/log/klog.go | 100 ---------------------------- utils/log/klog_test.go | 148 ----------------------------------------- 5 files changed, 121 insertions(+), 310 deletions(-) delete mode 100644 utils/log/klog.go delete mode 100644 utils/log/klog_test.go diff --git a/agent/cmd/cmd.go b/agent/cmd/cmd.go index 67de707b7..80b81f206 100644 --- a/agent/cmd/cmd.go +++ b/agent/cmd/cmd.go @@ -133,11 +133,11 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { config = *overrides.config } else { if err := configutil.Load(flags.ConfigFile, &config); err != nil { - panic(err) + return err } if flags.SecretsFile != "" { if err := configutil.Load(flags.SecretsFile, &config); err != nil { - panic(err) + return err } } } @@ -157,7 +157,7 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { if stats == nil { s, closer, err := metrics.New(config.Metrics, flags.KrakenCluster) if err != nil { - return fmt.Errorf("init metrics: %w", err) + return fmt.Errorf("failed to init metrics: %s", err) } stats = s defer closers.Close(closer) @@ -166,7 +166,7 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { if flags.PeerIP == "" { localIP, err := netutil.GetLocalIP() if err != nil { - return fmt.Errorf("get local IP: %w", err) + return fmt.Errorf("error getting local ip: %s", err) } flags.PeerIP = localIP } @@ -178,40 +178,40 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { pctx, err := core.NewPeerContext( config.PeerIDFactory, flags.Zone, flags.KrakenCluster, flags.PeerIP, flags.PeerPort, false) if err != nil { - return fmt.Errorf("create peer context: %w", err) + return fmt.Errorf("failed to create peer context: %s", err) } cads, err := store.NewCADownloadStore(config.CADownloadStore, stats) if err != nil { - return fmt.Errorf("create CA download store: %w", err) + return fmt.Errorf("failed to create local store: %s", err) } netevents, err := networkevent.NewProducer(config.NetworkEvent) if err != nil { - return fmt.Errorf("create network event producer: %w", err) + return fmt.Errorf("failed to create network event producer: %s", err) } trackers, err := config.Tracker.Build() if err != nil { - return fmt.Errorf("create tracker client: %w", err) + return fmt.Errorf("error building tracker upstream: %s", err) } - go trackers.Monitor(nil) + go trackers.Monitor(ctx.Done()) tls, err := config.TLS.BuildClient() if err != nil { - return fmt.Errorf("build TLS config: %w", err) + return fmt.Errorf("error building client tls config: %s", err) } announceClient := announceclient.New(pctx, trackers, tls) sched, err := scheduler.NewAgentScheduler( config.Scheduler, stats, pctx, cads, netevents, trackers, announceClient, tls) if err != nil { - return fmt.Errorf("create scheduler: %w", err) + return fmt.Errorf("error creating scheduler: %s", err) } buildIndexes, err := config.BuildIndex.Build() if err != nil { - return fmt.Errorf("build build-index upstream: %w", err) + return fmt.Errorf("error building build-index upstream: %s", err) } tagClient := tagclient.NewClusterClient(buildIndexes, tls) @@ -220,7 +220,7 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { registry, err := config.Registry.Build(config.Registry.ReadOnlyParameters(transferer, cads, stats)) if err != nil { - return fmt.Errorf("init registry: %w", err) + return fmt.Errorf("failed to init registry: %s", err) } registryAddr := fmt.Sprintf("127.0.0.1:%d", flags.AgentRegistryPort) @@ -232,14 +232,14 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { } containerRuntimeFactory, err := containerruntime.NewFactory(containerRuntimeCfg, registryAddr) if err != nil { - return fmt.Errorf("create container runtime factory: %w", err) + return fmt.Errorf("failed to create container runtime factory: %s", err) } agentServer := agentserver.New( config.AgentServer, stats, cads, sched, tagClient, announceClient, containerRuntimeFactory) addr := fmt.Sprintf(":%d", flags.AgentServerPort) - log.InfoS("starting agent server", "addr", addr) - errCh := make(chan error, 2) + log.Infof("Starting agent server on %s", addr) + errCh := make(chan error, 3) heartbeatTicker := &timeTicker{inner: time.NewTicker(10 * time.Second)} heartbeatDone := make(chan struct{}) var heartbeatStop sync.Once @@ -252,45 +252,60 @@ func Run(ctx context.Context, flags *Flags, opts ...Option) error { go heartbeat(stats, heartbeatTicker, heartbeatDone) defer stopHeartbeat() + + httpServer := &http.Server{Addr: addr, Handler: agentServer.Handler()} go func() { - if err := http.ListenAndServe(addr, agentServer.Handler()); err != nil { + defer cancel() + // ErrServerClosed is returned by ListenAndServe when Shutdown() is + // called during a clean shutdown — it is expected, not a real error. + if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed { stopHeartbeat() - log.ErrorS(err, "agent server exited") - cancel() + log.Errorf("agent server exited: %s", err) errCh <- err } }() log.Info("Starting registry...") go func() { + defer cancel() if err := registry.ListenAndServe(); err != nil { stopHeartbeat() - log.ErrorS(err, "registry exited") - cancel() + log.Errorf("registry exited: %s", err) errCh <- err } }() - if err := nginx.Run(config.Nginx, map[string]interface{}{ - "allowed_cidrs": config.AllowedCidrs, - "port": flags.AgentRegistryPort, - "registry_server": nginx.GetServer( - config.Registry.Docker.HTTP.Net, config.Registry.Docker.HTTP.Addr), - "agent_server": fmt.Sprintf("127.0.0.1:%d", flags.AgentServerPort), - "registry_backup": config.RegistryBackup}, - nginx.WithTLS(config.TLS)); err != nil { - stopHeartbeat() - log.ErrorS(err, "nginx exited") - cancel() - errCh <- err - } + go func() { + defer cancel() + if err := nginx.RunContext(ctx, config.Nginx, map[string]interface{}{ + "allowed_cidrs": config.AllowedCidrs, + "port": flags.AgentRegistryPort, + "registry_server": nginx.GetServer( + config.Registry.Docker.HTTP.Net, config.Registry.Docker.HTTP.Addr), + "agent_server": fmt.Sprintf("127.0.0.1:%d", flags.AgentServerPort), + "registry_backup": config.RegistryBackup}, + nginx.WithTLS(config.TLS)); err != nil { + stopHeartbeat() + log.Errorf("nginx exited: %s", err) + errCh <- err + } + }() - return waitForShutdown(ctx, errCh) + runErr := waitForShutdown(ctx, errCh) + // Drain in-flight HTTP requests before returning. cancel() is called first + // so Shutdown does not block indefinitely if we are exiting due to an error + // rather than a signal. + cancel() + if err := httpServer.Shutdown(context.Background()); err != nil { + log.Errorf("agent server shutdown: %s", err) + } + return runErr } // waitForShutdown blocks until ctx is cancelled or an error arrives on errCh. -// When both are ready simultaneously (i.e. a goroutine called cancel then sent -// to errCh), the error is returned instead of nil so callers see the root cause. +// Goroutines always send to errCh before calling cancel(), so by the time +// ctx.Done() is observed the error is already buffered and the non-blocking +// drain will always retrieve it. func waitForShutdown(ctx context.Context, errCh <-chan error) error { select { case <-ctx.Done(): @@ -299,7 +314,7 @@ func waitForShutdown(ctx context.Context, errCh <-chan error) error { return err default: } - log.InfoS("shutting down", "reason", ctx.Err()) + log.Infof("shutting down: %s", ctx.Err()) return nil case err := <-errCh: return err diff --git a/agent/cmd/cmd_test.go b/agent/cmd/cmd_test.go index 65076c62a..51b6cc760 100644 --- a/agent/cmd/cmd_test.go +++ b/agent/cmd/cmd_test.go @@ -4,7 +4,8 @@ import ( "context" "errors" "flag" - "fmt" + "net" + "net/http" "os" "path/filepath" "runtime" @@ -149,7 +150,7 @@ func TestRunUsesProvidedConfig(t *testing.T) { assert.True(t, called, "effect should be invoked") } -func TestRunPanicsWhenConfigLoadFails(t *testing.T) { +func TestRunReturnsErrorWhenConfigLoadFails(t *testing.T) { missing := filepath.Join(t.TempDir(), "missing.yaml") flags := &Flags{ @@ -159,19 +160,17 @@ func TestRunPanicsWhenConfigLoadFails(t *testing.T) { ConfigFile: missing, } - expected := fmt.Sprintf("open %s: no such file or directory", missing) - - assert.PanicsWithError(t, expected, func() { - _ = Run( //nolint:errcheck - context.Background(), - flags, - WithMetrics(tally.NewTestScope("", nil)), - WithLogger(zap.NewNop()), - ) - }) + err := Run( + context.Background(), + flags, + WithMetrics(tally.NewTestScope("", nil)), + WithLogger(zap.NewNop()), + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "no such file or directory") } -func TestRunPanicsWhenSecretsLoadFails(t *testing.T) { +func TestRunReturnsErrorWhenSecretsLoadFails(t *testing.T) { _, filename, _, ok := runtime.Caller(0) require.True(t, ok) @@ -189,16 +188,14 @@ func TestRunPanicsWhenSecretsLoadFails(t *testing.T) { SecretsFile: missingSecrets, } - expected := fmt.Sprintf("open %s: no such file or directory", missingSecrets) - - assert.PanicsWithError(t, expected, func() { - _ = Run( //nolint:errcheck - context.Background(), - flags, - WithMetrics(tally.NewTestScope("", nil)), - WithLogger(zap.NewNop()), - ) - }) + err := Run( + context.Background(), + flags, + WithMetrics(tally.NewTestScope("", nil)), + WithLogger(zap.NewNop()), + ) + require.Error(t, err) + assert.Contains(t, err.Error(), "no such file or directory") } func TestValidateRequiredPorts(t *testing.T) { @@ -325,3 +322,29 @@ func TestWaitForShutdown_ErrorWinsWhenBothReady(t *testing.T) { require.Equal(t, sentinel, err, "error must not be swallowed when both ctx and errCh are ready") } } + +// TestHTTPServerGracefulShutdown verifies the pattern used in Run: +// ListenAndServe stops cleanly when Shutdown is called, and http.ErrServerClosed +// is not surfaced as a fatal error. +func TestHTTPServerGracefulShutdown(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + srv := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})} + + serveErr := make(chan error, 1) + go func() { + serveErr <- srv.Serve(ln) + }() + + // Shutdown should cause Serve to return http.ErrServerClosed — not a real error. + require.NoError(t, srv.Shutdown(context.Background())) + + err = <-serveErr + assert.Equal(t, http.ErrServerClosed, err, "Serve should return ErrServerClosed after Shutdown") + + // Verify the ErrServerClosed guard used in Run filters it correctly. + if err != nil && err != http.ErrServerClosed { + t.Fatal("guard would have incorrectly treated ErrServerClosed as a fatal error") + } +} diff --git a/nginx/nginx.go b/nginx/nginx.go index cd77d5047..b3377b3ab 100644 --- a/nginx/nginx.go +++ b/nginx/nginx.go @@ -15,12 +15,14 @@ package nginx import ( "bytes" + "context" "errors" "fmt" "os" "os/exec" "path" "path/filepath" + "syscall" "text/template" "github.com/uber/kraken/nginx/config" @@ -168,6 +170,13 @@ func WithTLS(tls httputil.TLSConfig) Option { // Run injects params into an nginx configuration template and runs it. func Run(config Config, params map[string]interface{}, opts ...Option) error { + return RunContext(context.Background(), config, params, opts...) +} + +// RunContext is like Run but sends SIGQUIT to the nginx process when ctx is +// cancelled, triggering nginx's graceful shutdown (drain in-flight connections) +// rather than an immediate kill. +func RunContext(ctx context.Context, config Config, params map[string]interface{}, opts ...Option) error { if err := config.applyDefaults(); err != nil { return fmt.Errorf("invalid config: %s", err) } @@ -240,7 +249,19 @@ func Run(config Config, params map[string]interface{}, opts ...Option) error { cmd := exec.Command(args[0], args[1:]...) cmd.Stdout = stdout cmd.Stderr = stdout - return cmd.Run() + if err := cmd.Start(); err != nil { + return fmt.Errorf("start nginx: %s", err) + } + + go func() { + <-ctx.Done() + if cmd.Process != nil { + if err := cmd.Process.Signal(syscall.SIGQUIT); err != nil { + log.Errorf("nginx SIGQUIT: %s", err) + } + } + }() + return cmd.Wait() } func populateTemplate(tmpl string, args map[string]interface{}) ([]byte, error) { diff --git a/utils/log/klog.go b/utils/log/klog.go deleted file mode 100644 index 71ed33053..000000000 --- a/utils/log/klog.go +++ /dev/null @@ -1,100 +0,0 @@ -// Copyright (c) 2016-2019 Uber Technologies, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package log - -import ( - "sync/atomic" - - "go.uber.org/zap" -) - -// Level is a verbosity level used with V(). Higher values are more verbose. -type Level int32 - -var _verbosity int32 // atomic; default 0 = disabled - -// SetVerbosity sets the global V() verbosity threshold at runtime. -func SetVerbosity(l Level) { - atomic.StoreInt32(&_verbosity, int32(l)) -} - -// Verbose is returned by V(). Logging methods are no-ops when the requested -// level exceeds the current verbosity threshold. -type Verbose struct { - enabled bool -} - -// Enabled reports whether this verbosity level is active. -func (v Verbose) Enabled() bool { return v.enabled } - -// Info calls the global Info when enabled. -func (v Verbose) Info(args ...interface{}) { - if v.enabled { - Info(args...) - } -} - -// Infof calls the global Infof when enabled. -func (v Verbose) Infof(format string, args ...interface{}) { - if v.enabled { - Infof(format, args...) - } -} - -// InfoS calls the global InfoS when enabled. -func (v Verbose) InfoS(msg string, keysAndValues ...interface{}) { - if v.enabled { - InfoS(msg, keysAndValues...) - } -} - -// V returns a Verbose for the given level, mirroring klog.V. -// -// log.V(2).Info("detailed trace") -// log.V(4).InfoS("very verbose", "key", val) -func V(level Level) Verbose { - return Verbose{enabled: level <= Level(atomic.LoadInt32(&_verbosity))} -} - -// InfoS structured-logs to INFO with key-value pairs, mirroring klog.InfoS. -// -// log.InfoS("pod updated", "pod", name, "namespace", ns) -func InfoS(msg string, keysAndValues ...interface{}) { - zap.S().Infow(msg, keysAndValues...) -} - -// ErrorS structured-logs to ERROR with an error and key-value pairs, -// mirroring klog.ErrorS. -// -// log.ErrorS(err, "dial failed", "host", host, "attempt", n) -func ErrorS(err error, msg string, keysAndValues ...interface{}) { - zap.S().With(zap.Error(err)).Errorw(msg, keysAndValues...) -} - -// WithValues returns a child logger that always emits the given key-value pairs. -// -// reqLog := log.WithValues("traceID", id) -// reqLog.Info("handled") -func WithValues(keysAndValues ...interface{}) *zap.SugaredLogger { - return zap.S().With(keysAndValues...) -} - -// WithName returns a child logger with the given name attached. -// -// sched := log.WithName("scheduler") -// sched.Info("tick") -func WithName(name string) *zap.SugaredLogger { - return zap.S().Named(name) -} diff --git a/utils/log/klog_test.go b/utils/log/klog_test.go deleted file mode 100644 index 784e14d97..000000000 --- a/utils/log/klog_test.go +++ /dev/null @@ -1,148 +0,0 @@ -// Copyright (c) 2016-2019 Uber Technologies, Inc. -// -// Licensed under the Apache License, Version 2.0 (the "License"); -// you may not use this file except in compliance with the License. -// You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, software -// distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -// See the License for the specific language governing permissions and -// limitations under the License. - -package log - -import ( - "context" - "errors" - "testing" - - "github.com/stretchr/testify/assert" - "go.uber.org/zap" - "go.uber.org/zap/zapcore" - "go.uber.org/zap/zaptest/observer" -) - -// obs wires the package global to an in-memory observer and returns a restore func. -func obs(t *testing.T) (*observer.ObservedLogs, func()) { - t.Helper() - core, logs := observer.New(zapcore.DebugLevel) - SetGlobalLogger(zap.New(core).Sugar()) - return logs, func() { SetGlobalLogger(zap.NewNop().Sugar()) } -} - -// --- InfoS / ErrorS --- - -func TestInfoS(t *testing.T) { - logs, restore := obs(t) - defer restore() - - InfoS("pod updated", "pod", "kubedns", "namespace", "kube-system") - - assert.Equal(t, 1, logs.Len()) - e := logs.All()[0] - assert.Equal(t, "pod updated", e.Message) - assert.Equal(t, "kubedns", e.ContextMap()["pod"]) - assert.Equal(t, "kube-system", e.ContextMap()["namespace"]) -} - -func TestErrorS(t *testing.T) { - logs, restore := obs(t) - defer restore() - - ErrorS(errors.New("refused"), "dial failed", "host", "redis:6379") - - e := logs.All()[0] - assert.Equal(t, zapcore.ErrorLevel, e.Level) - assert.Equal(t, "dial failed", e.Message) - assert.Equal(t, "redis:6379", e.ContextMap()["host"]) - _, hasErr := e.ContextMap()["error"] - assert.True(t, hasErr, "expected 'error' field") -} - -// --- WithValues / WithName --- - -func TestWithValues(t *testing.T) { - logs, restore := obs(t) - defer restore() - - WithValues("traceID", "t1").Info("hit") - - assert.Equal(t, "t1", logs.All()[0].ContextMap()["traceID"]) -} - -func TestWithName(t *testing.T) { - logs, restore := obs(t) - defer restore() - - WithName("scheduler").Info("tick") - - assert.Equal(t, "scheduler", logs.All()[0].LoggerName) -} - -// --- V / Verbose --- - -func TestV_DisabledByDefault(t *testing.T) { - logs, restore := obs(t) - defer restore() - - SetVerbosity(0) - V(1).Info("should be silent") - - assert.Equal(t, 0, logs.Len()) -} - -func TestV_Enabled(t *testing.T) { - logs, restore := obs(t) - defer restore() - - SetVerbosity(3) - defer SetVerbosity(0) - - V(2).Info("within threshold") - V(3).InfoS("at threshold", "k", "v") - V(4).Info("above threshold – silent") - - assert.Equal(t, 2, logs.Len()) -} - -func TestV_Enabled_Check(t *testing.T) { - SetVerbosity(2) - defer SetVerbosity(0) - - assert.True(t, V(1).Enabled()) - assert.True(t, V(2).Enabled()) - assert.False(t, V(3).Enabled()) -} - -// --- FromContext / NewContext --- - -func TestFromContext_FallsBackToGlobal(t *testing.T) { - ctx := context.Background() - assert.NotNil(t, FromContext(ctx), "should return global, not nil") -} - -func TestNewContext_FromContext(t *testing.T) { - logs, restore := obs(t) - defer restore() - - ctx := NewContext(context.Background(), WithValues("requestID", "req-1")) - FromContext(ctx).Infow("handled") - - assert.Equal(t, "req-1", logs.All()[0].ContextMap()["requestID"]) -} - -func TestFromContext_InheritedByChildContext(t *testing.T) { - logs, restore := obs(t) - defer restore() - - parent := NewContext(context.Background(), WithValues("traceID", "t1")) - child, cancel := context.WithCancel(parent) - defer cancel() - - FromContext(child).Info("child log") - - assert.Equal(t, "t1", logs.All()[0].ContextMap()["traceID"]) -}