feat: add UpdateRegistrationStatus to agent store

internal/core/registration.go

@@ -6,11 +6,17 @@ import "github.com/openkcm/krypton/internal/clock"
 type AgentRegistrationStatus string
 
 const (
+	// AgentRegistrationStatusRegistered indicates that the registration is active but has not yet been marked as healthy or unhealthy.
+	AgentRegistrationStatusRegistered AgentRegistrationStatus = "registered"
+
 	// AgentRegistrationStatusHealthy indicates that the registration is active and healthy.
 	AgentRegistrationStatusHealthy AgentRegistrationStatus = "healthy"
 
 	// AgentRegistrationStatusUnhealthy indicates that the registration is active but unhealthy.
 	AgentRegistrationStatusUnhealthy AgentRegistrationStatus = "unhealthy"
+
+	// AgentRegistrationStatusDeregistered indicates that the registration has been de-registered and is no longer active.
+	AgentRegistrationStatusDeregistered AgentRegistrationStatus = "de-registered"
 )
 
 // AgentRegistration entries are used to track the status of agents and their heartbeats.

pkg/api/agents/handler.go

@@ -73,7 +73,7 @@ func (a *agent) register(w http.ResponseWriter, r *http.Request) {
 		Registration: core.AgentRegistration{
 			Name:          xAgentName,
 			InstanceID:    xAgentID,
-			Status:        core.AgentRegistrationStatusHealthy,
+			Status:        core.AgentRegistrationStatusRegistered,
 			LastHeartbeat: clock.Now(),
 		},
 	})

pkg/api/agents/register_test.go

@@ -167,7 +167,7 @@ func TestAgentRegister(t *testing.T) {
 			assert.Equal(t, core.AgentRegistration{
 				Name:          expAgentName,
 				InstanceID:    expAgentID,
-				Status:        core.AgentRegistrationStatusHealthy,
+				Status:        core.AgentRegistrationStatusRegistered,
 				LastHeartbeat: result.Registration.LastHeartbeat,
 				CreatedAt:     result.Registration.CreatedAt,
 				UpdatedAt:     result.Registration.UpdatedAt,

pkg/store/agent.go

@@ -3,6 +3,7 @@ package store
 import (
 	"context"
 	"errors"
+	"time"
 
 	"github.com/openkcm/krypton/internal/core"
 )
@@ -12,6 +13,7 @@ var ErrAgentRegistrationNotFound = errors.New("agent registration not found")
 type Agent interface {
 	Register(ctx context.Context, query RegisterAgentQuery) (RegisterAgentResult, error)
 	Get(ctx context.Context, query GetAgentQuery) (GetAgentResult, error)
+	UpdateRegistrationStatus(ctx context.Context, query UpdateRegistrationStatusQuery) error
 }
 
 type (
@@ -31,4 +33,12 @@ type (
 	GetAgentResult struct {
 		Registration core.AgentRegistration
 	}
+
+	UpdateRegistrationStatusQuery struct {
+		Name               string
+		InstanceID         string
+		FromStatus         []core.AgentRegistrationStatus
+		ToStatus           core.AgentRegistrationStatus
+		HeartbeatThreshold time.Duration
+	}
 )

pkg/store/sql/agent.go

@@ -3,6 +3,7 @@ package sql
 import (
 	"context"
 	"database/sql"
+	"strings"
 
 	"github.com/openkcm/krypton/internal/clock"
 	"github.com/openkcm/krypton/internal/core"
@@ -109,3 +110,30 @@ func (s *AgentStore) Get(ctx context.Context, q store.GetAgentQuery) (store.GetA
 		Registration: reg,
 	}, nil
 }
+
+// UpdateRegistrationStatus implements [store.Agent].
+func (s *AgentStore) UpdateRegistrationStatus(ctx context.Context, query store.UpdateRegistrationStatusQuery) error {
+	var sb strings.Builder
+	now := clock.Now()
+	sb.WriteString(`UPDATE agent_registrations
+	SET status = $1, updated_at = $2
+	WHERE name = $3 AND instance_id = $4
+	AND status = ANY($5)`)
+	args := []any{
+		query.ToStatus,
+		now,
+		query.Name,
+		query.InstanceID,
+		query.FromStatus,
+	}
+	if query.HeartbeatThreshold != 0 {
+		sb.WriteString(` AND ($6 - last_heartbeat) > $7`)
+		args = append(args, now)
+		args = append(args, query.HeartbeatThreshold.Nanoseconds())
+	}
+
+	_, err := s.db.ExecContext(ctx, sb.String(),
+		args...,
+	)
+	return err
+}

pkg/store/sql/agent_test.go

@@ -5,6 +5,8 @@ import (
 	"database/sql"
 	"os"
 	"testing"
+	"testing/synctest"
+	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
@@ -172,6 +174,255 @@ func TestGet(t *testing.T) {
 	})
 }
 
+func TestUpdateRegistrationStatus(t *testing.T) {
+	// given
+	ctx := t.Context()
+
+	db, err := sql.Open("postgres", pgConnStr)
+	require.NoError(t, err)
+
+	t.Cleanup(func() {
+		db.Close()
+	})
+
+	subj, err := storesql.NewAgentStore(ctx, db)
+	require.NoError(t, err)
+
+	t.Run("UpdateRegistrationStatus", func(t *testing.T) {
+		tests := []struct {
+			name                string
+			sleepDuration       time.Duration
+			expectStatusChanged bool
+		}{
+			{
+				name:                "should update status when heartbeat threshold exceeded",
+				sleepDuration:       21 * time.Second, // threshold + 1s
+				expectStatusChanged: true,
+			},
+			{
+				name:          "should not update status when heartbeat threshold not exceeded",
+				sleepDuration: 10 * time.Second, // threshold - 10s
+
+				expectStatusChanged: false,
+			},
+		}
+		heartbeatThreshold := 15 * time.Second
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				synctest.Test(t, func(t *testing.T) {
+					// given
+					registration := core.AgentRegistration{
+						Name:          "test-agent",
+						InstanceID:    uuid.New().String(),
+						Status:        core.AgentRegistrationStatusHealthy,
+						LastHeartbeat: clock.Now(),
+					}
+
+					prevResult, err := subj.Register(ctx, store.RegisterAgentQuery{
+						Registration: registration,
+					})
+
+					require.NoError(t, err)
+
+					time.Sleep(tt.sleepDuration)
+					synctest.Wait()
+
+					updateQuery := store.UpdateRegistrationStatusQuery{
+						Name:       registration.Name,
+						InstanceID: registration.InstanceID,
+						FromStatus: []core.AgentRegistrationStatus{
+							core.AgentRegistrationStatusHealthy,
+						},
+						ToStatus:           core.AgentRegistrationStatusUnhealthy,
+						HeartbeatThreshold: heartbeatThreshold,
+					}
+
+					// when
+					err = subj.UpdateRegistrationStatus(ctx, updateQuery)
+
+					// then
+					assert.NoError(t, err)
+
+					getResult, err := subj.Get(ctx, store.GetAgentQuery{
+						Name:       registration.Name,
+						InstanceID: registration.InstanceID,
+					})
+					require.NoError(t, err)
+
+					assert.Equal(t, prevResult.Registration.Name, getResult.Registration.Name)
+					assert.Equal(t, prevResult.Registration.InstanceID, getResult.Registration.InstanceID)
+					assert.Equal(t, prevResult.Registration.LastHeartbeat, getResult.Registration.LastHeartbeat)
+					assert.Equal(t, prevResult.Registration.CreatedAt, getResult.Registration.CreatedAt)
+					if tt.expectStatusChanged {
+						assert.Equal(t, updateQuery.ToStatus, getResult.Registration.Status)
+						assert.NotEqual(t, prevResult.Registration.UpdatedAt, getResult.Registration.UpdatedAt)
+						assert.Greater(t, getResult.Registration.UpdatedAt, prevResult.Registration.UpdatedAt)
+					} else {
+						assert.Equal(t, prevResult.Registration.Status, getResult.Registration.Status)
+						assert.Equal(t, prevResult.Registration.UpdatedAt, getResult.Registration.UpdatedAt)
+					}
+				})
+			})
+		}
+	})
+
+	t.Run("should deregister without a LastHeartbeat threshold", func(t *testing.T) {
+		tests := []struct {
+			name       string
+			fromStatus core.AgentRegistrationStatus
+		}{
+			{
+				name:       "from healthy",
+				fromStatus: core.AgentRegistrationStatusHealthy,
+			},
+			{
+				name:       "from unhealthy",
+				fromStatus: core.AgentRegistrationStatusUnhealthy,
+			},
+			{
+				name:       "from registered",
+				fromStatus: core.AgentRegistrationStatusRegistered,
+			},
+		}
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				// given
+				registration := core.AgentRegistration{
+					Name:          "test-agent",
+					InstanceID:    uuid.New().String(),
+					Status:        tt.fromStatus,
+					LastHeartbeat: clock.Now(),
+				}
+
+				prevResult, err := subj.Register(ctx, store.RegisterAgentQuery{
+					Registration: registration,
+				})
+				require.NoError(t, err)
+
+				updateQuery := store.UpdateRegistrationStatusQuery{
+					Name:       registration.Name,
+					InstanceID: registration.InstanceID,
+					FromStatus: []core.AgentRegistrationStatus{
+						tt.fromStatus,
+					},
+					ToStatus: core.AgentRegistrationStatusDeregistered,
+				}
+
+				// when
+				err = subj.UpdateRegistrationStatus(ctx, updateQuery)
+
+				// then
+				assert.NoError(t, err)
+
+				getResult, err := subj.Get(ctx, store.GetAgentQuery{
+					Name:       registration.Name,
+					InstanceID: registration.InstanceID,
+				})
+				require.NoError(t, err)
+
+				assert.Equal(t, updateQuery.ToStatus, getResult.Registration.Status)
+				assert.Greater(t, getResult.Registration.UpdatedAt, prevResult.Registration.UpdatedAt)
+
+				assert.Equal(t, prevResult.Registration.Name, getResult.Registration.Name)
+				assert.Equal(t, prevResult.Registration.InstanceID, getResult.Registration.InstanceID)
+				assert.Equal(t, prevResult.Registration.LastHeartbeat, getResult.Registration.LastHeartbeat)
+				assert.Equal(t, prevResult.Registration.CreatedAt, getResult.Registration.CreatedAt)
+			})
+		}
+	})
+
+	t.Run("should not update status if query does not match registration", func(t *testing.T) {
+		// given
+		registration := core.AgentRegistration{
+			Name:          "test-agent",
+			InstanceID:    uuid.New().String(),
+			Status:        core.AgentRegistrationStatusHealthy,
+			LastHeartbeat: clock.Now(),
+		}
+
+		prevResult, err := subj.Register(ctx, store.RegisterAgentQuery{
+			Registration: registration,
+		})
+
+		require.NoError(t, err)
+
+		tts := []struct {
+			name        string
+			updateQuery store.UpdateRegistrationStatusQuery
+		}{
+			{
+				name: "non-existent name",
+				updateQuery: store.UpdateRegistrationStatusQuery{
+					Name:       "non-existent-agent",
+					InstanceID: registration.InstanceID,
+					FromStatus: []core.AgentRegistrationStatus{
+						registration.Status,
+					},
+					ToStatus: core.AgentRegistrationStatusDeregistered,
+				},
+			},
+			{
+				name: "non-existent instance ID",
+				updateQuery: store.UpdateRegistrationStatusQuery{
+					Name:       registration.Name,
+					InstanceID: uuid.New().String(),
+					FromStatus: []core.AgentRegistrationStatus{
+						registration.Status,
+					},
+					ToStatus: core.AgentRegistrationStatusDeregistered,
+				},
+			},
+			{
+				name: "from status does not match",
+				updateQuery: store.UpdateRegistrationStatusQuery{
+					Name:       registration.Name,
+					InstanceID: registration.InstanceID,
+					FromStatus: []core.AgentRegistrationStatus{
+						core.AgentRegistrationStatusUnhealthy,
+					},
+					ToStatus: core.AgentRegistrationStatusDeregistered,
+				},
+			},
+			{
+				name: "heartbeat threshold not exceeded",
+				updateQuery: store.UpdateRegistrationStatusQuery{
+					Name:       registration.Name,
+					InstanceID: registration.InstanceID,
+					FromStatus: []core.AgentRegistrationStatus{
+						registration.Status,
+					},
+					ToStatus:           core.AgentRegistrationStatusDeregistered,
+					HeartbeatThreshold: 100 * time.Second,
+				},
+			},
+		}
+		for _, tt := range tts {
+			t.Run(tt.name, func(t *testing.T) {
+				// when
+				err = subj.UpdateRegistrationStatus(ctx, tt.updateQuery)
+
+				// then
+				assert.NoError(t, err)
+
+				getResult, err := subj.Get(ctx, store.GetAgentQuery{
+					Name:       registration.Name,
+					InstanceID: registration.InstanceID,
+				})
+
+				require.NoError(t, err)
+
+				assert.Equal(t, registration.Status, getResult.Registration.Status)
+				assert.Equal(t, prevResult.Registration.Name, getResult.Registration.Name)
+				assert.Equal(t, prevResult.Registration.InstanceID, getResult.Registration.InstanceID)
+				assert.Equal(t, prevResult.Registration.LastHeartbeat, getResult.Registration.LastHeartbeat)
+				assert.Equal(t, prevResult.Registration.CreatedAt, getResult.Registration.CreatedAt)
+				assert.Equal(t, prevResult.Registration.UpdatedAt, getResult.Registration.UpdatedAt)
+			})
+		}
+	})
+}
+
 func setupPostgres() (func(), error) {
 	ctx := context.Background()