[dss/RID] GetISA: optional forUpdate flag to allow for early locking

interuss · Sep 12, 2024 · 29c0874 · 29c0874
1 parent 97bbbe5
commit 29c0874
Show file tree

Hide file tree

Showing 6 changed files with 19 additions and 10 deletions.
diff --git a/pkg/rid/application/isa.go b/pkg/rid/application/isa.go
@@ -38,7 +38,7 @@ func (a *app) GetISA(ctx context.Context, id dssmodels.ID) (*ridmodels.Identific
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "Unable to interact with store")
 	}
-	return repo.GetISA(ctx, id)
+	return repo.GetISA(ctx, id, false)
 }
 
 // SearchISAs for ISA within the volume bounds.
@@ -64,7 +64,7 @@ func (a *app) DeleteISA(ctx context.Context, id dssmodels.ID, owner dssmodels.Ow
 	)
 	// The following will automatically retry TXN retry errors.
 	err := a.Store.Transact(ctx, func(repo repos.Repository) error {
-		old, err := repo.GetISA(ctx, id)
+		old, err := repo.GetISA(ctx, id, true)
 		switch {
 		case err != nil:
 			return stacktrace.Propagate(err, "Error getting ISA")
@@ -106,7 +106,7 @@ func (a *app) InsertISA(ctx context.Context, isa *ridmodels.IdentificationServic
 	// The following will automatically retry TXN retry errors.
 	err := a.Store.Transact(ctx, func(repo repos.Repository) error {
 		// ensure it doesn't exist yet
-		old, err := repo.GetISA(ctx, isa.ID)
+		old, err := repo.GetISA(ctx, isa.ID, false)
 		if err != nil {
 			return stacktrace.Propagate(err, "Error getting ISA")
 		}
@@ -141,7 +141,7 @@ func (a *app) UpdateISA(ctx context.Context, isa *ridmodels.IdentificationServic
 	err := a.Store.Transact(ctx, func(repo repos.Repository) error {
 		var err error
 
-		old, err := repo.GetISA(ctx, isa.ID)
+		old, err := repo.GetISA(ctx, isa.ID, true)
 		switch {
 		case err != nil:
 			return stacktrace.Propagate(err, "Error getting ISA")

diff --git a/pkg/rid/application/isa_test.go b/pkg/rid/application/isa_test.go
@@ -31,7 +31,7 @@ type isaStore struct {
 	isas map[dssmodels.ID]*ridmodels.IdentificationServiceArea
 }
 
-func (store *isaStore) GetISA(ctx context.Context, id dssmodels.ID) (*ridmodels.IdentificationServiceArea, error) {
+func (store *isaStore) GetISA(ctx context.Context, id dssmodels.ID, forUpdate bool) (*ridmodels.IdentificationServiceArea, error) {
 	if isa, ok := store.isas[id]; ok {
 		return isa, nil
 	}

diff --git a/pkg/rid/repos/isa.go b/pkg/rid/repos/isa.go
@@ -12,7 +12,7 @@ import (
 // ISA is an interface to a storage layer for the ISA entity
 type ISA interface {
 	// Returns nil, nil if not found
-	GetISA(ctx context.Context, id dssmodels.ID) (*ridmodels.IdentificationServiceArea, error)
+	GetISA(ctx context.Context, id dssmodels.ID, forUpdate bool) (*ridmodels.IdentificationServiceArea, error)
 
 	// DeleteISA deletes the IdentificationServiceArea identified by "id" and owned by "owner".
 	// Returns the delete IdentificationServiceArea and all Subscriptions affected by the delete.

diff --git a/pkg/rid/store/cockroach/identification_service_area.go b/pkg/rid/store/cockroach/identification_service_area.go
@@ -101,12 +101,13 @@ func (c *isaRepo) processOne(ctx context.Context, query string, args ...interfac
 
 // GetISA returns the isa identified by "id".
 // Returns nil, nil if not found
-func (c *isaRepo) GetISA(ctx context.Context, id dssmodels.ID) (*ridmodels.IdentificationServiceArea, error) {
+func (c *isaRepo) GetISA(ctx context.Context, id dssmodels.ID, forUpdate bool) (*ridmodels.IdentificationServiceArea, error) {
 	var query = fmt.Sprintf(`
 		SELECT %s FROM
 			identification_service_areas
 		WHERE
-			id = $1`, isaFields)
+			id = $1
+        %s`, isaFields, dssql.ForUpdate(forUpdate))
 	uid, err := id.PgUUID()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "Failed to convert id to PgUUID")

diff --git a/pkg/rid/store/cockroach/identification_service_area_v3.go b/pkg/rid/store/cockroach/identification_service_area_v3.go
@@ -82,12 +82,13 @@ func (c *isaRepoV3) processOne(ctx context.Context, query string, args ...interf
 
 // GetISA returns the isa identified by "id".
 // Returns nil, nil if not found
-func (c *isaRepoV3) GetISA(ctx context.Context, id dssmodels.ID) (*ridmodels.IdentificationServiceArea, error) {
+func (c *isaRepoV3) GetISA(ctx context.Context, id dssmodels.ID, forUpdate bool) (*ridmodels.IdentificationServiceArea, error) {
 	var query = fmt.Sprintf(`
 		SELECT %s FROM
 			identification_service_areas
 		WHERE
-			id = $1`, isaFieldsV3)
+			id = $1
+        %s`, isaFieldsV3, dssql.ForUpdate(forUpdate))
 	uid, err := id.PgUUID()
 	if err != nil {
 		return nil, stacktrace.Propagate(err, "Failed to convert id to PgUUID")

diff --git a/pkg/sql/utils.go b/pkg/sql/utils.go
@@ -26,3 +26,10 @@ func CellUnionToCellIdsWithValidation(cu s2.CellUnion) ([]int64, error) {
 	}
 	return pgCids, nil
 }
+
+func ForUpdate(forUpdate bool) string {
+	if forUpdate {
+		return "FOR UPDATE"
+	}
+	return ""
+}