Merge pull request #44 from ActiDoo/codex/find-node-reboot-prevention-mechanism

Require linearizable contexts for etcd lock operations

Author: msander

pkg/cooldown/etcd.go

@@ -94,7 +94,9 @@ func (m *EtcdManager) Status(ctx context.Context) (Status, error) {
 		ctx = context.Background()
 	}
 
-	resp, err := m.client.Get(ctx, m.key)
+	linearizableCtx := clientv3.WithRequireLeader(ctx)
+
+	resp, err := m.client.Get(linearizableCtx, m.key)
 	if err != nil {
 		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
 			return Status{}, err
@@ -122,7 +124,8 @@ func (m *EtcdManager) Status(ctx context.Context) (Status, error) {
 		status.ExpiresAt = startedAt.Add(time.Duration(record.DurationSec) * time.Second)
 	}
 	if kv.Lease != 0 {
-		ttlResp, ttlErr := m.client.TimeToLive(ctx, clientv3.LeaseID(kv.Lease))
+		ttlCtx := clientv3.WithRequireLeader(ctx)
+		ttlResp, ttlErr := m.client.TimeToLive(ttlCtx, clientv3.LeaseID(kv.Lease))
 		if ttlErr != nil {
 			if errors.Is(ttlErr, context.Canceled) || errors.Is(ttlErr, context.DeadlineExceeded) {
 				return Status{}, ttlErr

pkg/lock/etcd_manager.go

@@ -127,6 +127,8 @@ func (m *EtcdManager) Acquire(ctx context.Context) (Lease, error) {
 		ctx = context.Background()
 	}
 
+	linearizableCtx := clientv3.WithRequireLeader(ctx)
+
 	session, err := concurrency.NewSession(m.client, concurrency.WithTTL(m.ttlSeconds))
 	if err != nil {
 		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
@@ -136,7 +138,7 @@ func (m *EtcdManager) Acquire(ctx context.Context) (Lease, error) {
 	}
 
 	mutex := concurrency.NewMutex(session, m.key)
-	if err := mutex.TryLock(ctx); err != nil {
+	if err := mutex.TryLock(linearizableCtx); err != nil {
 		_ = session.Close()
 		if errors.Is(err, concurrency.ErrLocked) {
 			return nil, ErrNotAcquired
@@ -147,8 +149,9 @@ func (m *EtcdManager) Acquire(ctx context.Context) (Lease, error) {
 		return nil, fmt.Errorf("try lock: %w", err)
 	}
 
-	if err := m.annotateLease(ctx, session, mutex); err != nil {
-		cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	if err := m.annotateLease(linearizableCtx, session, mutex); err != nil {
+		cleanupBase := clientv3.WithRequireLeader(context.Background())
+		cleanupCtx, cancel := context.WithTimeout(cleanupBase, 5*time.Second)
 		_ = mutex.Unlock(cleanupCtx)
 		cancel()
 		_ = session.Close()
@@ -173,6 +176,8 @@ func (l *etcdLease) Release(ctx context.Context) error {
 		ctx = context.Background()
 	}
 
+	ctx = clientv3.WithRequireLeader(ctx)
+
 	unlockErr := l.mutex.Unlock(ctx)
 	closeErr := l.session.Close()
 

pkg/lock/etcd_manager_test.go

@@ -174,7 +174,8 @@ func TestEtcdManagerAnnotatesLease(t *testing.T) {
 	}
 
 	key := internal.mutex.Key()
-	resp, err := etcdClient.Get(context.Background(), key)
+	ctx := clientv3.WithRequireLeader(context.Background())
+	resp, err := etcdClient.Get(ctx, key)
 	if err != nil {
 		t.Fatalf("failed to read metadata: %v", err)
 	}
@@ -206,7 +207,7 @@ func TestEtcdManagerAnnotatesLease(t *testing.T) {
 		t.Fatalf("failed to release lock: %v", err)
 	}
 
-	resp, err = etcdClient.Get(context.Background(), key)
+	resp, err = etcdClient.Get(ctx, key)
 	if err != nil {
 		t.Fatalf("failed to read metadata after release: %v", err)
 	}

pkg/orchestrator/loop.go

@@ -178,20 +178,14 @@ func (l *Loop) Run(ctx context.Context) error {
 				return err
 			}
 			if err := l.executor.Execute(ctx, outcome.Command); err != nil {
-				clearErr := error(nil)
-				if outcome.cooldownStarted {
-					clearErr = outcome.clearCooldown(context.Background())
-				}
-				releaseErr := outcome.ReleaseLock(context.Background())
+				// Keep the cooldown marker active so other nodes continue to back off; the reboot
+				// command may still be in progress even though it reported an error.
 				execErr := fmt.Errorf("execute reboot command: %w", err)
-				finalErr := execErr
-				if clearErr != nil {
-					finalErr = errors.Join(finalErr, clearErr)
-				}
+				releaseErr := outcome.ReleaseLock(context.Background())
 				if releaseErr != nil {
-					finalErr = errors.Join(finalErr, releaseErr)
+					return errors.Join(execErr, releaseErr)
 				}
-				return finalErr
+				return execErr
 			}
 			return nil
 		}

pkg/orchestrator/loop_test.go

@@ -159,7 +159,7 @@ func TestLoopStartsCooldownBeforeExecuting(t *testing.T) {
 	}
 }
 
-func TestLoopClearsCooldownWhenExecutorFails(t *testing.T) {
+func TestLoopKeepsCooldownWhenExecutorFails(t *testing.T) {
 	cfg := baseConfig()
 	startCalls := 0
 	clearCalls := 0
@@ -189,8 +189,8 @@ func TestLoopClearsCooldownWhenExecutorFails(t *testing.T) {
 	if startCalls != 1 {
 		t.Fatalf("expected cooldown start to be invoked once, got %d", startCalls)
 	}
-	if clearCalls != 1 {
-		t.Fatalf("expected cooldown clear to run on failure, got %d", clearCalls)
+	if clearCalls != 0 {
+		t.Fatalf("expected cooldown clear to be skipped on failure, got %d", clearCalls)
 	}
 }