From a1553999bb8bfc57d9962ca3ee3ed992640b8f2b Mon Sep 17 00:00:00 2001
From: Rob Prentiss <rob@prentiss.name>
Date: Mon, 23 Jan 2023 22:02:45 -0800
Subject: [PATCH] Make shutdown timeout configurable (#1806)

Currently, when shutting down the server (via SIGINT or SIGTERM), the
shutdown closes any open connections after only 10 seconds (via a
context.WithTimeout). This does not provie a lot of time for longer
operations, such as listing versions, or downloading a larger module zip
file.

When running in Kubernetes, and scaling instances or changing config,
this causes a lot of dropped connections and gateway errors. 10 seconds
is arguably much too short, and should be configurable.

This commit increases that default to 60 seconds, and adds a config
variable to allow users to specify their desired timeout.
---
 cmd/proxy/main.go         | 2 +-
 config.dev.toml           | 6 ++++++
 pkg/config/config.go      | 2 ++
 pkg/config/config_test.go | 1 +
 4 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go
index 751932e7..6c016c63 100644
--- a/cmd/proxy/main.go
+++ b/cmd/proxy/main.go
@@ -55,7 +55,7 @@ func main() {
 		<-sigint
 
 		// We received an interrupt signal, shut down.
-		ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*time.Duration(conf.ShutdownTimeout))
 		defer cancel()
 		if err := srv.Shutdown(ctx); err != nil {
 			log.Fatal(err)
diff --git a/config.dev.toml b/config.dev.toml
index e0089c8d..7f8acc9d 100755
--- a/config.dev.toml
+++ b/config.dev.toml
@@ -305,6 +305,12 @@ SingleFlightType = "memory"
 # Env override: ATHENS_INDEX_TYPE
 IndexType = "none"
 
+# ShutdownTimeout sets the timeout (in seconds) for open connections when shutting down
+# (via SIGINT or SIGTERM). Connections still open after the timeout will be dropped.
+# Defaults to 60
+# Env override: ATHENS_SHUTDOWN_TIMEOUT
+ShutdownTimeout = 60
+
 [SingleFlight]
     [SingleFlight.Etcd]
         # Endpoints are comma separated URLs that determine all distributed etcd servers.
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 74dc1a83..ca2a1b0e 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -56,6 +56,7 @@ type Config struct {
 	SingleFlightType string    `envconfig:"ATHENS_SINGLE_FLIGHT_TYPE"`
 	RobotsFile       string    `envconfig:"ATHENS_ROBOTS_FILE"`
 	IndexType        string    `envconfig:"ATHENS_INDEX_TYPE"`
+	ShutdownTimeout  int       `validate:"min=0" envconfig:"ATHENS_SHUTDOWN_TIMEOUT"`
 	SingleFlight     *SingleFlight
 	Storage          *Storage
 	Index            *Index
@@ -165,6 +166,7 @@ func defaultConfig() *Config {
 		NetworkMode:      "strict",
 		RobotsFile:       "robots.txt",
 		IndexType:        "none",
+		ShutdownTimeout:  60,
 		SingleFlight: &SingleFlight{
 			Etcd:  &Etcd{"localhost:2379,localhost:22379,localhost:32379"},
 			Redis: &Redis{"127.0.0.1:6379", "", DefaultRedisLockConfig()},
diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go
index f45084be..a74c3a36 100644
--- a/pkg/config/config_test.go
+++ b/pkg/config/config_test.go
@@ -314,6 +314,7 @@ func TestParseExampleConfig(t *testing.T) {
 		DownloadMode:     "sync",
 		RobotsFile:       "robots.txt",
 		IndexType:        "none",
+		ShutdownTimeout:  60,
 		Index:            &Index{},
 	}