Adds agent-serve-wait-time backend configuration (#4834)

* Add agent serve wait time configuration Add agent listen wait time configuration. Include unit test for middleware rename flag to agent serve wait time use durations all the way through Signed-off-by: Christian Kruse <[email protected]> * reuse apid middlewares Signed-off-by: Christian Kruse <[email protected]> fix tests Signed-off-by: Christian Kruse <[email protected]> * add /ready endpoint to agent listener Signed-off-by: Christian Kruse <[email protected]> * update changelog Signed-off-by: Christian Kruse <[email protected]> Signed-off-by: Christian Kruse <[email protected]>
sensu · Aug 16, 2022 · 00d356d · 00d356d
1 parent 7c477db
commit 00d356d
Show file tree

Hide file tree

Showing 6 changed files with 61 additions and 6 deletions.
diff --git a/CHANGELOG-6.md b/CHANGELOG-6.md
@@ -19,6 +19,10 @@ be marked as global resources.
 for a period of time after startup.
 - Added `/ready` endpoint to the sensu-go API. Returns 200 when the API is ready
 to serve traffic.
+- Added `--agent-serve-wait-time` backend flag to delay accepting agent
+connections for a period of time after startup.
+- Added `/ready` endpoint to the agent listener. Returns 200 when the listener
+is ready to accept agent connections.
 
 ### Fixed
 - Fixed a bug where sensu-backend could crash if the BackendIDGetter

diff --git a/backend/agentd/agentd.go b/backend/agentd/agentd.go
@@ -102,6 +102,8 @@ type Agentd struct {
 	client              *clientv3.Client
 	etcdClientTLSConfig *tls.Config
 	healthRouter        *routers.HealthRouter
+	serveWaitTime       time.Duration
+	ready               func()
 }
 
 // Config configures an Agentd.
@@ -113,6 +115,7 @@ type Config struct {
 	TLS                 *corev2.TLSOptions
 	RingPool            *ringv2.RingPool
 	WriteTimeout        int
+	ServeWaitTime       time.Duration
 	Client              *clientv3.Client
 	EtcdClientTLSConfig *tls.Config
 	Watcher             <-chan store.WatchEventEntityConfig
@@ -142,6 +145,7 @@ func New(c Config, opts ...Option) (*Agentd, error) {
 		watcher:             c.Watcher,
 		client:              c.Client,
 		etcdClientTLSConfig: c.EtcdClientTLSConfig,
+		serveWaitTime:       c.ServeWaitTime,
 	}
 
 	// prepare server TLS config
@@ -150,6 +154,11 @@ func New(c Config, opts ...Option) (*Agentd, error) {
 		return nil, err
 	}
 
+	awaitStart := &middlewares.AwaitStartupMiddleware{
+		ResponseText: "agentd temporarily unavailable during startup",
+	}
+	a.ready = awaitStart.Ready
+
 	// Configure the middlewares used by agentd's HTTP server by assigning them to
 	// public variables so they can be overriden from the enterprise codebase
 	AuthenticationMiddleware = a.AuthenticationMiddleware
@@ -170,7 +179,11 @@ func New(c Config, opts ...Option) (*Agentd, error) {
 
 	route := router.NewRoute().Subrouter()
 	route.HandleFunc("/", a.webSocketHandler)
-	route.Use(agentLimit, authenticate, authorize)
+	route.Use(awaitStart.Then, agentLimit, authenticate, authorize)
+
+	readySubRouter := router.NewRoute().Subrouter()
+	new(routers.ReadyRouter).Mount(readySubRouter)
+	readySubRouter.Use(awaitStart.Then)
 
 	a.httpServer = &http.Server{
 		Addr:         fmt.Sprintf("%s:%d", a.Host, a.Port),
@@ -211,6 +224,23 @@ func (a *Agentd) Start() error {
 		return fmt.Errorf("failed to start agentd: %s", err)
 	}
 
+	if a.serveWaitTime > 0 {
+		logger.Warnf("agentd waiting %s before accepting traffic", a.serveWaitTime)
+		timer := time.After(a.serveWaitTime)
+		go func() {
+			// wait for wait listen time to expire or stop signal
+			select {
+			case <-a.ctx.Done():
+				return
+			case <-timer:
+				a.ready()
+				logger.Warn("agentd now ready to accept traffic")
+			}
+		}()
+	} else {
+		a.ready()
+	}
+
 	a.wg.Add(1)
 
 	go func() {

diff --git a/backend/agentd/agentd_test.go b/backend/agentd/agentd_test.go
@@ -13,6 +13,7 @@ import (
 
 	corev2 "github.com/sensu/sensu-go/api/core/v2"
 	corev3 "github.com/sensu/sensu-go/api/core/v3"
+	"github.com/sensu/sensu-go/backend/apid/middlewares"
 	"github.com/sensu/sensu-go/backend/etcd"
 	"github.com/sensu/sensu-go/backend/store"
 	etcdstore "github.com/sensu/sensu-go/backend/store/etcd"
@@ -34,25 +35,36 @@ func TestAgentdMiddlewares(t *testing.T) {
 		agentName    string
 		username     string
 		group        string
+		isReady      bool
 		storeErr     error
 		expectedCode int
 	}{
 		{
+			description:  "Not ready",
+			isReady:      false,
+			namespace:    "test-rbac",
+			username:     "authorized-user",
+			group:        "group-test-rbac",
+			expectedCode: http.StatusServiceUnavailable,
+		}, {
 			description:  "Authorized request",
 			namespace:    "test-rbac",
 			username:     "authorized-user",
 			group:        "group-test-rbac",
+			isReady:      true,
 			expectedCode: http.StatusOK,
 		}, {
 			description:  "Unauthorized request",
 			namespace:    "super-secret",
 			username:     "unauthorized-user",
+			isReady:      true,
 			expectedCode: http.StatusForbidden,
 		}, {
 			description:  "Invalid user",
 			namespace:    "test-rbac",
 			username:     "nonexistent-user",
 			storeErr:     fmt.Errorf("user not found"),
+			isReady:      true,
 			expectedCode: http.StatusUnauthorized,
 		},
 	}
@@ -98,7 +110,11 @@ func TestAgentdMiddlewares(t *testing.T) {
 				},
 			}}, nil)
 		agentd := &Agentd{store: stor}
-		server := httptest.NewServer(agentd.AuthenticationMiddleware(agentd.AuthorizationMiddleware(testHandler)))
+		readyMiddleware := &middlewares.AwaitStartupMiddleware{}
+		if tc.isReady {
+			readyMiddleware.Ready()
+		}
+		server := httptest.NewServer(readyMiddleware.Then(agentd.AuthenticationMiddleware(agentd.AuthorizationMiddleware(testHandler))))
 		defer server.Close()
 		req, _ := http.NewRequest(http.MethodPost, server.URL, bytes.NewBuffer([]byte{}))
 		req.SetBasicAuth(tc.username, "password")

diff --git a/backend/backend.go b/backend/backend.go
@@ -651,6 +651,7 @@ func Initialize(ctx context.Context, config *Config) (*Backend, error) {
 		TLS:                 config.AgentTLSOptions,
 		RingPool:            b.RingPool,
 		WriteTimeout:        config.AgentWriteTimeout,
+		ServeWaitTime:       config.AgentServeWaitTime,
 		Client:              b.Client,
 		Watcher:             entityConfigWatcher,
 		EtcdClientTLSConfig: b.EtcdClientTLSConfig,

diff --git a/backend/cmd/start.go b/backend/cmd/start.go
@@ -46,6 +46,7 @@ const (
 	flagConfigFile            = "config-file"
 	flagAgentHost             = "agent-host"
 	flagAgentPort             = "agent-port"
+	flagAgentServeWaitTime    = "agent-serve-wait-time"
 	flagAPIListenAddress      = "api-listen-address"
 	flagAPIRequestLimit       = "api-request-limit"
 	flagAPIURL                = "api-url"
@@ -224,6 +225,7 @@ func StartCommand(initialize InitializeFunc) *cobra.Command {
 				AgentHost:             viper.GetString(flagAgentHost),
 				AgentPort:             viper.GetInt(flagAgentPort),
 				AgentWriteTimeout:     viper.GetInt(backend.FlagAgentWriteTimeout),
+				AgentServeWaitTime:    viper.GetDuration(flagAgentServeWaitTime),
 				APIListenAddress:      viper.GetString(flagAPIListenAddress),
 				APIRequestLimit:       viper.GetInt64(flagAPIRequestLimit),
 				APIURL:                viper.GetString(flagAPIURL),
@@ -526,6 +528,7 @@ func flagSet(server bool) *pflag.FlagSet {
 		// Main Flags
 		flagSet.String(flagAgentHost, viper.GetString(flagAgentHost), "agent listener host")
 		flagSet.Int(flagAgentPort, viper.GetInt(flagAgentPort), "agent listener port")
+		flagSet.Duration(flagAgentServeWaitTime, viper.GetDuration(flagAgentServeWaitTime), "wait time before accepting agent connections on startup")
 		flagSet.String(flagAPIListenAddress, viper.GetString(flagAPIListenAddress), "address to listen on for api traffic")
 		flagSet.Int64(flagAPIRequestLimit, viper.GetInt64(flagAPIRequestLimit), "maximum API request body size, in bytes")
 		flagSet.String(flagAPIURL, viper.GetString(flagAPIURL), "url of the api to connect to")

diff --git a/backend/config.go b/backend/config.go
@@ -51,10 +51,11 @@ type Config struct {
 	CacheDir string
 
 	// Agentd Configuration
-	AgentHost         string
-	AgentPort         int
-	AgentTLSOptions   *corev2.TLSOptions
-	AgentWriteTimeout int
+	AgentHost          string
+	AgentPort          int
+	AgentTLSOptions    *corev2.TLSOptions
+	AgentWriteTimeout  int
+	AgentServeWaitTime time.Duration
 
 	// Apid Configuration
 	APIListenAddress string