workloadmanager: prevent inactive session GC starvation

pkg/common/types/sandbox.go

@@ -40,7 +40,7 @@ type SandboxInfo struct {
 	// a raw nanosecond integer, making the persisted JSON unambiguous.
 	IdleTimeout metav1.Duration `json:"idleTimeout,omitempty"`
 	// LastActivityAt is populated transiently from the store's last-activity sorted set
-	// during ListInactiveSandboxes. It is intentionally excluded from JSON serialization.
+	// during ListInactiveSandboxes or GetSandboxWithLastActivity. It is excluded from JSON serialization.
 	LastActivityAt time.Time `json:"-"`
 	Status         string    `json:"status"`
 }

pkg/router/session_manager_test.go

@@ -742,3 +742,10 @@ func TestSetIdentityHeader_NoClaimsOrNoManager(t *testing.T) {
 		})
 	}
 }
+
+func (*fakeStoreClient) ScanInactiveSandboxIDs(_ context.Context, _ time.Time, _ uint64, _ int64) ([]string, uint64, error) {
+	return nil, 0, nil
+}
+func (*fakeStoreClient) GetSandboxWithLastActivity(_ context.Context, _ string) (*types.SandboxInfo, error) {
+	return nil, nil
+}

pkg/store/gc_scan.go

@@ -0,0 +1,43 @@
+/*
+Copyright The Volcano Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package store
+
+import (
+	"fmt"
+	"strconv"
+	"time"
+)
+
+// inactiveIDsFromScan decodes the member/score pairs returned by ZSCAN.
+// Filtering here avoids loading JSON for active sessions, while continuation
+// always comes from the raw scan, even when no members pass the filter.
+func inactiveIDsFromScan(elements []string, before time.Time) ([]string, error) {
+	if len(elements)%2 != 0 {
+		return nil, fmt.Errorf("ZSCAN returned an incomplete member/score pair")
+	}
+	ids := make([]string, 0, len(elements)/2)
+	for i := 0; i < len(elements); i += 2 {
+		score, err := strconv.ParseFloat(elements[i+1], 64)
+		if err != nil {
+			return nil, fmt.Errorf("ZSCAN score for session %s: %w", elements[i], err)
+		}
+		if score <= float64(before.Unix()) {
+			ids = append(ids, elements[i])
+		}
+	}
+	return ids, nil
+}

pkg/store/gc_scan_test.go

@@ -0,0 +1,150 @@
+/*
+Copyright The Volcano Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package store
+
+import (
+	"context"
+	"fmt"
+	"testing"
+	"time"
+
+	"github.com/alicebob/miniredis/v2"
+	"github.com/stretchr/testify/require"
+)
+
+func newGCStore(t *testing.T, provider string) (Store, *miniredis.Miniredis) {
+	t.Helper()
+	var client Store
+	var server *miniredis.Miniredis
+	if provider == "redis" {
+		client, server = newTestRedisClient(t)
+	} else {
+		client, server = newValkeyTestClient(t)
+	}
+	t.Cleanup(func() { require.NoError(t, client.Close()) })
+	return client, server
+}
+
+func TestGCStoreScan(t *testing.T) {
+	for _, provider := range []string{"redis", "valkey"} {
+		t.Run(provider, func(t *testing.T) {
+			client, server := newGCStore(t, provider)
+			ctx := context.Background()
+			cutoff := time.Now().Add(-time.Minute).Truncate(time.Second)
+			// Active IDs sort first in miniredis, exercising an empty filtered
+			// page with a nonzero cursor. Real servers need not sort ZSCAN output.
+			for i := 0; i < 50; i++ {
+				_, err := server.ZAdd("sandbox:last_activity", float64(cutoff.Add(time.Minute).Unix()), fmt.Sprintf("active-%03d", i))
+				require.NoError(t, err)
+			}
+			expected := make(map[string]bool)
+			for i := 0; i < 237; i++ {
+				id := fmt.Sprintf("inactive-%03d", i)
+				expected[id] = true
+				// Ties, inclusive cutoff, and IDs with no JSON record are all valid
+				// scan results. Hydration decides whether the session still exists.
+				_, err := server.ZAdd("sandbox:last_activity", float64(cutoff.Unix()), id)
+				require.NoError(t, err)
+			}
+			seen := make(map[string]bool)
+			var cursor uint64
+			for page := 0; ; page++ {
+				require.Less(t, page, 1000, "scan did not finish")
+				ids, next, err := client.ScanInactiveSandboxIDs(ctx, cutoff, cursor, 17)
+				require.NoError(t, err)
+				for _, id := range ids {
+					require.True(t, expected[id], "active session returned: %s", id)
+					seen[id] = true
+				}
+				cursor = next
+				if cursor == 0 {
+					break
+				}
+			}
+			require.Equal(t, expected, seen)
+			for _, count := range []int64{0, -1} {
+				_, next, err := client.ScanInactiveSandboxIDs(ctx, cutoff, 19, count)
+				require.Error(t, err)
+				require.Equal(t, uint64(19), next)
+			}
+		})
+	}
+}
+
+func TestGCStoreHydration(t *testing.T) {
+	for _, provider := range []string{"redis", "valkey"} {
+		t.Run(provider, func(t *testing.T) {
+			client, server := newGCStore(t, provider)
+			ctx := context.Background()
+			now := time.Now().Truncate(time.Second)
+			sandbox := newTestSandbox("sb", "session", now.Add(time.Hour))
+			require.NoError(t, client.StoreSandbox(ctx, sandbox))
+			oldActivity := now.Add(-time.Hour)
+			require.NoError(t, client.UpdateSessionLastActivity(ctx, sandbox.SessionID, oldActivity))
+			loaded, err := client.GetSandboxWithLastActivity(ctx, sandbox.SessionID)
+			require.NoError(t, err)
+			require.True(t, loaded.LastActivityAt.Equal(oldActivity))
+			// Activity comes from the current score, never the session JSON or
+			// the earlier scan. Changes to the JSON are also observed.
+			sandbox.Name = "updated-name"
+			require.NoError(t, client.UpdateSandbox(ctx, sandbox))
+			require.NoError(t, client.UpdateSessionLastActivity(ctx, sandbox.SessionID, now))
+			loaded, err = client.GetSandboxWithLastActivity(ctx, sandbox.SessionID)
+			require.NoError(t, err)
+			require.True(t, loaded.LastActivityAt.Equal(now))
+			require.Equal(t, "updated-name", loaded.Name)
+			server.Del("sandbox:last_activity")
+			_, err = client.GetSandboxWithLastActivity(ctx, sandbox.SessionID)
+			require.ErrorIs(t, err, ErrNotFound)
+			_, err = server.ZAdd("sandbox:last_activity", float64(oldActivity.Unix()), sandbox.SessionID)
+			require.NoError(t, err)
+			server.Del("session:" + sandbox.SessionID)
+			_, err = client.GetSandboxWithLastActivity(ctx, sandbox.SessionID)
+			require.ErrorIs(t, err, ErrNotFound)
+			require.NoError(t, server.Set("session:"+sandbox.SessionID, "invalid JSON"))
+			_, err = client.GetSandboxWithLastActivity(ctx, sandbox.SessionID)
+			require.Error(t, err)
+			require.NotErrorIs(t, err, ErrNotFound)
+		})
+	}
+}
+
+func TestGCStoreErrors(t *testing.T) {
+	for _, provider := range []string{"redis", "valkey"} {
+		t.Run(provider, func(t *testing.T) {
+			client, server := newGCStore(t, provider)
+			ctx, cancel := context.WithCancel(context.Background())
+			cancel()
+			_, next, err := client.ScanInactiveSandboxIDs(ctx, time.Now(), 13, 100)
+			require.ErrorIs(t, err, context.Canceled)
+			require.Equal(t, uint64(13), next)
+			_, err = client.GetSandboxWithLastActivity(ctx, "session")
+			require.ErrorIs(t, err, context.Canceled)
+
+			ctx = context.Background()
+			require.NoError(t, client.StoreSandbox(ctx, newTestSandbox("sb", "session", time.Now().Add(time.Hour))))
+			server.Del("sandbox:last_activity")
+			require.NoError(t, server.Set("sandbox:last_activity", "wrong type"))
+			_, next, err = client.ScanInactiveSandboxIDs(ctx, time.Now(), 13, 100)
+			require.Error(t, err)
+			require.Equal(t, uint64(13), next)
+			_, err = client.GetSandboxWithLastActivity(ctx, "session")
+			require.Error(t, err)
+			require.NotErrorIs(t, err, ErrNotFound)
+		})
+	}
+}

pkg/store/interface.go

@@ -38,6 +38,15 @@ type Store interface {
 	ListExpiredSandboxes(ctx context.Context, before time.Time, limit int64) ([]*types.SandboxInfo, error)
 	// ListInactiveSandboxes returns up to limit sandboxes with last-activity time before the given time
 	ListInactiveSandboxes(ctx context.Context, before time.Time, limit int64) ([]*types.SandboxInfo, error)
+	// ScanInactiveSandboxIDs scans one page of the activity index, returning IDs
+	// with scores at or before before.Unix(). Start with cursor zero and continue
+	// with the returned cursor until it is zero again. An empty page is not
+	// necessarily the end. count must be positive and is a work hint, not a limit;
+	// results are unordered and may contain duplicates or missing sessions.
+	ScanInactiveSandboxIDs(ctx context.Context, before time.Time, cursor uint64, count int64) ([]string, uint64, error)
+	// GetSandboxWithLastActivity reads the session and its current activity score.
+	// It returns ErrNotFound if the session or activity index entry is missing.
+	GetSandboxWithLastActivity(ctx context.Context, sessionID string) (*types.SandboxInfo, error)
 	// UpdateSessionLastActivity updates the last-activity index for the given session
 	UpdateSessionLastActivity(ctx context.Context, sessionID string, at time.Time) error
 	// Close releases all resources held by the store (e.g. connection pools)

pkg/store/store_redis.go

@@ -332,3 +332,36 @@ func (rs *redisStore) UpdateSessionLastActivity(ctx context.Context, sessionID s
 
 	return nil
 }
+
+// ScanInactiveSandboxIDs advances an opaque cursor over the activity index.
+func (rs *redisStore) ScanInactiveSandboxIDs(ctx context.Context, before time.Time, cursor uint64, count int64) ([]string, uint64, error) {
+	if count <= 0 {
+		return nil, cursor, fmt.Errorf("ScanInactiveSandboxIDs: count must be positive")
+	}
+	elements, next, err := rs.cli.ZScan(ctx, rs.lastActivityIndexKey, cursor, "", count).Result()
+	if err != nil {
+		return nil, cursor, fmt.Errorf("ScanInactiveSandboxIDs: ZSCAN failed: %w", err)
+	}
+	ids, err := inactiveIDsFromScan(elements, before)
+	if err != nil {
+		return nil, cursor, err
+	}
+	return ids, next, nil
+}
+
+// GetSandboxWithLastActivity refreshes activity after reading the session JSON.
+func (rs *redisStore) GetSandboxWithLastActivity(ctx context.Context, sessionID string) (*types.SandboxInfo, error) {
+	sandbox, err := rs.GetSandboxBySessionID(ctx, sessionID)
+	if err != nil {
+		return nil, err
+	}
+	score, err := rs.cli.ZScore(ctx, rs.lastActivityIndexKey, sessionID).Result()
+	if errors.Is(err, redisv9.Nil) {
+		return nil, ErrNotFound
+	}
+	if err != nil {
+		return nil, fmt.Errorf("GetSandboxWithLastActivity: ZSCORE failed: %w", err)
+	}
+	sandbox.LastActivityAt = time.Unix(int64(score), 0)
+	return sandbox, nil
+}

pkg/store/store_valkey.go

@@ -327,3 +327,36 @@ func (vs *valkeyStore) UpdateSessionLastActivity(ctx context.Context, sessionID
 	}
 	return nil
 }
+
+// ScanInactiveSandboxIDs advances an opaque cursor over the activity index.
+func (vs *valkeyStore) ScanInactiveSandboxIDs(ctx context.Context, before time.Time, cursor uint64, count int64) ([]string, uint64, error) {
+	if count <= 0 {
+		return nil, cursor, fmt.Errorf("ScanInactiveSandboxIDs: count must be positive")
+	}
+	page, err := vs.cli.Do(ctx, vs.cli.B().Zscan().Key(vs.lastActivityIndexKey).Cursor(cursor).Count(count).Build()).AsScanEntry()
+	if err != nil {
+		return nil, cursor, fmt.Errorf("ScanInactiveSandboxIDs: ZSCAN failed: %w", err)
+	}
+	ids, err := inactiveIDsFromScan(page.Elements, before)
+	if err != nil {
+		return nil, cursor, err
+	}
+	return ids, page.Cursor, nil
+}
+
+// GetSandboxWithLastActivity refreshes activity after reading the session JSON.
+func (vs *valkeyStore) GetSandboxWithLastActivity(ctx context.Context, sessionID string) (*types.SandboxInfo, error) {
+	sandbox, err := vs.GetSandboxBySessionID(ctx, sessionID)
+	if err != nil {
+		return nil, err
+	}
+	score, err := vs.cli.Do(ctx, vs.cli.B().Zscore().Key(vs.lastActivityIndexKey).Member(sessionID).Build()).AsFloat64()
+	if valkey.IsValkeyNil(err) {
+		return nil, ErrNotFound
+	}
+	if err != nil {
+		return nil, fmt.Errorf("GetSandboxWithLastActivity: ZSCORE failed: %w", err)
+	}
+	sandbox.LastActivityAt = time.Unix(int64(score), 0)
+	return sandbox, nil
+}