From 77db55bb1416890ae4117ca9bd2d33fed04dab4e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Fri, 23 Jun 2023 11:38:37 +0200
Subject: [PATCH] Manager: when worker signs off, only remember specific
 statuses

Limit which worker statuses are remembered (when they go offline) to
those that we want to restore when they come back online. This is now
set to `awake` and `asleep`. This prevents workers from being told to go
to states that they cannot handle, such as `error` or `starting`.
---
 internal/manager/api_impl/workers.go      | 15 +++++--
 internal/manager/api_impl/workers_test.go | 54 ++++++++++++++++++-----
 2 files changed, 56 insertions(+), 13 deletions(-)

diff --git a/internal/manager/api_impl/workers.go b/internal/manager/api_impl/workers.go
index 6626992d..47cb73e5 100644
--- a/internal/manager/api_impl/workers.go
+++ b/internal/manager/api_impl/workers.go
@@ -21,6 +21,15 @@ import (
 	"git.blender.org/flamenco/pkg/api"
 )
 
+// rememberableWorkerStates contains those worker statuses that should be
+// remembered when the worker signs off, so that it'll be sent to that state
+// again next time it signs on. Not every state has to be remembered like this;
+// 'error' and 'starting' are not states to send the worker into.
+var rememberableWorkerStates = map[api.WorkerStatus]bool{
+	api.WorkerStatusAsleep: true,
+	api.WorkerStatusAwake:  true,
+}
+
 // RegisterWorker registers a new worker and stores it in the database.
 func (f *Flamenco) RegisterWorker(e echo.Context) error {
 	logger := requestLogger(e)
@@ -163,10 +172,10 @@ func (f *Flamenco) SignOff(e echo.Context) error {
 		w.StatusChangeClear()
 	}
 
-	// Remember the previous status if an initial status exists
-	if w.StatusRequested == "" {
+	// Remember the previous status if an initial status exists.
+	if w.StatusRequested == "" && rememberableWorkerStates[prevStatus] {
 		w.StatusChangeRequest(prevStatus, false)
-  	}
+	}
 
 	// Pass a generic background context, as these changes should be stored even
 	// when the HTTP connection is aborted.
diff --git a/internal/manager/api_impl/workers_test.go b/internal/manager/api_impl/workers_test.go
index 439fba41..0e8b2944 100644
--- a/internal/manager/api_impl/workers_test.go
+++ b/internal/manager/api_impl/workers_test.go
@@ -244,12 +244,12 @@ func TestWorkerSignoffTaskRequeue(t *testing.T) {
 		Name:           worker.Name,
 		PreviousStatus: &prevStatus,
 		Status:         api.WorkerStatusOffline,
-		StatusChange:	&api.WorkerStatusChangeRequest{
-			IsLazy:		false,
-			Status:		api.WorkerStatusAwake,
+		StatusChange: &api.WorkerStatusChangeRequest{
+			IsLazy: false,
+			Status: api.WorkerStatusAwake,
 		},
-		Updated:        worker.UpdatedAt,
-		Version:        worker.Software,
+		Updated: worker.UpdatedAt,
+		Version: worker.Software,
 	})
 
 	err := mf.flamenco.SignOff(echo)
@@ -273,12 +273,12 @@ func TestWorkerRememberPreviousStatus(t *testing.T) {
 		Name:           worker.Name,
 		PreviousStatus: ptr(api.WorkerStatusAwake),
 		Status:         api.WorkerStatusOffline,
-		StatusChange:	&api.WorkerStatusChangeRequest{
-			IsLazy:		false,
-			Status:		api.WorkerStatusAwake,
+		StatusChange: &api.WorkerStatusChangeRequest{
+			IsLazy: false,
+			Status: api.WorkerStatusAwake,
 		},
-		Updated:        worker.UpdatedAt,
-		Version:        worker.Software,
+		Updated: worker.UpdatedAt,
+		Version: worker.Software,
 	})
 
 	savedWorker := worker
@@ -299,6 +299,40 @@ func TestWorkerRememberPreviousStatus(t *testing.T) {
 	assert.Equal(t, false, worker.LazyStatusRequest)
 }
 
+// Test that certain statuses (like 'starting') are not remembered.
+func TestWorkerDontRememberPreviousStatus(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+	mf := newMockedFlamenco(mockCtrl)
+
+	worker := testWorker()
+	worker.Status = api.WorkerStatusError
+	worker.StatusChangeRequest(api.WorkerStatusOffline, true)
+
+	mf.broadcaster.EXPECT().BroadcastWorkerUpdate(api.SocketIOWorkerUpdate{
+		Id:             worker.UUID,
+		Name:           worker.Name,
+		PreviousStatus: ptr(api.WorkerStatusError),
+		Status:         api.WorkerStatusOffline,
+		StatusChange:   nil,
+		Updated:        worker.UpdatedAt,
+		Version:        worker.Software,
+	})
+
+	savedWorker := worker
+	savedWorker.Status = api.WorkerStatusOffline
+	savedWorker.StatusChangeClear() // No status change should be queued.
+	mf.persistence.EXPECT().SaveWorkerStatus(gomock.Any(), &savedWorker).Return(nil)
+	mf.stateMachine.EXPECT().RequeueActiveTasksOfWorker(gomock.Any(), &worker, "worker signed off").Return(nil)
+	mf.persistence.EXPECT().WorkerSeen(gomock.Any(), &worker)
+
+	echo := mf.prepareMockedRequest(nil)
+	requestWorkerStore(echo, &worker)
+	err := mf.flamenco.SignOff(echo)
+	assert.NoError(t, err)
+	assertResponseNoContent(t, echo)
+}
+
 func TestWorkerState(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()