add condition to check lock

gravitational · Apr 24, 2024 · 4a60e0d · 4a60e0d
1 parent 20dbe4a
commit 4a60e0d
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 3 deletions.
diff --git a/lib/backend/helpers.go b/lib/backend/helpers.go
@@ -37,7 +37,7 @@ func FlagKey(parts ...string) []byte {
 	return internalKey(flagsPrefix, parts...)
 }
 
-func lockKey(parts ...string) []byte {
+func LockKey(parts ...string) []byte {
 	return internalKey(locksPrefix, parts...)
 }
 
@@ -88,7 +88,7 @@ func AcquireLock(ctx context.Context, cfg LockConfiguration) (Lock, error) {
 	if err != nil {
 		return Lock{}, trace.Wrap(err)
 	}
-	key := lockKey(cfg.LockName)
+	key := LockKey(cfg.LockName)
 	id, err := randomID()
 	if err != nil {
 		return Lock{}, trace.Wrap(err)

diff --git a/lib/services/local/externalauditstorage.go b/lib/services/local/externalauditstorage.go
@@ -20,6 +20,7 @@ package local
 
 import (
 	"context"
+	"time"
 
 	"github.com/gravitational/trace"
 	"github.com/sirupsen/logrus"
@@ -35,6 +36,8 @@ const (
 	externalAuditStoragePrefix      = "external_audit_storage"
 	externalAuditStorageDraftName   = "draft"
 	externalAuditStorageClusterName = "cluster"
+	externalAuditStorageLockName    = "external_audit_storage_lock"
+	externalAuditStorageLockTTL     = 10 * time.Second
 )
 
 var (
@@ -93,6 +96,13 @@ func (s *ExternalAuditStorageService) CreateDraftExternalAuditStorage(ctx contex
 	}
 
 	revision, err := s.backend.AtomicWrite(ctx, []backend.ConditionalAction{
+		{
+			// Make sure another auth server on an older minor/patch version is not holding the lock that was
+			// used before this switched to AtomicWrite.
+			Key:       backend.LockKey(externalAuditStorageLockName),
+			Condition: backend.NotExists(),
+			Action:    backend.Nop(),
+		},
 		{
 			// Make sure the AWS OIDC integration checked above hasn't changed.
 			Key:       integrationKey,
@@ -133,6 +143,13 @@ func (s *ExternalAuditStorageService) UpsertDraftExternalAuditStorage(ctx contex
 	}
 
 	revision, err := s.backend.AtomicWrite(ctx, []backend.ConditionalAction{
+		{
+			// Make sure another auth server on an older minor/patch version is not holding the lock that was
+			// used before this switched to AtomicWrite.
+			Key:       backend.LockKey(externalAuditStorageLockName),
+			Condition: backend.NotExists(),
+			Action:    backend.Nop(),
+		},
 		{
 			// Make sure the AWS OIDC integration checked above hasn't changed.
 			Key:       integrationKey,
@@ -218,6 +235,13 @@ func (s *ExternalAuditStorageService) PromoteToClusterExternalAuditStorage(ctx c
 	}
 
 	_, err = s.backend.AtomicWrite(ctx, []backend.ConditionalAction{
+		{
+			// Make sure another auth server on an older minor/patch version is not holding the lock that was
+			// used before this switched to AtomicWrite.
+			Key:       backend.LockKey(externalAuditStorageLockName),
+			Condition: backend.NotExists(),
+			Action:    backend.Nop(),
+		},
 		{
 			// Make sure the AWS OIDC integration checked above hasn't changed.
 			Key:       integrationKey,

diff --git a/lib/services/local/integrations.go b/lib/services/local/integrations.go
@@ -143,7 +143,13 @@ func (s *IntegrationsService) DeleteIntegration(ctx context.Context, name string
 // notReferencedByEAS returns a slice of ConditionalActions to use with a backend.AtomicWrite to ensure that
 // integration [name] is not referenced by any EAS (External Audit Storage) integration.
 func notReferencedByEAS(ctx context.Context, bk backend.Backend, name string) ([]backend.ConditionalAction, error) {
-	var conditionalActions []backend.ConditionalAction
+	conditionalActions := []backend.ConditionalAction{{
+		// Make sure another auth server on an older minor/patch version is not holding the lock that was used
+		// before this switched to AtomicWrite.
+		Key:       backend.LockKey(externalAuditStorageLockName),
+		Condition: backend.NotExists(),
+		Action:    backend.Nop(),
+	}}
 	for _, key := range [][]byte{draftExternalAuditStorageBackendKey, clusterExternalAuditStorageBackendKey} {
 		condition := backend.ConditionalAction{
 			Key:    key,