[DO NOT MERGE] refactor: CachedOptimizedRevisions

internal/datastore/crdb/crdb.go

@@ -164,23 +164,16 @@
 		keyer = noOverlapKeyer
 	}
 
-	maxRevisionStaleness := time.Duration(float64(config.revisionQuantization.Nanoseconds())*
-		config.maxRevisionStalenessPercent) * time.Nanosecond
-
 	headMigration, err := migrations.CRDBMigrations.HeadRevision()
 	if err != nil {
 		return nil, fmt.Errorf("invalid head migration found for cockroach: %w", err)
 	}
 
 	ds := &crdbDatastore{
-		RemoteClockRevisions: revisions.NewRemoteClockRevisions(
-			config.gcWindow,
-			maxRevisionStaleness,
-			config.followerReadDelay,
-			config.revisionQuantization,
-		),
 		CommonDecoder:                revisions.CommonDecoder{Kind: revisions.HybridLogicalClock},
 		MigrationValidator:           common.NewMigrationValidator(headMigration, config.allowedMigrations),
+		followerReadDelay:            config.followerReadDelay,
+		revisionQuantization:         config.revisionQuantization,
 		dburl:                        url,
 		acquireTimeout:               config.acquireTimeout,
 		watchBufferLength:            config.watchBufferLength,
@@ -198,8 +191,6 @@
 		watchEnabled:                 !config.watchDisabled,
 		schema:                       *schema.Schema(config.columnOptimizationOption, config.withIntegrity, false),
 	}
-	ds.SetNowFunc(ds.headRevisionInternal)
-	ds.SetNowOnlyFunc(ds.headRevisionInternalNoHash)
 
 	// this ctx and cancel is tied to the lifetime of the datastore
 	ds.ctx, ds.cancel = context.WithCancel(context.Background())
@@ -257,11 +248,68 @@
 	return datastore.NewSeparatingContextDatastoreProxy(ds), nil
 }
 
+// OptimizedRevision computes the uncached quantized revision for CockroachDB. It
+// reads the current HLC time and quantizes it in Go (CRDB's only revision
+// primitive is cluster_logical_timestamp()). Caching, deduplication, and jitter
+// are layered on by proxy.NewOptimizedRevisionProxy.
+func (cds *crdbDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
+	ctx, span := tracer.Start(ctx, "OptimizedRevision")
+	defer span.End()
+
+	nowRev, schemaHash, err := cds.headRevisionInternal(ctx)
+	if err != nil {
+		return datastore.NoRevision, 0, "", err
+	}
+
+	if nowRev == datastore.NoRevision {
+		return datastore.NoRevision, 0, "", datastore.NewInvalidRevisionErr(nowRev, datastore.CouldNotDetermineRevision)
+	}
+
+	nowTS, ok := nowRev.(revisions.WithTimestampRevision)
+	if !ok {
+		return datastore.NoRevision, 0, "", spiceerrors.MustBugf("expected with-timestamp revision, got %T", nowRev)
+	}
+
+	rev, validFor := revisions.QuantizeHLC(nowTS, cds.followerReadDelay, cds.revisionQuantization)
+	return rev, validFor, schemaHash, nil
+}
+
+// CheckRevision verifies the given revision is within the software GC window. It
+// lives on the datastore (not the caching layer) because it is not a caching
+// concern. It uses the hash-free head revision query since the schema hash is
+// not needed here.
+func (cds *crdbDatastore) CheckRevision(ctx context.Context, dsRevision datastore.Revision) error {
+	if dsRevision == datastore.NoRevision {
+		return datastore.NewInvalidRevisionErr(dsRevision, datastore.CouldNotDetermineRevision)
+	}
+
+	revision, ok := dsRevision.(revisions.WithTimestampRevision)
+	if !ok {
+		return spiceerrors.MustBugf("expected HLC revision, got %T", dsRevision)
+	}
+
+	ctx, span := tracer.Start(ctx, "CheckRevision")
+	defer span.End()
+
+	now, err := cds.headRevisionInternalNoHash(ctx)
+	if err != nil {
+		return err
+	}
+
+	nowTS, ok := now.(revisions.WithTimestampRevision)
+	if !ok {
+		return spiceerrors.MustBugf("expected HLC revision, got %T", now)
+	}
+
+	return revisions.CheckHLCGCWindow(nowTS, revision, cds.gcWindow)
+}
+
 type crdbDatastore struct {
-	*revisions.RemoteClockRevisions
 	revisions.CommonDecoder
 	*common.MigrationValidator
 
+	followerReadDelay            time.Duration
+	revisionQuantization         time.Duration
 	dburl                        string
 	readPool, writePool          *pool.RetryPool
 	collectors                   []prometheus.Collector

internal/datastore/crdb/crdb_test.go

@@ -167,13 +167,13 @@ func TestCRDBDatastoreWithFollowerReads(t *testing.T) {
 
 			// Revisions should be at least the follower read delay amount in the past
 			for start := time.Now(); time.Since(start) < 50*time.Millisecond; {
-				testRevisionResult, err := ds.OptimizedRevision(ctx)
+				testRevision, _, _, err := ds.OptimizedRevision(ctx)
 				require.NoError(t, err)
 
 				nowRevisionResult, err := ds.HeadRevision(ctx)
 				require.NoError(t, err)
 
-				diff := nowRevisionResult.Revision.(revisions.HLCRevision).TimestampNanoSec() - testRevisionResult.Revision.(revisions.HLCRevision).TimestampNanoSec()
+				diff := nowRevisionResult.Revision.(revisions.HLCRevision).TimestampNanoSec() - testRevision.(revisions.HLCRevision).TimestampNanoSec()
 				require.Greater(t, diff, followerReadDelay.Nanoseconds())
 			}
 		})

internal/datastore/memdb/memdb_test.go

@@ -140,7 +140,7 @@ func TestAnythingAfterCloseDoesNotPanic(t *testing.T) {
 	err = ds.CheckRevision(t.Context(), lowestRevision.Revision)
 	require.ErrorIs(err, ErrMemDBIsClosed)
 
-	_, err = ds.OptimizedRevision(t.Context())
+	_, _, _, err = ds.OptimizedRevision(t.Context())
 	require.ErrorIs(err, ErrMemDBIsClosed)
 
 	reader := ds.SnapshotReader(datastore.NoRevision)

internal/datastore/memdb/revisions.go

@@ -68,17 +68,20 @@ func (mdb *memdbDatastore) headRevisionNoLock() revisions.TimestampRevision {
 	return mdb.revisions[len(mdb.revisions)-1].revision
 }
 
-func (mdb *memdbDatastore) OptimizedRevision(_ context.Context) (datastore.RevisionWithSchemaHash, error) {
+func (mdb *memdbDatastore) OptimizedRevision(_ context.Context) (datastore.Revision, time.Duration, string, error) {
 	mdb.RLock()
 	defer mdb.RUnlock()
 	if err := mdb.checkNotClosed(); err != nil {
-		return datastore.RevisionWithSchemaHash{}, err
+		return datastore.NoRevision, 0, "", err
 	}
 
 	now := nowRevision()
 	var optimized revisions.TimestampRevision
+	var validFor time.Duration
 	if mdb.quantizationPeriod > 0 {
-		optimized = revisions.NewForTimestamp(now.TimestampNanoSec() - now.TimestampNanoSec()%mdb.quantizationPeriod)
+		afterLastQuantization := now.TimestampNanoSec() % mdb.quantizationPeriod
+		optimized = revisions.NewForTimestamp(now.TimestampNanoSec() - afterLastQuantization)
+		validFor = time.Duration(mdb.quantizationPeriod-afterLastQuantization) * time.Nanosecond
 	} else {
 		optimized = now
 	}
@@ -94,7 +97,7 @@ func (mdb *memdbDatastore) OptimizedRevision(_ context.Context) (datastore.Revis
 		}
 	}
 
-	return datastore.RevisionWithSchemaHash{Revision: optimized, SchemaHash: hash}, nil
+	return optimized, validFor, hash, nil
 }
 
 func (mdb *memdbDatastore) CheckRevision(_ context.Context, dr datastore.Revision) error {

internal/datastore/mysql/datastore.go

@@ -201,9 +201,6 @@ func newMySQLDatastore(ctx context.Context, uri string, replicaIndex int, option
 
 	gcCtx, cancelGc := context.WithCancel(context.Background())
 
-	maxRevisionStaleness := time.Duration(float64(config.revisionQuantization.Nanoseconds())*
-		config.maxRevisionStalenessPercent) * time.Nanosecond
-
 	quantizationPeriodNanos := max(config.revisionQuantization.Nanoseconds(), 1)
 
 	followerReadDelayNanos := max(config.followerReadDelay.Nanoseconds(), 0)
@@ -269,17 +266,12 @@ func newMySQLDatastore(ctx context.Context, uri string, replicaIndex int, option
 		maxRetries:                   config.maxRetries,
 		analyzeBeforeStats:           config.analyzeBeforeStats,
 		schema:                       *schema,
-		CachedOptimizedRevisions: revisions.NewCachedOptimizedRevisions(
-			maxRevisionStaleness,
-		),
 		CommonDecoder: revisions.CommonDecoder{
 			Kind: revisions.TransactionID,
 		},
 		filterMaximumIDCount: config.filterMaximumIDCount,
 	}
 
-	store.SetOptimizedRevisionFunc(store.optimizedRevisionFunc)
-
 	ctx, cancel := context.WithTimeout(context.Background(), seedingTimeout)
 	defer cancel()
 	err = store.seedDatabase(ctx)
@@ -475,7 +467,6 @@ func newMySQLExecutor(tx querier, explainable datastore.Explainable) common.Exec
 
 // Datastore is a MySQL-based implementation of the datastore.Datastore interface
 type mysqlDatastore struct {
-	*revisions.CachedOptimizedRevisions
 	*common.MigrationValidator
 
 	db                 *sql.DB

internal/datastore/mysql/datastore_test.go

@@ -764,9 +764,9 @@ func TransactionTimestampsTest(t *testing.T, ds datastore.Datastore) {
 	// Let's make sure both Now() and transactionCreated() have timezones aligned
 	req.Less(ts.Sub(startTimeUTC), 5*time.Minute)
 
-	revisionResult, err := ds.OptimizedRevision(ctx)
+	rev, _, _, err := ds.OptimizedRevision(ctx)
 	req.NoError(err)
-	req.Equal(revisions.NewForTransactionID(txID), revisionResult.Revision)
+	req.Equal(revisions.NewForTransactionID(txID), rev)
 }
 
 func TestMySQLMigrations(t *testing.T) {

internal/datastore/mysql/revisions.go

@@ -76,7 +76,7 @@ const (
 		) as unknown;`
 )
 
-func (mds *mysqlDatastore) optimizedRevisionFunc(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
+func (mds *mysqlDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
 	var rev uint64
 	var validForNanos time.Duration
 	var schemaHash []byte

internal/datastore/postgres/postgres.go

@@ -29,7 +29,6 @@ import (
 	pgxcommon "github.com/authzed/spicedb/internal/datastore/postgres/common"
 	"github.com/authzed/spicedb/internal/datastore/postgres/migrations"
 	"github.com/authzed/spicedb/internal/datastore/postgres/schema"
-	"github.com/authzed/spicedb/internal/datastore/revisions"
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/sharederrors"
 	"github.com/authzed/spicedb/pkg/datastore"
@@ -287,9 +286,6 @@ func newPostgresDatastore(
 		schema.ColSnapshot,
 	)
 
-	maxRevisionStaleness := time.Duration(float64(config.revisionQuantization.Nanoseconds())*
-		config.maxRevisionStalenessPercent) * time.Nanosecond
-
 	isolationLevel := pgx.Serializable
 	if config.relaxedIsolationLevel {
 		isolationLevel = pgx.RepeatableRead
@@ -298,9 +294,6 @@ func newPostgresDatastore(
 	startGarbageCollector := datastore.StartGarbageCollector
 
 	datastore := &pgDatastore{
-		CachedOptimizedRevisions: revisions.NewCachedOptimizedRevisions(
-			maxRevisionStaleness,
-		),
 		MigrationValidator:           common.NewMigrationValidator(headMigration, config.allowedMigrations),
 		dburl:                        pgURL,
 		readPool:                     pgxcommon.MustNewInterceptorPooler(readPool, config.queryInterceptor),
@@ -338,8 +331,6 @@ func newPostgresDatastore(
 		datastore.writePool = pgxcommon.MustNewInterceptorPooler(writePool, config.queryInterceptor)
 	}
 
-	datastore.SetOptimizedRevisionFunc(datastore.optimizedRevisionFunc)
-
 	// Start a goroutine for garbage collection and the revision heartbeat.
 	if isPrimary {
 		datastore.workerGroup, datastore.workerCtx = errgroup.WithContext(datastore.workerCtx)
@@ -368,7 +359,6 @@ func newPostgresDatastore(
 }
 
 type pgDatastore struct {
-	*revisions.CachedOptimizedRevisions
 	*common.MigrationValidator
 
 	dburl                          string

internal/datastore/postgres/postgres_shared_test.go

@@ -1045,7 +1045,7 @@ func assertRevisionLowerAndHigher(ctx context.Context, t *testing.T, ds datastor
 	var snapshot pgSnapshot
 	pgDS, ok := ds.(*pgDatastore)
 	require.True(t, ok)
-	rev, _, _, err := pgDS.optimizedRevisionFunc(ctx)
+	rev, _, _, err := pgDS.OptimizedRevision(ctx)
 	require.NoError(t, err)
 
 	pgRev, ok := rev.(postgresRevision)

internal/datastore/postgres/revisions.go

@@ -128,7 +128,7 @@ const (
 	queryLatestXID            = `SELECT max(xid)::text::integer FROM relation_tuple_transaction;`
 )
 
-func (pgd *pgDatastore) optimizedRevisionFunc(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
+func (pgd *pgDatastore) OptimizedRevision(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
 	var revision xid8
 	var snapshot pgSnapshot
 	var validForNanos time.Duration

internal/datastore/proxy/checkingreplicated_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
 
@@ -179,8 +180,8 @@ func (f fakeDatastore) ReadWriteTx(_ context.Context, _ datastore.TxUserFunc, _
 	return nil, nil
 }
 
-func (f fakeDatastore) OptimizedRevision(_ context.Context) (datastore.RevisionWithSchemaHash, error) {
-	return datastore.RevisionWithSchemaHash{}, nil
+func (f fakeDatastore) OptimizedRevision(_ context.Context) (datastore.Revision, time.Duration, string, error) {
+	return datastore.NoRevision, 0, "", nil
 }
 
 func (f fakeDatastore) HeadRevision(_ context.Context) (datastore.RevisionWithSchemaHash, error) {

internal/datastore/proxy/indexcheck/fakedatastore_test.go

@@ -3,6 +3,7 @@ package indexcheck
 import (
 	"context"
 	"fmt"
+	"time"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
@@ -39,8 +40,8 @@ func (f fakeDatastore) ReadWriteTx(ctx context.Context, fn datastore.TxUserFunc,
 	})
 }
 
-func (f fakeDatastore) OptimizedRevision(_ context.Context) (datastore.RevisionWithSchemaHash, error) {
-	return datastore.RevisionWithSchemaHash{}, nil
+func (f fakeDatastore) OptimizedRevision(_ context.Context) (datastore.Revision, time.Duration, string, error) {
+	return nil, 0, "", nil
 }
 
 func (f fakeDatastore) HeadRevision(_ context.Context) (datastore.RevisionWithSchemaHash, error) {

internal/datastore/proxy/indexcheck/indexcheck.go

@@ -3,6 +3,7 @@ package indexcheck
 import (
 	"context"
 	"fmt"
+	"time"
 
 	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
 
@@ -56,7 +57,7 @@ func (p *indexcheckingProxy) UniqueID(ctx context.Context) (string, error) {
 	return p.delegate.UniqueID(ctx)
 }
 
-func (p *indexcheckingProxy) OptimizedRevision(ctx context.Context) (datastore.RevisionWithSchemaHash, error) {
+func (p *indexcheckingProxy) OptimizedRevision(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
 	return p.delegate.OptimizedRevision(ctx)
 }
 

internal/datastore/proxy/indexcheck/indexcheck_test.go

@@ -97,9 +97,9 @@ func TestIndexCheckingProxyMethods(t *testing.T) {
 	})
 
 	t.Run("OptimizedRevision", func(t *testing.T) {
-		result, err := proxy.OptimizedRevision(t.Context())
+		rev, _, _, err := proxy.OptimizedRevision(t.Context())
 		require.NoError(t, err)
-		require.Nil(t, result.Revision)
+		require.Nil(t, rev)
 	})
 
 	t.Run("CheckRevision", func(t *testing.T) {

internal/datastore/proxy/observable.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"context"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
@@ -98,7 +99,7 @@ func (p *observableProxy) ReadWriteTx(
 	}, opts...)
 }
 
-func (p *observableProxy) OptimizedRevision(ctx context.Context) (datastore.RevisionWithSchemaHash, error) {
+func (p *observableProxy) OptimizedRevision(ctx context.Context) (datastore.Revision, time.Duration, string, error) {
 	ctx, closer := observe(ctx, "OptimizedRevision", "")
 	defer closer()
 

internal/datastore/proxy/observable_test.go

@@ -3,6 +3,7 @@ package proxy
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/prometheus/client_golang/prometheus"
 	dto "github.com/prometheus/client_model/go"
@@ -74,10 +75,10 @@ func TestObservableProxy_DatastoreMethodsWithMetrics(t *testing.T) {
 			name:     "OptimizedRevision",
 			metricOp: "OptimizedRevision",
 			setupMock: func(ds *proxy_test.MockDatastore) {
-				ds.On("OptimizedRevision").Return(datastore.RevisionWithSchemaHash{Revision: testRev}, nil).Once()
+				ds.On("OptimizedRevision").Return(testRev, time.Duration(0), "", nil).Once()
 			},
 			call: func(t *testing.T, ds datastore.Datastore) {
-				_, err := ds.OptimizedRevision(t.Context())
+				_, _, _, err := ds.OptimizedRevision(t.Context())
 				require.NoError(t, err)
 			},
 		},