better restart button

Author: snadrus

cmd/curio/rpc/rpc.go

@@ -98,10 +98,17 @@ func CurioHandler(
 	return ah
 }
 
+// TaskEngineRestarter begins an in-memory graceful restart: it cordons the
+// node (without persisting to the DB), drains running tasks, then restarts.
+type TaskEngineRestarter interface {
+	RequestRestart()
+}
+
 type CurioAPI struct {
 	*deps.Deps
 	paths.SectorIndex
 	ShutdownChan chan struct{}
+	restarter    TaskEngineRestarter
 }
 
 func (p *CurioAPI) Version(context.Context) ([]int, error) {
@@ -226,8 +233,14 @@ func (p *CurioAPI) Shutdown(context.Context) error {
 	return gracehttpsvc.TriggerShutdown()
 }
 
-// Restart triggers a zero-downtime process restart via SIGUSR2.
+// Restart initiates a graceful restart. When a task engine is wired in, it
+// cordons the node in-memory and drains running tasks before triggering the
+// zero-downtime restart; otherwise it triggers the restart immediately.
 func (p *CurioAPI) Restart(context.Context) error {
+	if p.restarter != nil {
+		p.restarter.RequestRestart()
+		return nil
+	}
 	return gracehttpsvc.TriggerRestart()
 }
 
@@ -457,7 +470,7 @@ func (p *CurioAPI) IndexSamples(ctx context.Context, pcid cid.Cid) ([]multihash.
 	return p.IndexStore.GetPieceHashRange(ctx, pcid, firstHash, chunk.NumberOfBlocks, true)
 }
 
-func ListenAndServe(ctx context.Context, dependencies *deps.Deps, shutdownChan chan struct{}, extraServers ...*http.Server) error {
+func ListenAndServe(ctx context.Context, dependencies *deps.Deps, shutdownChan chan struct{}, restarter TaskEngineRestarter, extraServers ...*http.Server) error {
 	fh := &paths.FetchHandler{Local: dependencies.LocalStore, PfHandler: &paths.DefaultPartialFileHandler{}}
 	remoteHandler := func(w http.ResponseWriter, r *http.Request) {
 		if !auth.HasPerm(r.Context(), nil, lapi.PermAdmin) {
@@ -489,7 +502,7 @@ func ListenAndServe(ctx context.Context, dependencies *deps.Deps, shutdownChan c
 		Handler: CurioHandler(
 			authVerify,
 			remoteHandler,
-			&CurioAPI{dependencies, dependencies.Si, shutdownChan},
+			&CurioAPI{dependencies, dependencies.Si, shutdownChan, restarter},
 			prometheusServiceDiscovery(ctx, dependencies),
 			permissioned),
 		ReadHeaderTimeout: time.Minute * 3,

cmd/curio/run.go

@@ -168,7 +168,7 @@ var runCmd = &cli.Command{
 			extraServers = append(extraServers, marketServer)
 		}
 
-		err = rpc.ListenAndServe(ctx, dependencies, shutdownChan, extraServers...)
+		err = rpc.ListenAndServe(ctx, dependencies, shutdownChan, taskEngine, extraServers...)
 		if err != nil {
 			return err
 		}

harmony/harmonytask/harmonytask.go

@@ -146,6 +146,11 @@ type TaskEngine struct {
 	// runtime flags
 	yieldBackground atomic.Bool
 
+	// restartRequested is an in-memory cordon set when a graceful restart is
+	// requested via RPC. It is intentionally not persisted to the DB so that
+	// the node returns to its previous cordon state after restarting.
+	restartRequested atomic.Bool
+
 	// synchronous to the single-threaded poller
 	lastFollowTime time.Time
 	lastCleanup    atomic.Value
@@ -289,6 +294,11 @@ func (e *TaskEngine) poller() {
 			continue
 		}
 
+		// An in-memory restart request cordons the node without touching the DB.
+		if e.restartRequested.Load() {
+			schedulable = false
+		}
+
 		e.yieldBackground.Store(!schedulable)
 
 		accepted := e.pollerTryAllWork(schedulable)
@@ -519,28 +529,86 @@ func (e *TaskEngine) checkNodeFlags() (bool, error) {
 	return !unschedulable, nil
 }
 
-func (e *TaskEngine) restartIfNoTasksPending(pendingSince time.Time) {
-	var tasksPending int
-	err := e.db.QueryRow(e.ctx, `SELECT COUNT(*) FROM harmony_task WHERE owner_id=$1`, e.ownerID).Scan(&tasksPending)
-	if err != nil {
-		log.Error("Unable to check for tasks pending: ", err)
+func (e *TaskEngine) activeTaskCount() int {
+	count := 0
+	for _, h := range e.handlers {
+		count += h.Max.ActiveThis()
+	}
+	return count
+}
+
+// RequestRestart begins an in-memory graceful restart. The node immediately
+// stops accepting new tasks (the same effect as cordoning), waits for the
+// currently-running tasks to drain, then triggers a zero-downtime restart.
+//
+// The cordon is kept in-memory only and is never written to harmony_machines,
+// so after the restart the node returns to whatever cordon state was persisted
+// in the database before the restart was requested.
+func (e *TaskEngine) RequestRestart() {
+	if !e.restartRequested.CompareAndSwap(false, true) {
+		// a restart is already in progress
 		return
 	}
-	if tasksPending == 0 {
-		log.Infow("no tasks pending, restarting", "ownerID", e.ownerID, "pendingSince", pendingSince, "took", time.Since(pendingSince))
+	log.Infow("restart requested; cordoning in-memory and draining tasks", "ownerID", e.ownerID)
+	go e.drainAndRestart()
+}
 
-		// unset the flags first
-		_, err = e.db.Exec(e.ctx, `UPDATE harmony_machines SET restart_request=NULL, unschedulable=FALSE WHERE host_and_port=$1`, e.hostAndPort)
-		if err != nil {
-			log.Error("Unable to unset restart request: ", err)
+func (e *TaskEngine) drainAndRestart() {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-e.ctx.Done():
 			return
+		case <-ticker.C:
+		}
+
+		if active := e.activeTaskCount(); active > 0 {
+			log.Infow("restart waiting for tasks to drain", "ownerID", e.ownerID, "activeTasks ", active)
+			continue
 		}
 
-		// zero-downtime restart via gracehttp; fall back to exit 100 for systemd
+		log.Infow("no tasks running, triggering graceful restart", "ownerID", e.ownerID)
 		if err := gracehttpsvc.TriggerRestart(); err != nil {
 			log.Errorw("graceful restart failed, falling back to exit", "error", err)
 			os.Exit(ExitStatusRestartRequest)
 		}
+		return
+	}
+}
+
+func (e *TaskEngine) restartIfNoTasksPending(pendingSince time.Time) {
+	var tasksPending int
+	err := e.db.QueryRow(e.ctx, `SELECT COUNT(*) FROM harmony_task WHERE owner_id=$1`, e.ownerID).Scan(&tasksPending)
+	if err != nil {
+		log.Error("Unable to check for tasks pending: ", err)
+		return
+	}
+
+	activeTasks := e.activeTaskCount()
+	if tasksPending > 0 || activeTasks > 0 {
+		log.Infow("restart waiting for tasks to finish",
+			"ownerID", e.ownerID,
+			"pendingSince", pendingSince,
+			"tasksPending", tasksPending,
+			"activeTasks", activeTasks)
+		return
+	}
+
+	log.Infow("no tasks pending, restarting", "ownerID", e.ownerID, "pendingSince", pendingSince, "took", time.Since(pendingSince))
+
+	// Clear restart_request only; stay cordoned until the operator uncordons after restart.
+	_, err = e.db.Exec(e.ctx, `UPDATE harmony_machines SET restart_request=NULL WHERE host_and_port=$1`, e.hostAndPort)
+	if err != nil {
+		log.Error("Unable to unset restart request: ", err)
+		return
+	}
+
+	// zero-downtime restart via gracehttp; fall back to exit 100 for systemd
+	if err := gracehttpsvc.TriggerRestart(); err != nil {
+		log.Errorw("graceful restart failed, falling back to exit", "error", err)
+		os.Exit(ExitStatusRestartRequest)
 	}
 }
 

itests/grace_restart_test.go

@@ -71,7 +71,7 @@ func TestGraceRestartDuringPDPPing(t *testing.T) {
 	}()
 
 	require.Eventually(t, func() bool {
-		return pingSuccesses.Load() >= 20
+		return pingSuccesses.Load() >= 5
 	}, 30*time.Second, 25*time.Millisecond, "expected initial successful PDP pings before restart")
 
 	require.NoError(t, gracehttpsvc.RestartFromPIDFile(pidPath))
@@ -85,7 +85,7 @@ func TestGraceRestartDuringPDPPing(t *testing.T) {
 	}, 2*time.Minute, 100*time.Millisecond, "expected curio.pid to reflect a new process after restart")
 
 	require.Eventually(t, func() bool {
-		return pingSuccesses.Load() >= 40
+		return pingSuccesses.Load() >= 10
 	}, 2*time.Minute, 25*time.Millisecond, "expected continued successful PDP pings through restart")
 
 	close(stopPinger)

itests/helpers/harness.go

@@ -100,7 +100,7 @@ func StartCurioHarness(
 		if marketServer != nil {
 			extra = append(extra, marketServer)
 		}
-		err := rpc.ListenAndServe(ctx, dependencies, shutdownChan, extra...)
+		err := rpc.ListenAndServe(ctx, dependencies, shutdownChan, taskEngine, extra...)
 		if err != nil {
 			t.Errorf("failed to start the Curio RPC server: %v", err)
 		}