fix: prevent stuck runner jobs when FetchTask performance is slower than client-side timeout (#13658)

Forgejo has recently experienced a significant increase in the number of Action jobs which appear stuck in "Set up job" and never progress.  This was supposed to be prevented by #11401.  This commit adds app-level locking on the runner request key, and when a concurrent request is received to send an error back to the client to allow it to retry.

Reviewed-on: https://codeberg.org/forgejo/forgejo/pulls/13658
Reviewed-by: Andreas Ahlenstorf <aahlenst@noreply.codeberg.org>
Reviewed-by: Gusted <gusted@noreply.codeberg.org>

Author: mfenniak

modules/cache/mutex_map.go

@@ -22,8 +22,27 @@ type refcountMutex struct {
 	sync.Mutex
 }
 
-// Locks the given key, and returns a function that must be invoked to unlock the key.
+// Locks the given key, and returns a function that must be invoked to unlock the key and cleanup the mutex.
 func (m *MutexMap) Lock(key string) func() {
+	mutex := m.getOrCreateMutex(key)
+	mutex.Lock()
+	return m.makeUnlock(mutex, key)
+}
+
+// Attempts to lock the given key, and returns whether the key was locked, and a function that must be invoked to unlock
+// the key and cleanup the mutex.  The returned function must be invoked even if the lock acqusition result was false.
+func (m *MutexMap) TryLock(key string) (bool, func()) {
+	mutex := m.getOrCreateMutex(key)
+	lockAcquired := mutex.TryLock()
+	if !lockAcquired {
+		return false, func() {
+			m.releaseMutex(mutex, key)
+		}
+	}
+	return true, m.makeUnlock(mutex, key)
+}
+
+func (m *MutexMap) getOrCreateMutex(key string) *refcountMutex {
 	m.mu.Lock()
 	if m.mutexMap == nil {
 		m.mutexMap = make(map[string]*refcountMutex)
@@ -36,8 +55,10 @@ func (m *MutexMap) Lock(key string) func() {
 	mutex.refCount++
 	m.mu.Unlock()
 
-	mutex.Lock()
+	return mutex
+}
 
+func (m *MutexMap) makeUnlock(mutex *refcountMutex, key string) func() {
 	unlockPending := true
 
 	return func() {
@@ -46,15 +67,18 @@ func (m *MutexMap) Lock(key string) func() {
 			// to detect and panic so that this programming error can be found closest to the source.
 			panic("MutexMap unlock invoked twice")
 		}
-
 		unlockPending = false
 		mutex.Unlock()
 
-		m.mu.Lock()
-		mutex.refCount--
-		if mutex.refCount == 0 {
-			delete(m.mutexMap, key)
-		}
-		m.mu.Unlock()
+		m.releaseMutex(mutex, key)
+	}
+}
+
+func (m *MutexMap) releaseMutex(mutex *refcountMutex, key string) {
+	m.mu.Lock()
+	mutex.refCount--
+	if mutex.refCount == 0 {
+		delete(m.mutexMap, key)
 	}
+	m.mu.Unlock()
 }

modules/cache/mutex_map_test.go

@@ -24,6 +24,31 @@ func TestMutexMap_BasicLockUnlock(t *testing.T) {
 	unlock2()
 }
 
+func TestMutexMap_BasicTryLockUnlock(t *testing.T) {
+	mm := &MutexMap{}
+
+	locked, unlock := mm.TryLock("test-key")
+	assert.True(t, locked)
+	unlock()
+
+	// Should be able to lock again
+	locked, unlock2 := mm.TryLock("test-key")
+	assert.True(t, locked)
+	unlock2()
+}
+
+func TestMutexMap_TryLock(t *testing.T) {
+	mm := &MutexMap{}
+
+	locked, unlock1 := mm.TryLock("test-key")
+	defer unlock1()
+	assert.True(t, locked)
+
+	locked, unlock2 := mm.TryLock("test-key")
+	defer unlock2()
+	assert.False(t, locked)
+}
+
 func TestMutexMap_ConcurrentSameKey(t *testing.T) {
 	mm := &MutexMap{}
 	var anotherLockActive atomic.Bool
@@ -91,6 +116,23 @@ func TestMutexMap_SimpleCleanup(t *testing.T) {
 	mm.mu.Unlock()
 }
 
+func TestMutexMap_TryLockCleanup(t *testing.T) {
+	mm := &MutexMap{}
+	_, unlock1 := mm.TryLock("test-key-1")
+	_, unlock2 := mm.TryLock("test-key-1")
+
+	mm.mu.Lock()
+	assert.Len(t, mm.mutexMap, 1)
+	mm.mu.Unlock()
+
+	unlock1()
+	unlock2()
+
+	mm.mu.Lock()
+	assert.Empty(t, mm.mutexMap)
+	mm.mu.Unlock()
+}
+
 func TestMutexMap_ConcurrentCleanup(t *testing.T) {
 	mm := &MutexMap{}
 	var foundRefGreaterThanOne atomic.Bool

routers/api/actions/runner/runner.go

@@ -13,6 +13,7 @@ import (
 	repo_model "forgejo.org/models/repo"
 	user_model "forgejo.org/models/user"
 	"forgejo.org/modules/actions"
+	"forgejo.org/modules/cache"
 	"forgejo.org/modules/log"
 	"forgejo.org/modules/setting"
 	"forgejo.org/modules/util"
@@ -36,10 +37,11 @@ var _ runnerv1connect.RunnerServiceClient = (*Service)(nil)
 
 type Service struct {
 	runnerv1connect.UnimplementedRunnerServiceHandler
+	runnerRequestKeyMutexMap cache.MutexMap
 }
 
 // Register for new runner.
-func (s *Service) Register(
+func (*Service) Register(
 	ctx context.Context,
 	req *connect.Request[runnerv1.RegisterRequest],
 ) (*connect.Response[runnerv1.RegisterResponse], error) {
@@ -109,7 +111,7 @@ func (s *Service) Register(
 	return res, nil
 }
 
-func (s *Service) Declare(
+func (*Service) Declare(
 	ctx context.Context,
 	req *connect.Request[runnerv1.DeclareRequest],
 ) (*connect.Response[runnerv1.DeclareResponse], error) {
@@ -142,6 +144,23 @@ func (s *Service) FetchTask(
 
 	requestKey := getRequestKey(ctx)
 	if requestKey != nil {
+		// It's possible for Forgejo to receive multiple concurrent requests for a given request key if the client made
+		// a request (A), request (A) took longer than the client's HTTP timeout, request (A) continues to run on
+		// Forgejo, and the client sends request (B).  In that case, we need to protect against reading from the
+		// database and sending only *some* of the tasks for the request key back to the runner, as they get assigned
+		// and committed to the database from request (A), but while request (A) is still running and request (B) is
+		// received.  To do this, we lock on the request key with a MutexMap.
+		//
+		// The lock must be held for the entirety of `FetchTask`, so even if the request key isn't used to recover
+		// tasks, the lock is held while new tasks are picked.
+		locked, cleanup := s.runnerRequestKeyMutexMap.TryLock(*requestKey)
+		defer cleanup()
+		if !locked {
+			// Another goroutine is currently processing some work for this request key.  Provide an error to the
+			// client.  This will allow the client to retry with the same request key at its typical fetch interval.
+			return nil, connect.NewError(connect.CodeInternal, errors.New("request key is currently locked; retry soon"))
+		}
+
 		recoveredTasks, err := recoverTasks(ctx, runner, *requestKey)
 		if err != nil {
 			return nil, connect.NewError(connect.CodeInternal, err)
@@ -201,7 +220,7 @@ func (s *Service) FetchTask(
 	return res, nil
 }
 
-func (s *Service) FetchSingleTask(
+func (*Service) FetchSingleTask(
 	ctx context.Context,
 	req *connect.Request[runnerv1.FetchSingleTaskRequest],
 ) (*connect.Response[runnerv1.FetchSingleTaskResponse], error) {
@@ -253,7 +272,7 @@ func (s *Service) FetchSingleTask(
 }
 
 // UpdateTask updates the task status.
-func (s *Service) UpdateTask(
+func (*Service) UpdateTask(
 	ctx context.Context,
 	req *connect.Request[runnerv1.UpdateTaskRequest],
 ) (*connect.Response[runnerv1.UpdateTaskResponse], error) {
@@ -334,7 +353,7 @@ func (s *Service) UpdateTask(
 }
 
 // UpdateLog uploads log of the task.
-func (s *Service) UpdateLog(
+func (*Service) UpdateLog(
 	ctx context.Context,
 	req *connect.Request[runnerv1.UpdateLogRequest],
 ) (*connect.Response[runnerv1.UpdateLogResponse], error) {

tests/integration/actions_fetch_task_test.go

@@ -5,6 +5,7 @@ package integration
 
 import (
 	"net/url"
+	"sync"
 	"testing"
 	"testing/fstest"
 
@@ -13,10 +14,12 @@ import (
 	unit_model "forgejo.org/models/unit"
 	"forgejo.org/models/unittest"
 	user_model "forgejo.org/models/user"
+	"forgejo.org/modules/container"
 	"forgejo.org/modules/setting"
 	"forgejo.org/modules/util"
 	"forgejo.org/tests/forgery"
 
+	runnerv1 "code.forgejo.org/forgejo/actions-proto/runner/v1"
 	"code.forgejo.org/xorm/xorm/convert"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -201,6 +204,91 @@ jobs:
 	})
 }
 
+func TestActionFetchTask_IdempotentConcurrent(t *testing.T) {
+	if !setting.Database.Type.IsSQLite3() {
+		// mock repo runner only supported on SQLite testing
+		t.Skip()
+	}
+
+	onApplicationRun(t, func(t *testing.T, u *url.URL) {
+		user2 := unittest.AssertExistsAndLoadBean(t, &user_model.User{ID: 2})
+
+		// create the repo
+		repo := createFetchTaskTestRepository(t, user2, "matrix.yml", `
+on:
+  push:
+jobs:
+  job1:
+    strategy:
+      matrix:
+        d1: [a, b, c, d, e]
+        d2: [a, b, c, d, e]
+    runs-on: ubuntu-latest
+    steps:
+      - run: sleep 2
+`)
+
+		runner := newMockRunner()
+		runner.registerAsRepoRunner(t, user2.Name, repo.Name, "mock-runner", []string{"ubuntu-latest"})
+
+		runner.setRequestKey("c6dacc80-dace-4cea-9aad-f0e266355d8e")
+
+		// If we make two simultaneous requests with the same runner request key, we should get either the error
+		// "request key is currently locked; retry soon", or, the same tasks from both requests.
+		concurrentCount := 15
+		type fetchResult struct {
+			index     int
+			task      *runnerv1.Task
+			addtTasks []*runnerv1.Task
+			err       error
+		}
+		fetchResults := make(chan fetchResult, concurrentCount)
+
+		var wg sync.WaitGroup
+		for i := range concurrentCount {
+			wg.Go(func() {
+				// Larger task capacity is used to make the successful call take longer, cause higher chance of problems if
+				// concurrency isn't handled correctly
+				task, addtTasks, err := runner.fetchTaskOrError(t, 10)
+				fetchResults <- fetchResult{index: i, task: task, addtTasks: addtTasks, err: err}
+			})
+		}
+
+		wg.Wait()
+		close(fetchResults)
+
+		var firstResponseTaskIDs container.Set[int64]
+		for res := range fetchResults {
+			t.Logf("res = %#v", res)
+			if res.task != nil {
+				// This response had tasks, so let's ensure they're always the same for every response.
+				taskIDs := container.Set[int64]{}
+				taskIDs.Add(res.task.GetId())
+				for _, extraTask := range res.addtTasks {
+					assert.True(t, taskIDs.Add(extraTask.GetId()))
+				}
+				if firstResponseTaskIDs == nil {
+					// first response with tasks -- record the IDs
+					firstResponseTaskIDs = taskIDs
+					assert.Len(t, taskIDs, 10)
+				} else {
+					// we've already found one response with tasks, so assert that they're all the same
+					d1 := firstResponseTaskIDs.Difference(taskIDs)
+					assert.Empty(t, d1, "first response taskIDs minus current response taskIDs should be empty")
+					d2 := taskIDs.Difference(firstResponseTaskIDs)
+					assert.Empty(t, d2, "current response taskIDs minus first response taskIDs should be empty")
+				}
+			} else if res.err != nil {
+				require.ErrorContains(t, res.err, "request key is currently locked")
+			} else {
+				assert.Fail(t, "unexpected condition - res.task = nil, res.err = nil")
+			}
+		}
+
+		assert.NotNil(t, firstResponseTaskIDs, "at least one response should return tasks")
+	})
+}
+
 func TestActionFetchTask_RequestedJob(t *testing.T) {
 	if !setting.Database.Type.IsSQLite3() {
 		// mock repo runner only supported on SQLite testing

tests/integration/actions_runner_test.go

@@ -141,6 +141,18 @@ func (r *mockRunner) registerAsEphemeralRepoRunner(t *testing.T, ownerName, repo
 	r.doRegisterEphemeral(t, runnerName, registrationToken.Token, labels)
 }
 
+func (r *mockRunner) fetchTaskOrError(t *testing.T, taskCapacity int64) (*runnerv1.Task, []*runnerv1.Task, error) {
+	resp, err := r.client.runnerServiceClient.FetchTask(t.Context(), connect.NewRequest(&runnerv1.FetchTaskRequest{
+		TasksVersion: r.lastTasksVersion,
+		TaskCapacity: &taskCapacity,
+	}))
+	if err != nil {
+		return nil, nil, err
+	}
+	r.lastTasksVersion = resp.Msg.TasksVersion
+	return resp.Msg.Task, resp.Msg.AdditionalTasks, nil
+}
+
 func (r *mockRunner) maybeFetchTask(t *testing.T) *runnerv1.Task {
 	resp, err := r.client.runnerServiceClient.FetchTask(t.Context(), connect.NewRequest(&runnerv1.FetchTaskRequest{
 		TasksVersion: r.lastTasksVersion,