From f2060db5c0126c5837031fd2717b8c69129ade60 Mon Sep 17 00:00:00 2001
From: Lukas <lukas@slucky.de>
Date: Wed, 18 Sep 2024 18:26:18 +0200
Subject: [PATCH] Fix agent last work update throttling (#4124)

---
 server/grpc/rpc.go      |  4 ++--
 server/grpc/rpc_test.go | 50 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/server/grpc/rpc.go b/server/grpc/rpc.go
index 934d87469..52edf1f99 100644
--- a/server/grpc/rpc.go
+++ b/server/grpc/rpc.go
@@ -525,8 +525,8 @@ func (s *RPC) getHostnameFromContext(ctx context.Context) (string, error) {
 }
 
 func (s *RPC) updateAgentLastWork(agent *model.Agent) error {
-	// only update agent.LastWork if not done recently
-	if time.Unix(agent.LastWork, 0).Add(updateAgentLastWorkDelay).Before(time.Now()) {
+	// only update agent.LastWork if not recently updated
+	if time.Unix(agent.LastWork, 0).Add(updateAgentLastWorkDelay).After(time.Now()) {
 		return nil
 	}
 
diff --git a/server/grpc/rpc_test.go b/server/grpc/rpc_test.go
index 1fd850b8a..01c57676b 100644
--- a/server/grpc/rpc_test.go
+++ b/server/grpc/rpc_test.go
@@ -17,9 +17,11 @@ package grpc
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/franela/goblin"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
 	"google.golang.org/grpc/metadata"
 
 	"go.woodpecker-ci.org/woodpecker/v2/server/model"
@@ -107,3 +109,51 @@ func TestRegisterAgent(t *testing.T) {
 		})
 	})
 }
+
+func TestUpdateAgentLastWork(t *testing.T) {
+	t.Run("When last work was never updated it should update last work timestamp", func(t *testing.T) {
+		agent := model.Agent{
+			LastWork: 0,
+		}
+		store := mocks_store.NewStore(t)
+		rpc := RPC{
+			store: store,
+		}
+		store.On("AgentUpdate", mock.Anything).Once().Return(nil)
+
+		err := rpc.updateAgentLastWork(&agent)
+		assert.NoError(t, err)
+
+		assert.NotZero(t, agent.LastWork)
+	})
+
+	t.Run("When last work was updated over a minute ago it should update last work timestamp", func(t *testing.T) {
+		lastWork := time.Now().Add(-time.Hour).Unix()
+		agent := model.Agent{
+			LastWork: lastWork,
+		}
+		store := mocks_store.NewStore(t)
+		rpc := RPC{
+			store: store,
+		}
+		store.On("AgentUpdate", mock.Anything).Once().Return(nil)
+
+		err := rpc.updateAgentLastWork(&agent)
+		assert.NoError(t, err)
+
+		assert.NotEqual(t, lastWork, agent.LastWork)
+	})
+
+	t.Run("When last work was updated in the last minute it should not update last work timestamp again", func(t *testing.T) {
+		lastWork := time.Now().Add(-time.Second * 30).Unix()
+		agent := model.Agent{
+			LastWork: lastWork,
+		}
+		rpc := RPC{}
+
+		err := rpc.updateAgentLastWork(&agent)
+		assert.NoError(t, err)
+
+		assert.Equal(t, lastWork, agent.LastWork)
+	})
+}