Fix lock behavior

Author: yolken-segment

pkg/store/lock.go

@@ -82,8 +82,8 @@ type KubeLocker struct {
 	namespace          string
 	objLock            sync.Mutex
 	lockCancellations  map[string]context.CancelFunc
+	lockCompletions    map[string]chan struct{}
 	coordinationClient coordv1.CoordinationV1Interface
-	done               chan struct{}
 }
 
 // NewKubeLocker returns a Locker that is backed by a lock in Kubernetes.
@@ -102,14 +102,13 @@ func NewKubeLocker(
 	}
 
 	coordinationClient := client.CoordinationV1()
-	done := make(chan struct{}, 1)
 
 	return &KubeLocker{
 		id:                 id,
 		namespace:          namespace,
 		lockCancellations:  map[string]context.CancelFunc{},
+		lockCompletions:    map[string]chan struct{}{},
 		coordinationClient: coordinationClient,
-		done:               done,
 	}, nil
 }
 
@@ -124,6 +123,7 @@ func (k *KubeLocker) Acquire(ctx context.Context, name string) error {
 	// Create a separate context for the lock itself
 	lockCtx, lockCancel := context.WithCancel(context.Background())
 	k.lockCancellations[name] = lockCancel
+	k.lockCompletions[name] = make(chan struct{}, 1)
 	k.objLock.Unlock()
 
 	leaseName := fmt.Sprintf("kubeapply-lock-%s", name)
@@ -154,8 +154,10 @@ func (k *KubeLocker) Acquire(ctx context.Context, name string) error {
 					acquired <- struct{}{}
 				},
 				OnStoppedLeading: func() {
+					k.objLock.Lock()
+					defer k.objLock.Unlock()
 					log.Warn("Lock lost")
-					k.done <- struct{}{}
+					k.lockCompletions[name] <- struct{}{}
 				},
 			},
 		},
@@ -180,30 +182,43 @@ func (k *KubeLocker) Acquire(ctx context.Context, name string) error {
 
 // Release releases the lock with the argument name.
 func (k *KubeLocker) Release(name string) error {
-	k.objLock.Lock()
-	defer k.objLock.Unlock()
-
-	log.Infof("Releasing lock with name %s", name)
-
-	cancel, ok := k.lockCancellations[name]
-	if !ok {
-		return fmt.Errorf("Lock was not acquired for name %s", name)
+	// Do this in a separate function to prevent deadlock on k.objLock.
+	if err := k.releaseHelper(name); err != nil {
+		return err
 	}
 
-	cancel()
-	delete(k.lockCancellations, name)
-
 	log.Infof("Waiting for lock to be released")
 	releaseCtx, releaseCancel := context.WithTimeout(
 		context.Background(),
 		kubeLockerReleaseTimeout,
 	)
 	defer releaseCancel()
 
+	k.objLock.Lock()
+	completionChan := k.lockCompletions[name]
+	k.objLock.Unlock()
+
 	select {
-	case <-k.done:
+	case <-completionChan:
+		delete(k.lockCompletions, name)
 		return nil
 	case <-releaseCtx.Done():
 		return releaseCtx.Err()
 	}
 }
+
+func (k *KubeLocker) releaseHelper(name string) error {
+	k.objLock.Lock()
+	defer k.objLock.Unlock()
+
+	log.Infof("Releasing lock with name %s", name)
+
+	cancel, ok := k.lockCancellations[name]
+	if !ok {
+		return fmt.Errorf("Lock was not acquired for name %s", name)
+	}
+
+	cancel()
+	delete(k.lockCancellations, name)
+	return nil
+}