From be0b10400f44f813c293f3aca8a668d905fca44a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@blender.org>
Date: Thu, 16 Jun 2022 10:39:42 +0200
Subject: [PATCH] Manager: count workers as 'seen' even when there is no task

Fix a bug where a worker would only be counted as 'seen' by the task
scheduler if it actually got a task assigned.
---
 internal/manager/api_impl/workers.go      | 13 +++++++++----
 internal/manager/api_impl/workers_test.go | 23 +++++++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/internal/manager/api_impl/workers.go b/internal/manager/api_impl/workers.go
index cd62f134..e9cef72a 100644
--- a/internal/manager/api_impl/workers.go
+++ b/internal/manager/api_impl/workers.go
@@ -257,11 +257,20 @@ func (f *Flamenco) WorkerStateChanged(e echo.Context) error {
 func (f *Flamenco) ScheduleTask(e echo.Context) error {
 	logger := requestLogger(e)
 	worker := requestWorkerOrPanic(e)
+	ctx := e.Request().Context()
 	logger.Debug().Msg("worker requesting task")
 
 	f.taskSchedulerMutex.Lock()
 	defer f.taskSchedulerMutex.Unlock()
 
+	// The worker is actively asking for a task, so note that it was seen
+	// regardless of any failures below, or whether there actually is a task to
+	// run.
+	if err := f.workerSeen(ctx, logger, worker); err != nil {
+		return sendAPIError(e, http.StatusInternalServerError,
+			"error storing worker 'last seen' timestamp in database")
+	}
+
 	// Check that this worker is actually allowed to do work.
 	requiredStatusToGetTask := api.WorkerStatusAwake
 	if worker.Status != api.WorkerStatusAwake {
@@ -303,13 +312,9 @@ func (f *Flamenco) ScheduleTask(e echo.Context) error {
 	}
 
 	// Start timeout measurement as soon as the Worker gets the task assigned.
-	ctx := e.Request().Context()
 	if err := f.workerPingedTask(ctx, logger, dbTask); err != nil {
 		return sendAPIError(e, http.StatusInternalServerError, "internal error updating task for timeout calculation: %v", err)
 	}
-	if err := f.workerSeen(ctx, logger, worker); err != nil {
-		return sendAPIError(e, http.StatusInternalServerError, "error storing worker 'last seen' timestamp in database")
-	}
 
 	// Convert database objects to API objects:
 	apiCommands := []api.Command{}
diff --git a/internal/manager/api_impl/workers_test.go b/internal/manager/api_impl/workers_test.go
index 409fe564..ecab9fe0 100644
--- a/internal/manager/api_impl/workers_test.go
+++ b/internal/manager/api_impl/workers_test.go
@@ -57,6 +57,29 @@ func TestTaskScheduleHappy(t *testing.T) {
 	assert.Equal(t, http.StatusOK, resp.StatusCode)
 }
 
+func TestTaskScheduleNoTaskAvailable(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	mf := newMockedFlamenco(mockCtrl)
+	worker := testWorker()
+
+	echo := mf.prepareMockedRequest(nil)
+	requestWorkerStore(echo, &worker)
+
+	// Expect a call into the persistence layer, which should return nil.
+	ctx := echo.Request().Context()
+	mf.persistence.EXPECT().ScheduleTask(ctx, &worker).Return(nil, nil)
+
+	// This call should still trigger a "worker seen" call, as the worker is
+	// actively asking for tasks.
+	mf.persistence.EXPECT().WorkerSeen(ctx, &worker)
+
+	err := mf.flamenco.ScheduleTask(echo)
+	assert.NoError(t, err)
+	assertResponseEmpty(t, echo)
+}
+
 func TestTaskScheduleNonActiveStatus(t *testing.T) {
 	mockCtrl := gomock.NewController(t)
 	defer mockCtrl.Finish()