[rid] Add lock on subscriptions

pkg/rid/application/isa.go

@@ -78,6 +78,12 @@ func (a *app) DeleteISA(ctx context.Context, id dssmodels.ID, owner dssmodels.Ow
 				"ISA owned by %s, but %s attempted to delete", old.Owner, owner)
 		}
 
+		err = repo.LockSubscriptionsOnCells(ctx, old.Cells)
+
+		if err != nil {
+			return stacktrace.Propagate(err, "Error locking cells")
+		}
+
 		ret, err = repo.DeleteISA(ctx, old)
 		if err != nil {
 			return stacktrace.Propagate(err, "Error deleting ISA")
@@ -105,6 +111,15 @@ func (a *app) InsertISA(ctx context.Context, isa *ridmodels.IdentificationServic
 	)
 	// The following will automatically retry TXN retry errors.
 	err := a.store.Transact(ctx, func(ctx context.Context, repo repos.Repository) error {
+
+		var err error
+
+		err = repo.LockSubscriptionsOnCells(ctx, isa.Cells)
+
+		if err != nil {
+			return stacktrace.Propagate(err, "Error locking cells")
+		}
+
 		// ensure it doesn't exist yet
 		old, err := repo.GetISA(ctx, isa.ID, false)
 		if err != nil {
@@ -139,6 +154,7 @@ func (a *app) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServic
 	)
 	// The following will automatically retry TXN retry errors.
 	err := a.store.Transact(ctx, func(ctx context.Context, repo repos.Repository) error {
+
 		var err error
 
 		old, err := repo.GetISA(ctx, isa.ID, true)
@@ -159,15 +175,24 @@ func (a *app) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServic
 			return stacktrace.Propagate(err, "Error adjusting time range")
 		}
 
+		// TODO steeling, we should change this to a Custom type, to obfuscate
+		// some of these metrics and prevent us from doing the wrong thing.
+		cells := s2.CellUnionFromUnion(old.Cells, isa.Cells)
+		geo.Levelify(&cells)
+
+		// NB: There is a tradeoff between locking first without old cells or now.
+		// Best option hasn't been valided, this just follow SCD logic that does the same
+		err = repo.LockSubscriptionsOnCells(ctx, cells)
+
+		if err != nil {
+			return stacktrace.Propagate(err, "Error locking cells")
+		}
+
 		ret, err = repo.UpdateISA(ctx, isa)
 		if err != nil {
 			return stacktrace.Propagate(err, "Error updating ISA")
 		}
 
-		// TODO steeling, we should change this to a Custom type, to obfuscate
-		// some of these metrics and prevent us from doing the wrong thing.
-		cells := s2.CellUnionFromUnion(old.Cells, isa.Cells)
-		geo.Levelify(&cells)
 		// UpdateNotificationIdxsInCells is done in a Txn along with insert since
 		// they are both modifying the db. Insert a susbcription alone does
 		// not do this, so that does not need to use a txn (in subscription.go).

pkg/rid/application/subscription_test.go

@@ -167,6 +167,11 @@ func (store *subscriptionStore) CountSubscriptions(ctx context.Context) (int64,
 	return int64(len(store.subs)), nil
 }
 
+// Don't perform anything
+func (store *subscriptionStore) LockSubscriptionsOnCells(ctx context.Context, cells s2.CellUnion) error {
+	return nil
+}
+
 func TestBadOwner(t *testing.T) {
 	ctx := context.Background()
 	app, cleanup := setUpSubApp(ctx, t)

pkg/rid/repos/subscription.go

@@ -44,4 +44,7 @@ type Subscription interface {
 
 	// Count the number of existing subscriptions
 	CountSubscriptions(ctx context.Context) (int64, error)
+
+	// LockSubscriptionsOnCells locks the subscriptions of interest on specific cells
+	LockSubscriptionsOnCells(ctx context.Context, cells s2.CellUnion) error
 }

pkg/rid/store/raftstore/subscriptions.go

@@ -50,3 +50,7 @@ func (r *repo) ListExpiredSubscriptions(_ context.Context, writer string, thresh
 func (r *repo) CountSubscriptions(_ context.Context) (int64, error) {
 	return 0, stacktrace.NewErrorWithCode(dsserr.NotImplemented, "CountSubscriptions not implemented for raftstore")
 }
+
+func (r *repo) LockSubscriptionsOnCells(_ context.Context, cells s2.CellUnion) error {
+	return stacktrace.NewErrorWithCode(dsserr.NotImplemented, "LockSubscriptionsOnCells not implemented for raftstore")
+}

pkg/rid/store/sqlstore/subscriptions.go

@@ -6,10 +6,12 @@ import (
 	"time"
 
 	dsserr "github.com/interuss/dss/pkg/errors"
+	"github.com/interuss/dss/pkg/logging"
 	dssmodels "github.com/interuss/dss/pkg/models"
 	ridmodels "github.com/interuss/dss/pkg/rid/models"
 	dssql "github.com/interuss/dss/pkg/sql"
 	"github.com/interuss/stacktrace"
+	"go.uber.org/zap"
 
 	"github.com/golang/geo/s2"
 	"github.com/jackc/pgx/v5/pgtype"
@@ -18,6 +20,8 @@ import (
 const (
 	subscriptionFields       = "id, owner, url, notification_index, cells, starts_at, ends_at, writer, updated_at"
 	updateSubscriptionFields = "id, url, notification_index, cells, starts_at, ends_at, writer, updated_at"
+	// This threshold keep lock diagnostics low-noise in production while still surfacing likely bottlenecks.
+	lockQuerySlowThreshold = 4 * time.Second
 )
 
 // process a query that should return one or many subscriptions.
@@ -304,3 +308,43 @@ func (r *repo) CountSubscriptions(ctx context.Context) (int64, error) {
 	err := r.QueryRow(ctx, "SELECT COUNT(*) FROM subscriptions").Scan(&count)
 	return count, err
 }
+
+func (r *repo) LockSubscriptionsOnCells(ctx context.Context, cells s2.CellUnion) error {
+	logger := logging.WithValuesFromContext(ctx, logging.Logger)
+
+	const query = `
+        SELECT
+            id
+        FROM
+            subscriptions
+        WHERE
+			(
+				cells && $1
+				AND ends_at >= $2
+			)
+        FOR UPDATE
+    `
+
+	start := time.Now()
+	_, err := r.Exec(ctx, query, dssql.CellUnionToCellIds(cells), r.clock.Now())
+	duration := time.Since(start)
+
+	if err != nil {
+		logger.Warn("RID subscription lock query failed",
+			zap.Duration("duration", duration),
+			zap.Int("cell_count", len(cells)),
+			zap.Error(err),
+		)
+		return stacktrace.Propagate(err, "Error in query: %s", query)
+	}
+
+	if duration >= lockQuerySlowThreshold {
+		logger.Warn("Expensive RID lock detected",
+			zap.Bool("global_lock", false),
+			zap.Duration("duration", duration),
+			zap.Int("cell_count", len(cells)),
+		)
+	}
+
+	return nil
+}