fix: scope offloaded node status query to the page. Fixes #16611

Signed-off-by: HsiuChuanHsu <hchsu2106@gmail.com>

Author: HsiuChuanHsu

persist/sqldb/explosive_offload_node_status_repo.go

@@ -26,7 +26,7 @@ func (n *explosiveOffloadNodeStatusRepo) Get(context.Context, string, string) (w
 	return nil, ErrOffloadNotSupported
 }
 
-func (n *explosiveOffloadNodeStatusRepo) List(context.Context, string) (map[UUIDVersion]wfv1.Nodes, error) {
+func (n *explosiveOffloadNodeStatusRepo) List(context.Context, string, []UUIDVersion) (map[UUIDVersion]wfv1.Nodes, error) {
 	return nil, ErrOffloadNotSupported
 }
 

persist/sqldb/mocks/OffloadNodeStatusRepo.go

@@ -26,7 +26,7 @@ type UUIDVersion struct {
 type OffloadNodeStatusRepo interface {
 	Save(ctx context.Context, uid, namespace string, nodes wfv1.Nodes) (string, error)
 	Get(ctx context.Context, uid, version string) (wfv1.Nodes, error)
-	List(ctx context.Context, namespace string) (map[UUIDVersion]wfv1.Nodes, error)
+	List(ctx context.Context, namespace string, keys []UUIDVersion) (map[UUIDVersion]wfv1.Nodes, error)
 	ListOldOffloads(ctx context.Context, namespace string) (map[string][]string, error)
 	Delete(ctx context.Context, uid, version string) error
 	IsEnabled() bool
@@ -149,8 +149,28 @@ func (wdc *nodeOffloadRepo) Get(ctx context.Context, uid, version string) (wfv1.
 	return nodes, nil
 }
 
-func (wdc *nodeOffloadRepo) List(ctx context.Context, namespace string) (map[UUIDVersion]wfv1.Nodes, error) {
-	wdc.log.WithFields(logging.Fields{"namespace": namespace}).Debug(ctx, "Listing offloaded nodes")
+// uuidVersionIn matches exactly the given (uid, version) pairs. Written as OR-of-ANDs rather
+// than a row value `(uid, version) IN ((?,?),...)` so that it behaves the same on MySQL,
+// MariaDB, Postgres and SQLite.
+//
+// Each pair costs two placeholders, so this tops out at roughly 32k pairs on MySQL. A page
+// that large is not a real scenario; batch the keys here if that ever changes.
+func uuidVersionIn(keys []UUIDVersion) db.LogicalExpr {
+	conds := make([]db.LogicalExpr, len(keys))
+	for i, key := range keys {
+		conds[i] = db.And(db.Cond{"uid": key.UID}, db.Cond{"version": key.Version})
+	}
+	return db.Or(conds...)
+}
+
+// List returns the offloaded nodes for the given keys only.
+func (wdc *nodeOffloadRepo) List(ctx context.Context, namespace string, keys []UUIDVersion) (map[UUIDVersion]wfv1.Nodes, error) {
+	// This is not merely an optimisation: db.Or() with no arguments is an empty condition,
+	// so the query would silently widen back to every nodes blob in the namespace.
+	if len(keys) == 0 {
+		return map[UUIDVersion]wfv1.Nodes{}, nil
+	}
+	wdc.log.WithFields(logging.Fields{"namespace": namespace, "keys": len(keys)}).Debug(ctx, "Listing offloaded nodes")
 	var res map[UUIDVersion]wfv1.Nodes
 	err := wdc.sessionProxy.With(ctx, func(s db.Session) error {
 		var records []nodesRecord
@@ -159,6 +179,7 @@ func (wdc *nodeOffloadRepo) List(ctx context.Context, namespace string) (map[UUI
 			From(wdc.tableName).
 			Where(db.Cond{"clustername": wdc.clusterName}).
 			And(namespaceEqual(namespace)).
+			And(uuidVersionIn(keys)).
 			All(&records)
 		if err != nil {
 			return err

persist/sqldb/offload_node_status_repo.go

@@ -0,0 +1,76 @@
+//go:build !windows
+
+package sqldb
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	wfv1 "github.com/argoproj/argo-workflows/v4/pkg/apis/workflow/v1alpha1"
+	"github.com/argoproj/argo-workflows/v4/util/logging"
+	usqldb "github.com/argoproj/argo-workflows/v4/util/sqldb"
+)
+
+// setupMySQLOffloadTest starts a MySQL or MariaDB container and returns an offload repository.
+func setupMySQLOffloadTest(ctx context.Context, t *testing.T, v usqldb.MySQLVariant) OffloadNodeStatusRepo {
+	t.Helper()
+	repo, err := NewOffloadNodeStatusRepo(ctx, logging.RequireLoggerFromContext(ctx), setupMySQLTest(ctx, t, v), "test", "argo_workflows")
+	require.NoError(t, err)
+	return repo
+}
+
+// saveOffload writes a node status and returns the version it was stored under.
+func saveOffload(ctx context.Context, t *testing.T, repo OffloadNodeStatusRepo, uid, nodeName string) string {
+	t.Helper()
+	version, err := repo.Save(ctx, uid, "argo", wfv1.Nodes{nodeName: wfv1.NodeStatus{ID: nodeName}})
+	require.NoError(t, err)
+	return version
+}
+
+// TestMySQLListOnlyReturnsRequestedKeys covers the behaviour the list path depends on: the
+// query is scoped to the keys it is given, so a caller that needs one page of workflows does
+// not pull every offloaded blob in the namespace.
+func TestMySQLListOnlyReturnsRequestedKeys(t *testing.T) {
+	for name, variant := range usqldb.MySQLVariants {
+		t.Run(name, func(t *testing.T) {
+			ctx := logging.TestContext(t.Context())
+			repo := setupMySQLOffloadTest(ctx, t, variant)
+
+			wantedA := UUIDVersion{UID: "uid-a", Version: saveOffload(ctx, t, repo, "uid-a", "node-a")}
+			wantedB := UUIDVersion{UID: "uid-b", Version: saveOffload(ctx, t, repo, "uid-b", "node-b")}
+			unwanted := UUIDVersion{UID: "uid-c", Version: saveOffload(ctx, t, repo, "uid-c", "node-c")}
+
+			got, err := repo.List(ctx, "argo", []UUIDVersion{wantedA, wantedB})
+			require.NoError(t, err)
+
+			assert.Len(t, got, 2)
+			assert.Contains(t, got, wantedA)
+			assert.Contains(t, got, wantedB)
+			assert.NotContains(t, got, unwanted, "a key that was not asked for must not come back")
+		})
+	}
+}
+
+// TestMySQLListExcludesSupersededVersions pins down why the version has to be part of the
+// condition. Older offloads of the same workflow stay in the table until they are garbage
+// collected, so filtering on the uid alone would drag those blobs back too.
+// Dialect portability is already covered on both engines by TestMySQLListOnlyReturnsRequestedKeys,
+// so one engine is enough here.
+func TestMySQLListExcludesSupersededVersions(t *testing.T) {
+	ctx := logging.TestContext(t.Context())
+	repo := setupMySQLOffloadTest(ctx, t, usqldb.MySQLVariants["MySQL"])
+
+	oldVersion := saveOffload(ctx, t, repo, "uid-a", "node-old")
+	newVersion := saveOffload(ctx, t, repo, "uid-a", "node-new")
+	require.NotEqual(t, oldVersion, newVersion, "the two writes must produce different versions")
+
+	got, err := repo.List(ctx, "argo", []UUIDVersion{{UID: "uid-a", Version: newVersion}})
+	require.NoError(t, err)
+
+	assert.Len(t, got, 1)
+	assert.Contains(t, got, UUIDVersion{UID: "uid-a", Version: newVersion})
+	assert.NotContains(t, got, UUIDVersion{UID: "uid-a", Version: oldVersion})
+}

persist/sqldb/offload_node_status_repo_mysql_test.go

@@ -7,8 +7,18 @@ import (
 	"github.com/stretchr/testify/require"
 
 	wfv1 "github.com/argoproj/argo-workflows/v4/pkg/apis/workflow/v1alpha1"
+	"github.com/argoproj/argo-workflows/v4/util/logging"
 )
 
+// Test_ListWithoutKeys guards the empty-keys early return. Without it the generated condition
+// is empty, which would widen the query back to the whole namespace. A zero-value repo has no
+// session, so anything past the guard panics rather than quietly querying.
+func Test_ListWithoutKeys(t *testing.T) {
+	got, err := (&nodeOffloadRepo{}).List(logging.TestContext(t.Context()), "argo", nil)
+	require.NoError(t, err)
+	assert.Empty(t, got)
+}
+
 func Test_nodeStatusVersion(t *testing.T) {
 	t.Run("Empty", func(t *testing.T) {
 		marshalled, version, err := nodeStatusVersion(nil)

persist/sqldb/offload_node_status_repo_test.go

@@ -27,6 +27,13 @@ import (
 // setupMySQLArchiveTest starts a MySQL or MariaDB container, runs migrations, and returns a WorkflowArchive.
 func setupMySQLArchiveTest(ctx context.Context, t *testing.T, v usqldb.MySQLVariant) WorkflowArchive {
 	t.Helper()
+	return NewWorkflowArchive(setupMySQLTest(ctx, t, v), "test", "", instanceid.NewService(""))
+}
+
+// setupMySQLTest starts a MySQL or MariaDB container, runs migrations, and returns a session
+// proxy the caller can build any repository on.
+func setupMySQLTest(ctx context.Context, t *testing.T, v usqldb.MySQLVariant) *usqldb.SessionProxy {
+	t.Helper()
 
 	c, err := testmysql.Run(ctx,
 		v.Image,
@@ -73,7 +80,7 @@ func setupMySQLArchiveTest(ctx context.Context, t *testing.T, v usqldb.MySQLVari
 
 	t.Cleanup(func() { proxy.Close() })
 
-	return NewWorkflowArchive(proxy, "test", "", instanceid.NewService(""))
+	return proxy
 }
 
 // TestMySQLListWorkflows verifies that JSON_EXTRACT/JSON_UNQUOTE queries in

persist/sqldb/workflow_archive_mysql_test.go

@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"maps"
+	"slices"
 	"sort"
 	"sync"
 	"time"
@@ -285,19 +287,22 @@ func (s *workflowServer) ListWorkflows(ctx context.Context, req *workflowpkg.Wor
 	}
 
 	cleaner := fields.NewCleaner(req.Fields)
-	logger := logging.RequireLoggerFromContext(ctx)
 	if s.offloadNodeStatusRepo.IsEnabled() && !cleaner.WillExclude("items.status.nodes") {
-		offloadedNodes, err := s.offloadNodeStatusRepo.List(ctx, req.Namespace)
-		if err != nil {
-			return nil, sutils.ToStatusError(err, codes.Internal)
-		}
+		// This page is already resolved, so we know exactly which offloaded rows we need.
+		offloaded := map[int]sqldb.UUIDVersion{}
 		for i, wf := range wfs {
 			if wf.Status.IsOffloadNodeStatus() {
-				if s.offloadNodeStatusRepo.IsEnabled() {
-					wfs[i].Status.Nodes = offloadedNodes[sqldb.UUIDVersion{UID: string(wf.UID), Version: wf.GetOffloadNodeStatusVersion()}]
-				} else {
-					logger.WithFields(logging.Fields{"namespace": wf.Namespace, "name": wf.Name}).Warn(ctx, sqldb.OffloadNodeStatusDisabled)
-				}
+				offloaded[i] = sqldb.UUIDVersion{UID: string(wf.UID), Version: wf.GetOffloadNodeStatusVersion()}
+			}
+		}
+		// Nothing on this page is offloaded, so there is nothing to fetch.
+		if len(offloaded) > 0 {
+			offloadedNodes, err := s.offloadNodeStatusRepo.List(ctx, req.Namespace, slices.Collect(maps.Values(offloaded)))
+			if err != nil {
+				return nil, sutils.ToStatusError(err, codes.Internal)
+			}
+			for i, key := range offloaded {
+				wfs[i].Status.Nodes = offloadedNodes[key]
 			}
 		}
 	}

server/workflow/workflow_server.go

@@ -0,0 +1,109 @@
+package workflow
+
+import (
+	"context"
+	"testing"
+
+	"github.com/go-jose/go-jose/v4/jwt"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+	authorizationv1 "k8s.io/api/authorization/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/kubernetes/fake"
+	ktesting "k8s.io/client-go/testing"
+
+	"github.com/argoproj/argo-workflows/v4/persist/sqldb"
+	"github.com/argoproj/argo-workflows/v4/persist/sqldb/mocks"
+	workflowpkg "github.com/argoproj/argo-workflows/v4/pkg/apiclient/workflow"
+	"github.com/argoproj/argo-workflows/v4/pkg/apis/workflow/v1alpha1"
+	v1alpha "github.com/argoproj/argo-workflows/v4/pkg/client/clientset/versioned/fake"
+	"github.com/argoproj/argo-workflows/v4/server/auth"
+	authtypes "github.com/argoproj/argo-workflows/v4/server/auth/types"
+	"github.com/argoproj/argo-workflows/v4/server/clusterworkflowtemplate"
+	"github.com/argoproj/argo-workflows/v4/server/workflow/store"
+	"github.com/argoproj/argo-workflows/v4/server/workflowtemplate"
+	"github.com/argoproj/argo-workflows/v4/util/instanceid"
+	"github.com/argoproj/argo-workflows/v4/util/logging"
+)
+
+// offloadedWorkflow builds a workflow whose node status lives in the offload table.
+func offloadedWorkflow(name, uid, version string) v1alpha1.Workflow {
+	return v1alpha1.Workflow{
+		ObjectMeta: metav1.ObjectMeta{Name: name, Namespace: "argo", UID: types.UID(uid)},
+		Status:     v1alpha1.WorkflowStatus{OffloadNodeStatusVersion: version},
+	}
+}
+
+// offloadTestServer builds the smallest server that can serve ListWorkflows, and hands back
+// the offload mock so tests can assert on how it was called. It deliberately does not reuse
+// getWorkflowServer, which is shared by many other tests and does not expose the mock.
+func offloadTestServer(t *testing.T, wfs ...v1alpha1.Workflow) (Server, context.Context, *mocks.OffloadNodeStatusRepo) {
+	t.Helper()
+
+	offloadNodeStatusRepo := &mocks.OffloadNodeStatusRepo{}
+	offloadNodeStatusRepo.On("IsEnabled", mock.Anything).Return(true)
+
+	archivedRepo := &mocks.WorkflowArchive{}
+	archivedRepo.On("CountWorkflows", mock.Anything, mock.Anything).Return(int64(0), nil)
+	archivedRepo.On("ListWorkflows", mock.Anything, mock.Anything).Return(v1alpha1.Workflows{}, nil)
+	archivedRepo.On("HasMoreWorkflows", mock.Anything, mock.Anything).Return(false, nil)
+
+	kubeClientSet := fake.NewClientset()
+	kubeClientSet.PrependReactor("create", "selfsubjectaccessreviews", func(action ktesting.Action) (handled bool, ret runtime.Object, err error) {
+		return true, &authorizationv1.SelfSubjectAccessReview{
+			Status: authorizationv1.SubjectAccessReviewStatus{Allowed: true},
+		}, nil
+	})
+
+	wfClientset := v1alpha.NewClientset()
+	ctx := logging.TestContext(t.Context())
+	ctx = context.WithValue(context.WithValue(context.WithValue(ctx, auth.WfKey, wfClientset), auth.KubeKey, kubeClientSet), auth.ClaimsKey, &authtypes.Claims{Claims: jwt.Claims{Subject: "my-sub"}})
+
+	// An empty instance ID keeps the store from requiring an instance-id label on the fixtures.
+	instanceIDSvc := instanceid.NewService("")
+	wfStore, err := store.NewSQLiteStore(instanceIDSvc)
+	require.NoError(t, err)
+	for i := range wfs {
+		require.NoError(t, wfStore.Add(&wfs[i]))
+	}
+
+	namespace := "argo"
+	server := NewServer(ctx, instanceIDSvc, offloadNodeStatusRepo, archivedRepo, wfClientset, wfStore, nil, workflowtemplate.NewClientStore(), clusterworkflowtemplate.NewClientStore(), nil, &namespace, nil)
+	return server, ctx, offloadNodeStatusRepo
+}
+
+// TestListWorkflows_PassesOnlyPageKeys asserts that the offload query is scoped to the
+// workflows on this page, rather than to the whole namespace.
+func TestListWorkflows_PassesOnlyPageKeys(t *testing.T) {
+	wantedA := offloadedWorkflow("offloaded-a", "uid-a", "v1")
+	wantedB := offloadedWorkflow("offloaded-b", "uid-b", "v2")
+	server, ctx, offloadNodeStatusRepo := offloadTestServer(t, wantedA, wantedB, offloadedWorkflow("inline", "uid-c", ""))
+
+	offloadNodeStatusRepo.On("List", mock.Anything, mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
+
+	_, err := server.ListWorkflows(ctx, &workflowpkg.WorkflowListRequest{Namespace: "argo"})
+	require.NoError(t, err)
+
+	offloadNodeStatusRepo.AssertNumberOfCalls(t, "List", 1)
+	call := offloadNodeStatusRepo.Calls[len(offloadNodeStatusRepo.Calls)-1]
+	require.Equal(t, "List", call.Method)
+	assert.Equal(t, "argo", call.Arguments[0])
+	assert.ElementsMatch(t, []sqldb.UUIDVersion{
+		{UID: "uid-a", Version: "v1"},
+		{UID: "uid-b", Version: "v2"},
+	}, call.Arguments[1], "List must be given exactly the offloaded keys on this page")
+}
+
+// TestListWorkflows_SkipsQueryWhenPageHasNoOffload asserts that a page with nothing offloaded
+// issues no offload query at all. This is the common case for most users.
+func TestListWorkflows_SkipsQueryWhenPageHasNoOffload(t *testing.T) {
+	server, ctx, offloadNodeStatusRepo := offloadTestServer(t, offloadedWorkflow("inline-a", "uid-a", ""), offloadedWorkflow("inline-b", "uid-b", ""))
+
+	_, err := server.ListWorkflows(ctx, &workflowpkg.WorkflowListRequest{Namespace: "argo"})
+	require.NoError(t, err)
+
+	offloadNodeStatusRepo.AssertNotCalled(t, "List")
+}

server/workflow/workflow_server_offload_test.go

@@ -598,7 +598,7 @@ func getWorkflowServer(t *testing.T) (workflowpkg.WorkflowServiceServer, context
 
 	offloadNodeStatusRepo := &mocks.OffloadNodeStatusRepo{}
 	offloadNodeStatusRepo.On("IsEnabled", mock.Anything).Return(true)
-	offloadNodeStatusRepo.On("List", mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
+	offloadNodeStatusRepo.On("List", mock.Anything, mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
 
 	archivedRepo := &mocks.WorkflowArchive{}
 
@@ -1501,7 +1501,7 @@ func getWorkflowServerWithArtifacts(t *testing.T, template runtime.Object, defau
 
 	offloadNodeStatusRepo := &mocks.OffloadNodeStatusRepo{}
 	offloadNodeStatusRepo.On("IsEnabled", mock.Anything).Return(true)
-	offloadNodeStatusRepo.On("List", mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
+	offloadNodeStatusRepo.On("List", mock.Anything, mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
 
 	archivedRepo := &mocks.WorkflowArchive{}
 

server/workflow/workflow_server_test.go

@@ -35,7 +35,7 @@ func Test_archivedWorkflowServer(t *testing.T) {
 	wfClient := &argofake.Clientset{}
 	offloadNodeStatusRepo := &mocks.OffloadNodeStatusRepo{}
 	offloadNodeStatusRepo.On("IsEnabled", mock.Anything).Return(true)
-	offloadNodeStatusRepo.On("List", mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
+	offloadNodeStatusRepo.On("List", mock.Anything, mock.Anything).Return(map[sqldb.UUIDVersion]v1alpha1.Nodes{}, nil)
 	w := NewWorkflowArchiveServer(repo, offloadNodeStatusRepo, nil)
 	allowed := true
 	kubeClient.AddReactor("create", "selfsubjectaccessreviews", func(action k8stesting.Action) (handled bool, ret runtime.Object, err error) {