diff --git a/pkg/rid/application/isa.go b/pkg/rid/application/isa.go index ea402af60..1c2445bb2 100644 --- a/pkg/rid/application/isa.go +++ b/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). diff --git a/pkg/rid/application/subscription_test.go b/pkg/rid/application/subscription_test.go index e83fee2ff..cfcac3f5b 100644 --- a/pkg/rid/application/subscription_test.go +++ b/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) diff --git a/pkg/rid/repos/subscription.go b/pkg/rid/repos/subscription.go index 1ec849f72..aa0a823e0 100644 --- a/pkg/rid/repos/subscription.go +++ b/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 } diff --git a/pkg/rid/store/raftstore/subscriptions.go b/pkg/rid/store/raftstore/subscriptions.go index 0c9c261dd..5415c385d 100644 --- a/pkg/rid/store/raftstore/subscriptions.go +++ b/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") +} diff --git a/pkg/rid/store/sqlstore/subscriptions.go b/pkg/rid/store/sqlstore/subscriptions.go index 219cb4630..e39fbef70 100644 --- a/pkg/rid/store/sqlstore/subscriptions.go +++ b/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 +}