Merge pull request #3370 from nosammai/disconnect-after-uptime

Add disconnect-after-uptime flag to set a max lifetime for agents

Author: moskyb

agent/agent_configuration.go

@@ -49,6 +49,7 @@ type AgentConfiguration struct {
 	HealthCheckAddr              string
 	DisconnectAfterJob           bool
 	DisconnectAfterIdleTimeout   int
+	DisconnectAfterUptime        int
 	CancelGracePeriod            int
 	SignalGracePeriod            time.Duration
 	EnableJobLogTmpfile          bool

agent/agent_worker.go

@@ -106,6 +106,9 @@ type AgentWorker struct {
 	state        agentWorkerState
 	currentJobID string
 
+	// The time when this agent worker started
+	startTime time.Time
+
 	// disable the delay between pings, to speed up certain testing scenarios
 	noWaitBetweenPingsForTesting bool
 }
@@ -213,6 +216,9 @@ func (a *AgentWorker) statusCallback(context.Context) (any, error) {
 
 // Starts the agent worker
 func (a *AgentWorker) Start(ctx context.Context, idleMonitor *IdleMonitor) error {
+	// Record the start time for max agent lifetime tracking
+	a.startTime = time.Now()
+
 	a.metrics = a.metricsCollector.Scope(metrics.Tags{
 		"agent_name": a.agent.Name,
 	})
@@ -331,6 +337,7 @@ func (a *AgentWorker) runPingLoop(ctx context.Context, idleMonitor *IdleMonitor)
 	//   ping action isn't "pause",
 	// * the agent is in disconnect-after-idle-timeout mode, has been idle for
 	//   longer than the idle timeout, and the ping action isn't "pause".
+	// * the agent has exceeded its disconnect-after-uptime and the ping action isn't "pause".
 	for {
 		setStat("😴 Waiting until next ping interval tick")
 		select {
@@ -411,6 +418,18 @@ func (a *AgentWorker) runPingLoop(ctx context.Context, idleMonitor *IdleMonitor)
 			return nil
 		}
 
+		// Exit after disconnect-after-uptime is exceeded.
+		if a.agentConfiguration.DisconnectAfterUptime > 0 {
+			maxUptime := time.Second * time.Duration(a.agentConfiguration.DisconnectAfterUptime)
+			if time.Since(a.startTime) >= maxUptime {
+				if job != nil {
+					a.logger.Error("Agent ping dispatched a job (id %q) but agent has exceeded max uptime of %v!", job.ID, maxUptime)
+				}
+				a.logger.Info("Agent has exceeded max uptime of %v. Disconnecting...", maxUptime)
+				return nil
+			}
+		}
+
 		// Note that Ping only returns a job if err == nil.
 		if job == nil {
 			if disconnectAfterIdleTimeout == 0 {

agent/agent_worker_test.go

@@ -487,6 +487,92 @@ func TestAgentWorker_DisconnectAfterJob_Start_Pause_Unpause(t *testing.T) {
 	}
 }
 
+func TestAgentWorker_DisconnectAfterUptime(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	server := NewFakeAPIServer()
+	defer server.Close()
+
+	// Create a job that the agent could potentially accept
+	job := server.AddJob(map[string]string{
+		"BUILDKITE_COMMAND": "echo hello",
+	})
+
+	// Pre-register the agent.
+	const agentSessionToken = "alpacas"
+	agent := server.AddAgent(agentSessionToken)
+
+	pingCount := 0
+	agent.PingHandler = func(*http.Request) (api.Ping, error) {
+		pingCount++
+		// Always offer the job to test that the agent stops accepting jobs after max lifetime
+		return api.Ping{
+			Job: job.Job,
+		}, nil
+	}
+
+	server.Assign(agent, job)
+
+	apiClient := api.NewClient(logger.Discard, api.Config{
+		Endpoint: server.URL,
+		Token:    "llamas",
+	})
+
+	l := logger.NewConsoleLogger(logger.NewTestPrinter(t), func(int) {})
+
+	worker := NewAgentWorker(
+		l,
+		&api.AgentRegisterResponse{
+			UUID:              uuid.New().String(),
+			Name:              "agent-1",
+			AccessToken:       agentSessionToken,
+			Endpoint:          server.URL,
+			PingInterval:      1,
+			JobStatusInterval: 1,
+			HeartbeatInterval: 10,
+		},
+		metrics.NewCollector(logger.Discard, metrics.CollectorConfig{}),
+		apiClient,
+		AgentWorkerConfig{
+			SpawnIndex: 1,
+			AgentConfiguration: AgentConfiguration{
+				BootstrapScript:       "./dummy_bootstrap.sh",
+				BuildPath:             filepath.Join(os.TempDir(), t.Name(), "build"),
+				HooksPath:             filepath.Join(os.TempDir(), t.Name(), "hooks"),
+				DisconnectAfterUptime: 1, // 1 second max uptime
+			},
+		},
+	)
+	worker.noWaitBetweenPingsForTesting = true
+
+	idleMonitor := NewIdleMonitor(1)
+
+	// Record start time
+	startTime := time.Now()
+
+	if err := worker.Start(ctx, idleMonitor); err != nil {
+		t.Errorf("worker.Start() = %v", err)
+	}
+
+	// Check that the agent disconnected after approximately 1 second
+	elapsed := time.Since(startTime)
+	if elapsed < 900*time.Millisecond || elapsed > 2*time.Second {
+		t.Errorf("Agent should have disconnected after ~1 second, but took %v", elapsed)
+	}
+
+	// The agent should have made at least one ping before disconnecting
+	if pingCount == 0 {
+		t.Error("Agent should have made at least one ping before disconnecting")
+	}
+
+	// The agent should have made at least one ping and should have disconnected
+	// due to max uptime being exceeded. The important thing is that the agent
+	// disconnected properly with the uptime check, which we verified above.
+}
+
 func TestAgentWorker_SetEndpointDuringRegistration(t *testing.T) {
 	// The registration request is made in clicommand.AgentStartCommand, and the response
 	// is passed into agent.NewAgentWorker(...), so we'll just test the response handling.

clicommand/agent_start.go

@@ -104,6 +104,7 @@ type AgentStartConfig struct {
 	AcquireJob                 string `cli:"acquire-job"`
 	DisconnectAfterJob         bool   `cli:"disconnect-after-job"`
 	DisconnectAfterIdleTimeout int    `cli:"disconnect-after-idle-timeout"`
+	DisconnectAfterUptime      int    `cli:"disconnect-after-uptime"`
 	CancelGracePeriod          int    `cli:"cancel-grace-period"`
 	SignalGracePeriodSeconds   int    `cli:"signal-grace-period-seconds"`
 
@@ -235,6 +236,10 @@ func (asc AgentStartConfig) Features(ctx context.Context) []string {
 		features = append(features, "disconnect-after-idle")
 	}
 
+	if asc.DisconnectAfterUptime != 0 {
+		features = append(features, "disconnect-after-uptime")
+	}
+
 	if asc.NoPlugins {
 		features = append(features, "no-plugins")
 	}
@@ -363,6 +368,12 @@ var AgentStartCommand = cli.Command{
 			Usage:  "The maximum idle time in seconds to wait for a job before disconnecting. The default of 0 means no timeout",
 			EnvVar: "BUILDKITE_AGENT_DISCONNECT_AFTER_IDLE_TIMEOUT",
 		},
+		cli.IntFlag{
+			Name:   "disconnect-after-uptime",
+			Value:  0,
+			Usage:  "The maximum uptime in seconds before the agent stops accepting new jobs and shuts down after any running jobs complete. The default of 0 means no timeout",
+			EnvVar: "BUILDKITE_AGENT_DISCONNECT_AFTER_UPTIME",
+		},
 		cancelGracePeriodFlag,
 		cli.BoolFlag{
 			Name:   "enable-job-log-tmpfile",
@@ -1014,6 +1025,7 @@ var AgentStartCommand = cli.Command{
 			TimestampLines:               cfg.TimestampLines,
 			DisconnectAfterJob:           cfg.DisconnectAfterJob,
 			DisconnectAfterIdleTimeout:   cfg.DisconnectAfterIdleTimeout,
+			DisconnectAfterUptime:        cfg.DisconnectAfterUptime,
 			CancelGracePeriod:            cfg.CancelGracePeriod,
 			SignalGracePeriod:            signalGracePeriod,
 			EnableJobLogTmpfile:          cfg.EnableJobLogTmpfile,
@@ -1109,6 +1121,10 @@ var AgentStartCommand = cli.Command{
 			l.Info("Agents will disconnect after %d seconds of inactivity", agentConf.DisconnectAfterIdleTimeout)
 		}
 
+		if agentConf.DisconnectAfterUptime > 0 {
+			l.Info("Agents will disconnect after %d seconds of uptime and shut down after any running jobs complete", agentConf.DisconnectAfterUptime)
+		}
+
 		if len(cfg.AllowedRepositories) > 0 {
 			agentConf.AllowedRepositories = make([]*regexp.Regexp, 0, len(cfg.AllowedRepositories))
 			for _, v := range cfg.AllowedRepositories {