perf(postgres): bulk cache session ownership

Identity preparation previously issued serialized owner lookups for each candidate and repeated many of them during alias and relationship resolution. The latency dominated full pushes to remote PostgreSQL even when later fingerprint comparison skipped the session.\n\nPreload every candidate, legacy-prefix, canonical, and alias owner in one array query. Cache both rows and absences for the full push so out-of-window relationship lookups are queried at most once.

Author: wesm

internal/postgres/push.go

@@ -14,6 +14,7 @@ import (
 	"slices"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	"go.kenn.io/agentsview/internal/artifact"
@@ -354,6 +355,13 @@ func (s *Sync) Push(
 	if err != nil {
 		return result, err
 	}
+	ctx, err = s.preloadPGSessionOwners(ctx, s.pushIdentityOwnerCandidateIDs(
+		sessionByID, artifactImportedSessions, localArtifactOrigin,
+		markerID, legacyMarkerMachines,
+	))
+	if err != nil {
+		return result, err
+	}
 	for id, sess := range sessionByID {
 		_, artifactImported := artifactImportedSessions[sess.ID]
 		identity, err := s.resolvePushedSessionIdentity(
@@ -1855,6 +1863,31 @@ func (s *Sync) markRelationshipConflicts(
 	return nil
 }
 
+func initialPushedSessionIdentity(
+	sess db.Session,
+	fallbackMachine string,
+	localArtifactOrigin string,
+	artifactImported bool,
+	markerID string,
+) (pushedSessionIdentity, bool) {
+	identity := pushedSessionIdentity{
+		ID:      sess.ID,
+		Machine: pushedSessionMachine(sess, fallbackMachine),
+	}
+	if id, machine, ownerMarker, ok := artifactPushIdentity(
+		sess, localArtifactOrigin, artifactImported,
+	); ok {
+		identity.ID = id
+		identity.Machine = machine
+		identity.OwnerMarker = ownerMarker
+		identity.LegacyOwnerMarkers = []string{markerID}
+		identity.ArtifactReplica = artifactImported
+		identity.AliasIDs = artifactPushAliasIDs(sess, id, machine)
+		return identity, true
+	}
+	return identity, false
+}
+
 // resolvePushedSessionIdentity decides the PG id a local session is stored
 // under. A session this sync owns -- by matching push marker, or an adoptable
 // legacy/ownerless row (see sameSessionOwner) -- is updated in place: an
@@ -1873,22 +1906,9 @@ func (s *Sync) resolvePushedSessionIdentity(
 	markerID string,
 	legacyMarkerMachines []string,
 ) (pushedSessionIdentity, error) {
-	identity := pushedSessionIdentity{
-		ID:      sess.ID,
-		Machine: pushedSessionMachine(sess, s.machine),
-	}
-	artifactIdentity := false
-	if id, machine, ownerMarker, ok := artifactPushIdentity(
-		sess, localArtifactOrigin, artifactImported,
-	); ok {
-		artifactIdentity = true
-		identity.ID = id
-		identity.Machine = machine
-		identity.OwnerMarker = ownerMarker
-		identity.LegacyOwnerMarkers = []string{markerID}
-		identity.ArtifactReplica = artifactImported
-		identity.AliasIDs = artifactPushAliasIDs(sess, id, machine)
-	}
+	identity, artifactIdentity := initialPushedSessionIdentity(
+		sess, s.machine, localArtifactOrigin, artifactImported, markerID,
+	)
 	canonicalID := identity.ID
 	id, err := s.resolveOwnedPushIdentityID(
 		ctx, identity.ID, identity, markerID, legacyMarkerMachines,
@@ -1920,6 +1940,42 @@ func (s *Sync) resolvePushedSessionIdentity(
 	return identity, nil
 }
 
+func (s *Sync) pushIdentityOwnerCandidateIDs(
+	sessionByID map[string]db.Session,
+	artifactImportedSessions map[string]struct{},
+	localArtifactOrigin string,
+	markerID string,
+	legacyMarkerMachines []string,
+) []string {
+	ids := make(map[string]struct{}, len(sessionByID)*3)
+	for _, sess := range sessionByID {
+		_, artifactImported := artifactImportedSessions[sess.ID]
+		identity, _ := initialPushedSessionIdentity(
+			sess, s.machine, localArtifactOrigin, artifactImported, markerID,
+		)
+		ids[identity.ID] = struct{}{}
+		for _, machine := range pushIDMachinePrefixes(
+			identity.Machine, legacyMarkerMachines,
+		) {
+			candidateID := prefixedSessionID(machine, identity.ID)
+			if candidateID != identity.ID {
+				ids[candidateID] = struct{}{}
+			}
+		}
+		for _, aliasID := range uniqueNonEmptyStrings(identity.AliasIDs) {
+			ids[aliasID] = struct{}{}
+		}
+	}
+	result := make([]string, 0, len(ids))
+	for id := range ids {
+		if id != "" {
+			result = append(result, id)
+		}
+	}
+	sort.Strings(result)
+	return result
+}
+
 // artifactLegacyDuplicateCandidate recognizes the narrow upgrade state where
 // an importer already created the stable artifact id while this origin still
 // owns its pre-artifact bare row. Both rows must already have the exact owners
@@ -2076,13 +2132,85 @@ func pushIDMachinePrefixes(machine string, legacyMarkerMachines []string) []stri
 	return prefixes
 }
 
+type pgSessionOwnerRecord struct {
+	machine     string
+	ownerMarker string
+	exists      bool
+}
+
+type pgSessionOwnerCache struct {
+	mu      sync.Mutex
+	entries map[string]pgSessionOwnerRecord
+}
+
+type pgSessionOwnerCacheContextKey struct{}
+
+func (c *pgSessionOwnerCache) get(id string) (pgSessionOwnerRecord, bool) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	record, ok := c.entries[id]
+	return record, ok
+}
+
+func (c *pgSessionOwnerCache) put(id string, record pgSessionOwnerRecord) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.entries[id] = record
+}
+
+// preloadPGSessionOwners resolves a candidate set in one PG round trip and
+// records both hits and misses. pgSessionOwner reuses this cache and memoizes
+// any relationship targets discovered later in the same push.
+func (s *Sync) preloadPGSessionOwners(
+	ctx context.Context, ids []string,
+) (context.Context, error) {
+	unique := uniqueNonEmptyStrings(ids)
+	cache := &pgSessionOwnerCache{
+		entries: make(map[string]pgSessionOwnerRecord, len(unique)),
+	}
+	for _, id := range unique {
+		cache.entries[id] = pgSessionOwnerRecord{}
+	}
+	if len(unique) == 0 {
+		return context.WithValue(ctx, pgSessionOwnerCacheContextKey{}, cache), nil
+	}
+	rows, err := s.pg.QueryContext(ctx, `
+		SELECT id, machine, owner_marker
+		FROM sessions
+		WHERE id = ANY($1)`, unique)
+	if err != nil {
+		return ctx, fmt.Errorf("preloading pg session owners: %w", err)
+	}
+	defer rows.Close()
+	for rows.Next() {
+		var id, machine string
+		var ownerMarker sql.NullString
+		if err := rows.Scan(&id, &machine, &ownerMarker); err != nil {
+			return ctx, fmt.Errorf("scanning pg session owner: %w", err)
+		}
+		cache.entries[id] = pgSessionOwnerRecord{
+			machine: machine, ownerMarker: ownerMarker.String, exists: true,
+		}
+	}
+	if err := rows.Err(); err != nil {
+		return ctx, fmt.Errorf("iterating pg session owners: %w", err)
+	}
+	return context.WithValue(ctx, pgSessionOwnerCacheContextKey{}, cache), nil
+}
+
 // pgSessionOwner returns the machine and owner_marker of a PG session row, and
 // whether it exists. owner_marker is empty for legacy rows pushed before the
 // marker model.
 func (s *Sync) pgSessionOwner(
 	ctx context.Context,
 	id string,
 ) (string, string, bool, error) {
+	cache, _ := ctx.Value(pgSessionOwnerCacheContextKey{}).(*pgSessionOwnerCache)
+	if cache != nil {
+		if record, ok := cache.get(id); ok {
+			return record.machine, record.ownerMarker, record.exists, nil
+		}
+	}
 	var machine string
 	var ownerMarker sql.NullString
 	err := s.pg.QueryRowContext(ctx,
@@ -2091,13 +2219,21 @@ func (s *Sync) pgSessionOwner(
 	).Scan(&machine, &ownerMarker)
 	if err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
+			if cache != nil {
+				cache.put(id, pgSessionOwnerRecord{})
+			}
 			return "", "", false, nil
 		}
 		return "", "", false, fmt.Errorf(
 			"reading pg session owner for %s: %w",
 			id, err,
 		)
 	}
+	if cache != nil {
+		cache.put(id, pgSessionOwnerRecord{
+			machine: machine, ownerMarker: ownerMarker.String, exists: true,
+		})
+	}
 	return machine, ownerMarker.String, true, nil
 }
 

internal/postgres/push_test.go

@@ -1194,6 +1194,13 @@ type pushSessionProbeState struct {
 	aliases             map[string]string
 	excludedIDs         map[string]bool
 	existingExcluded    map[string]bool
+	ownerQueries        int
+	owners              map[string]pushSessionProbeOwner
+}
+
+type pushSessionProbeOwner struct {
+	machine string
+	marker  string
 }
 
 var (
@@ -1310,7 +1317,29 @@ func (c *pushSessionProbeConn) QueryContext(
 	defer c.state.mu.Unlock()
 
 	switch {
+	case strings.Contains(normalized, "select id, machine, owner_marker"):
+		c.state.ownerQueries++
+		values := [][]driver.Value{}
+		for _, id := range namedValueStrings(args) {
+			if owner, ok := c.state.owners[id]; ok {
+				values = append(values, []driver.Value{id, owner.machine, owner.marker})
+			}
+		}
+		return &pushSessionProbeRows{
+			columns: []string{"id", "machine", "owner_marker"},
+			values:  values,
+		}, nil
 	case strings.Contains(normalized, "select machine, owner_marker"):
+		c.state.ownerQueries++
+		if len(args) > 0 {
+			id, _ := args[0].Value.(string)
+			if owner, ok := c.state.owners[id]; ok {
+				return &pushSessionProbeRows{
+					columns: []string{"machine", "owner_marker"},
+					values:  [][]driver.Value{{owner.machine, owner.marker}},
+				}, nil
+			}
+		}
 		return &pushSessionProbeRows{
 			columns: []string{"machine", "owner_marker"},
 		}, nil
@@ -1347,6 +1376,71 @@ func (c *pushSessionProbeConn) QueryContext(
 	}
 }
 
+func TestPreloadPGSessionOwnersUsesOneQueryAndCachesMisses(t *testing.T) {
+	state := &pushSessionProbeState{owners: map[string]pushSessionProbeOwner{
+		"owned-a": {machine: "desk", marker: "marker-a"},
+		"owned-b": {machine: "laptop", marker: "marker-b"},
+	}}
+	sync := &Sync{pg: newPushSessionProbeDB(t, state)}
+
+	ctx, err := sync.preloadPGSessionOwners(
+		context.Background(), []string{"owned-a", "owned-b", "missing"},
+	)
+	require.NoError(t, err)
+	for _, tc := range []struct {
+		id      string
+		machine string
+		marker  string
+		exists  bool
+	}{
+		{id: "owned-a", machine: "desk", marker: "marker-a", exists: true},
+		{id: "owned-b", machine: "laptop", marker: "marker-b", exists: true},
+		{id: "missing"},
+	} {
+		machine, marker, exists, lookupErr := sync.pgSessionOwner(ctx, tc.id)
+		require.NoError(t, lookupErr)
+		assert.Equal(t, tc.machine, machine)
+		assert.Equal(t, tc.marker, marker)
+		assert.Equal(t, tc.exists, exists)
+	}
+	assert.Equal(t, 1, state.ownerQueries,
+		"preloaded hits and misses must use one owner query")
+
+	_, _, exists, err := sync.pgSessionOwner(ctx, "late-miss")
+	require.NoError(t, err)
+	assert.False(t, exists)
+	_, _, exists, err = sync.pgSessionOwner(ctx, "late-miss")
+	require.NoError(t, err)
+	assert.False(t, exists)
+	assert.Equal(t, 2, state.ownerQueries,
+		"an owner first discovered after preload must be memoized")
+}
+
+func TestPushIdentityOwnerCandidateIDsCoverLegacyAndArtifactAliases(t *testing.T) {
+	sync := &Sync{machine: "desk"}
+	sessions := map[string]db.Session{
+		"plain": {
+			ID: "plain", Machine: "local",
+		},
+		"remote-a1b2c3~imported": {
+			ID: "remote-a1b2c3~imported", Machine: "remote-a1b2c3",
+		},
+	}
+	imported := map[string]struct{}{"remote-a1b2c3~imported": {}}
+
+	got := sync.pushIdentityOwnerCandidateIDs(
+		sessions, imported, "desk-origin", "marker", []string{"old-desk"},
+	)
+	assert.ElementsMatch(t, []string{
+		"desk-origin~plain",
+		"old-desk~desk-origin~plain",
+		"plain",
+		"remote-a1b2c3~imported",
+		"old-desk~remote-a1b2c3~imported",
+		"imported",
+	}, got)
+}
+
 func (pushSessionProbeTx) Commit() error { return nil }
 
 func (pushSessionProbeTx) Rollback() error { return nil }