A hash trie based pebble differ

pkg/dotc1z/engine/pebble/adapter.go

@@ -230,6 +230,18 @@ func (a *Adapter) EndSync(ctx context.Context) error {
 			zap.Error(err),
 		)
 	}
+	// Build the per-entitlement grant digests at seal time, after all
+	// grants are written and while still on the fresh-sync NoSync path
+	// (the EndFreshSync flush below hardens the nodes). Non-fatal, like
+	// the stats sidecar: a missing/stale digest only forces a diff
+	// consumer onto the on-demand index fold, and the on-Open migration
+	// backfills it next time the file opens writable.
+	if err := a.engine.BuildAllGrantDigests(ctx, existing.GetSyncId()); err != nil {
+		ctxzap.Extract(ctx).Warn("pebble: build grant digests failed; grant-diff callers will fall back to on-demand index folds until the next Open backfills them",
+			zap.String("sync_id", existing.GetSyncId()),
+			zap.Error(err),
+		)
+	}
 	// Single flush + WAL fsync at sync end. This is the durability
 	// boundary — counterpart to MarkFreshSync at StartNewSync. After
 	// this returns, all writes from the sync are on disk.

pkg/dotc1z/engine/pebble/adapter_clone_sync.go

@@ -110,6 +110,8 @@ func cloneSync(
 		{GrantByPrincipalSyncLowerBound(syncIDBytes), GrantByPrincipalSyncUpperBound(syncIDBytes)},
 		{GrantByPrincipalResourceTypeSyncLowerBound(syncIDBytes), GrantByPrincipalResourceTypeSyncUpperBound(syncIDBytes)},
 		{GrantByNeedsExpansionSyncLowerBound(syncIDBytes), GrantByNeedsExpansionSyncUpperBound(syncIDBytes)},
+		{GrantByEntPrincHashSyncLowerBound(syncIDBytes), GrantByEntPrincHashSyncUpperBound(syncIDBytes)},
+		{DigestSyncLowerBound(syncIDBytes), DigestSyncUpperBound(syncIDBytes)},
 		{encodeAssetPrefix(syncIDBytes), upperBoundOf(encodeAssetPrefix(syncIDBytes))},
 		// Stats sidecar — single key per sync; copyRange's [lo, hi)
 		// shape requires a half-open range, so we synthesize one

pkg/dotc1z/engine/pebble/adapter_diff.go

@@ -23,13 +23,17 @@ import (
 // deletions are NOT captured — that matches the SQLite behavior
 // in pkg/dotc1z/diff.go (additions-only diff).
 //
-// Strategy: for each record-type-bearing keyspace under
-// appliedSyncID, iterate the source keys, recompute the same
-// record's primary key under baseSyncID, Get from base; if base
-// returns ErrNotFound, write the value under diffSyncID's
-// keyspace in the same record type. Index entries are
-// recomputed on write so the diff sync has its own (fresh)
-// indexes that match the records that landed.
+// Strategy: for the small record types (resource_types, resources,
+// entitlements, assets), iterate the source keys under appliedSyncID,
+// recompute the same record's primary key under baseSyncID, Get from
+// base; if base returns ErrNotFound, write the value under
+// diffSyncID's keyspace in the same record type. GRANTS — the type
+// that dominates every real file — are diffed via the per-entitlement
+// grant digests instead (see diffGrants): unchanged entitlements are
+// skipped with a single root comparison, and only the principal-hash
+// buckets that actually differ are scanned. Index entries are
+// recomputed on write so the diff sync has its own (fresh) indexes
+// that match the records that landed.
 //
 // Returns the diff sync's ID.
 func generateSyncDiff(ctx context.Context, a *Adapter, baseSyncID, appliedSyncID string) (string, error) {
@@ -106,7 +110,7 @@ func generateSyncDiff(ctx context.Context, a *Adapter, baseSyncID, appliedSyncID
 	if err := diffEntitlements(ctx, a, baseBytes, appliedBytes, diffSyncID); err != nil {
 		return "", fmt.Errorf("generate-diff: entitlements: %w", err)
 	}
-	if err := diffGrants(ctx, a, baseBytes, appliedBytes, diffSyncID); err != nil {
+	if err := diffGrants(ctx, a, baseBytes, appliedBytes, baseSyncID, appliedSyncID, diffSyncID); err != nil {
 		return "", fmt.Errorf("generate-diff: grants: %w", err)
 	}
 	if err := diffAssets(ctx, a, baseBytes, appliedBytes, diffSyncID); err != nil {
@@ -190,7 +194,84 @@ func diffEntitlements(ctx context.Context, a *Adapter, baseBytes, appliedBytes [
 	})
 }
 
-func diffGrants(ctx context.Context, a *Adapter, baseBytes, appliedBytes []byte, diffSyncID string) error {
+// diffGrants computes the grants set difference using the
+// per-entitlement grant digests instead of scanning every applied
+// grant.
+//
+// Walk: enumerate the distinct entitlement_ids in the applied sync's
+// hash index (one seek each), compare each entitlement's digest between
+// base and applied — one root read per side when nothing changed, the
+// overwhelmingly common case — and materialize only the principal-hash
+// buckets the comparison flags as dirty. Grants in dirty buckets are
+// probed against base's PRIMARY keyspace by external_id, which
+// preserves the additions-only contract exactly: a grant whose
+// external_id exists in base under a different entitlement/principal
+// is still "present in base" (not emitted), which is why the probe
+// targets the primary key rather than comparing index keys.
+//
+// Why pruning is sound: equal (count, digest) node pairs mean the two
+// sides hold identical grant content-hash sets, and the content hash
+// folds external_id — so a clean entitlement/bucket cannot contain an
+// applied external_id that base lacks. Dirtiness over-approximates
+// additions (it also fires for removals and source-set changes, which
+// the base probe then filters out), never under-approximates them.
+//
+// NOTE: grants without an entitlement or principal ref have no
+// hash-index entry and are invisible to the digest. They are silently
+// skipped. A future O(1) coverage check (e.g. a stored grant count in
+// the sync run record) will restore detection of this case.
+func diffGrants(ctx context.Context, a *Adapter, baseBytes, appliedBytes []byte, baseSyncID, appliedSyncID, diffSyncID string) error {
+	eng := a.engine
+
+	ents, err := eng.distinctDigestPartitions(ctx, grantDigestSpec, appliedBytes)
+	if err != nil {
+		return err
+	}
+	for _, ent := range ents {
+		if err := ctx.Err(); err != nil {
+			return err
+		}
+		// Entitlements only present in base (fully removed) are not in
+		// `ents` and are correctly skipped: they cannot contain
+		// additions. Digests missing on either side (e.g. a ghost
+		// entitlement with grants but no entitlement record) degrade
+		// to an on-demand index fold inside DirtyEntitlementBuckets.
+		dirty, err := eng.DirtyEntitlementBuckets(ctx, baseSyncID, eng, appliedSyncID, ent)
+		if err != nil {
+			return err
+		}
+		for _, bucket := range dirty {
+			var innerErr error
+			err := eng.IterateGrantsByEntitlementBucket(ctx, appliedSyncID, ent, bucket, func(rec *v3.GrantRecord) bool {
+				exists, probeErr := existsAt(eng.DB(), encodeGrantKey(baseBytes, rec.GetExternalId()))
+				if probeErr != nil {
+					innerErr = probeErr
+					return false
+				}
+				if exists {
+					return true
+				}
+				if putErr := eng.PutGrantRecord(ctx, rec); putErr != nil {
+					innerErr = putErr
+					return false
+				}
+				return true
+			})
+			if err != nil {
+				return err
+			}
+			if innerErr != nil {
+				return innerErr
+			}
+		}
+	}
+	return nil
+}
+
+// diffGrantsFullScan is the O(applied grants) path: iterate every
+// applied grant, probe base by external_id, emit on miss. Used
+// directly by benchmarks and available as a correctness reference.
+func diffGrantsFullScan(ctx context.Context, a *Adapter, baseBytes, appliedBytes []byte, diffSyncID string) error {
 	srcPrefix := encodeGrantPrefix(appliedBytes)
 	return iterDiff(ctx, a.engine.DB(), srcPrefix, upperBoundOf(srcPrefix), func(_ []byte, val []byte) error {
 		var rec v3.GrantRecord

pkg/dotc1z/engine/pebble/adapter_diff_trie_test.go

@@ -0,0 +1,201 @@
+package pebble
+
+import (
+	"context"
+	"testing"
+
+	"github.com/cockroachdb/pebble/v2"
+
+	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
+	v3 "github.com/conductorone/baton-sdk/pb/c1/storage/v3"
+	"github.com/conductorone/baton-sdk/pkg/connectorstore"
+)
+
+func mkV2Ent(id string) *v2.Entitlement {
+	return v2.Entitlement_builder{
+		Id: id,
+		Resource: v2.Resource_builder{
+			Id: v2.ResourceId_builder{
+				ResourceType: "app",
+				Resource:     "github",
+			}.Build(),
+		}.Build(),
+	}.Build()
+}
+
+// runSync starts a sync, writes the given entitlements + grants, and
+// seals it (EndSync builds the per-entitlement tries). Returns the
+// sync id.
+func runSync(t *testing.T, a *Adapter, ents []*v2.Entitlement, grants []*v2.Grant) string {
+	t.Helper()
+	ctx := context.Background()
+	syncID, err := a.StartNewSync(ctx, connectorstore.SyncTypeFull, "")
+	if err != nil {
+		t.Fatalf("StartNewSync: %v", err)
+	}
+	if len(ents) > 0 {
+		if err := a.PutEntitlements(ctx, ents...); err != nil {
+			t.Fatalf("PutEntitlements: %v", err)
+		}
+	}
+	if len(grants) > 0 {
+		if err := a.PutGrants(ctx, grants...); err != nil {
+			t.Fatalf("PutGrants: %v", err)
+		}
+	}
+	if err := a.EndSync(ctx); err != nil {
+		t.Fatalf("EndSync: %v", err)
+	}
+	return syncID
+}
+
+// diffGrantIDs runs generateSyncDiff and returns the set of grant
+// external_ids that landed in the diff sync.
+func diffGrantIDs(t *testing.T, a *Adapter, baseSyncID, appliedSyncID string) map[string]bool {
+	t.Helper()
+	ctx := context.Background()
+	diffID, err := generateSyncDiff(ctx, a, baseSyncID, appliedSyncID)
+	if err != nil {
+		t.Fatalf("generateSyncDiff: %v", err)
+	}
+	got := map[string]bool{}
+	if err := a.engine.IterateGrantsBySync(ctx, diffID, func(r *v3.GrantRecord) bool {
+		got[r.GetExternalId()] = true
+		return true
+	}); err != nil {
+		t.Fatalf("IterateGrantsBySync(diff): %v", err)
+	}
+	return got
+}
+
+func requireDiffIDs(t *testing.T, got map[string]bool, want ...string) {
+	t.Helper()
+	wantSet := map[string]bool{}
+	for _, id := range want {
+		wantSet[id] = true
+		if !got[id] {
+			t.Errorf("diff missing expected grant %q", id)
+		}
+	}
+	for id := range got {
+		if !wantSet[id] {
+			t.Errorf("diff contains unexpected grant %q", id)
+		}
+	}
+}
+
+// TestGenerateSyncDiffTrieMultiEntitlement exercises the trie-driven
+// grants diff across the full semantic matrix in one pair of syncs:
+// unchanged entitlements (pruned at the root), an addition inside an
+// existing entitlement, an addition under a brand-new entitlement, a
+// removal (must not be emitted), and a grant that MOVED to a different
+// entitlement while keeping its external_id (dirty bucket on both
+// entitlements, but per the additions-only-by-external_id contract it
+// must not be emitted).
+func TestGenerateSyncDiffTrieMultiEntitlement(t *testing.T) {
+	a := newAdapter(t)
+
+	base := runSync(t, a,
+		[]*v2.Entitlement{mkV2Ent("ent-A"), mkV2Ent("ent-B"), mkV2Ent("ent-C")},
+		[]*v2.Grant{
+			mkV2Grant("g1", "ent-A", "user", "alice"),
+			mkV2Grant("g2", "ent-A", "user", "bob"),
+			mkV2Grant("g3", "ent-B", "user", "carol"),
+			mkV2Grant("g4", "ent-B", "user", "dave"),
+			mkV2Grant("g5", "ent-C", "user", "eve"),
+			mkV2Grant("g6", "ent-A", "user", "frank"),
+		})
+	applied := runSync(t, a,
+		[]*v2.Entitlement{mkV2Ent("ent-A"), mkV2Ent("ent-B"), mkV2Ent("ent-C"), mkV2Ent("ent-D")},
+		[]*v2.Grant{
+			mkV2Grant("g1", "ent-A", "user", "alice"), // unchanged
+			mkV2Grant("g2", "ent-A", "user", "bob"),   // unchanged
+			mkV2Grant("g3", "ent-B", "user", "carol"), // unchanged
+			mkV2Grant("g4", "ent-B", "user", "dave"),  // unchanged
+			// g5 removed (ent-C empty in applied) — removals not emitted.
+			mkV2Grant("g6", "ent-B", "user", "frank"), // moved ent-A→ent-B, same external_id — NOT an addition
+			mkV2Grant("g7", "ent-B", "user", "grace"), // addition in an existing entitlement
+			mkV2Grant("g8", "ent-D", "user", "heidi"), // addition under a new entitlement
+		})
+
+	got := diffGrantIDs(t, a, base, applied)
+	requireDiffIDs(t, got, "g7", "g8")
+}
+
+// TestGenerateSyncDiffOrphanGrantSkipped documents that a grant
+// without entitlement/principal refs has no hash-index entry and is
+// invisible to the trie-driven diff path. It is silently skipped.
+// Normal grants in the same sync (g2) are still emitted correctly.
+// TODO: restore detection via an O(1) coverage check (e.g. a stored
+// grant count in the sync run record).
+func TestGenerateSyncDiffOrphanGrantSkipped(t *testing.T) {
+	ctx := context.Background()
+	a := newAdapter(t)
+
+	base := runSync(t, a,
+		[]*v2.Entitlement{mkV2Ent("ent-A")},
+		[]*v2.Grant{mkV2Grant("g1", "ent-A", "user", "alice")})
+
+	applied, err := a.StartNewSync(ctx, connectorstore.SyncTypeFull, "")
+	if err != nil {
+		t.Fatalf("StartNewSync: %v", err)
+	}
+	if err := a.PutEntitlements(ctx, mkV2Ent("ent-A")); err != nil {
+		t.Fatalf("PutEntitlements: %v", err)
+	}
+	if err := a.PutGrants(ctx,
+		mkV2Grant("g1", "ent-A", "user", "alice"),
+		mkV2Grant("g2", "ent-A", "user", "bob"),
+	); err != nil {
+		t.Fatalf("PutGrants: %v", err)
+	}
+	orphan := v3.GrantRecord_builder{
+		ExternalId: "g-orphan",
+	}.Build()
+	if err := a.engine.PutGrantRecord(ctx, orphan); err != nil {
+		t.Fatalf("PutGrantRecord(orphan): %v", err)
+	}
+	if err := a.EndSync(ctx); err != nil {
+		t.Fatalf("EndSync: %v", err)
+	}
+
+	got := diffGrantIDs(t, a, base, applied)
+	// g-orphan has no hash-index entry — silently skipped by trie diff.
+	requireDiffIDs(t, got, "g2")
+}
+
+// TestGenerateSyncDiffTrieWithoutTrees simulates a file whose tries
+// are missing (e.g. written before they existed and not yet
+// backfilled): the hash index is intact, so the trie path runs, and
+// DirtyEntitlementBuckets degrades to the on-demand index fold per
+// entitlement. The diff must still be exact.
+func TestGenerateSyncDiffTrieWithoutTrees(t *testing.T) {
+	a := newAdapter(t)
+
+	base := runSync(t, a,
+		[]*v2.Entitlement{mkV2Ent("ent-A"), mkV2Ent("ent-B")},
+		[]*v2.Grant{
+			mkV2Grant("g1", "ent-A", "user", "alice"),
+			mkV2Grant("g2", "ent-B", "user", "bob"),
+		})
+	applied := runSync(t, a,
+		[]*v2.Entitlement{mkV2Ent("ent-A"), mkV2Ent("ent-B")},
+		[]*v2.Grant{
+			mkV2Grant("g1", "ent-A", "user", "alice"),
+			mkV2Grant("g2", "ent-B", "user", "bob"),
+			mkV2Grant("g3", "ent-B", "user", "carol"),
+		})
+
+	for _, sid := range []string{base, applied} {
+		idBytes, err := a.engine.resolveSyncBytes(sid)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if err := a.engine.db.DeleteRange(DigestSyncLowerBound(idBytes), DigestSyncUpperBound(idBytes), pebble.Sync); err != nil {
+			t.Fatalf("DeleteRange(digest %s): %v", sid, err)
+		}
+	}
+
+	got := diffGrantIDs(t, a, base, applied)
+	requireDiffIDs(t, got, "g3")
+}

pkg/dotc1z/engine/pebble/cleanup.go

@@ -35,6 +35,8 @@ func syncScopedRanges(syncIDBytes []byte) [][2][]byte {
 		{GrantByPrincipalSyncLowerBound(syncIDBytes), GrantByPrincipalSyncUpperBound(syncIDBytes)},
 		{GrantByPrincipalResourceTypeSyncLowerBound(syncIDBytes), GrantByPrincipalResourceTypeSyncUpperBound(syncIDBytes)},
 		{GrantByNeedsExpansionSyncLowerBound(syncIDBytes), GrantByNeedsExpansionSyncUpperBound(syncIDBytes)},
+		{GrantByEntPrincHashSyncLowerBound(syncIDBytes), GrantByEntPrincHashSyncUpperBound(syncIDBytes)},
+		{DigestSyncLowerBound(syncIDBytes), DigestSyncUpperBound(syncIDBytes)},
 		{encodeAssetPrefix(syncIDBytes), upperBoundOf(encodeAssetPrefix(syncIDBytes))},
 		// Stats sidecar — single key per sync; the half-open range
 		// shape contains exactly the one key for this sync.