From 6fee0516025df18f2026e33b320370fa199f4846 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=B8ren=20Louv-Jansen?= <soren.louv@elastic.co>
Date: Thu, 8 May 2025 19:40:20 +0200
Subject: [PATCH] [LockManager] Update token and metadata when an expired lock
 is re-acquired (#220476)

Related: https://github.com/elastic/kibana/pull/216397

This fixes a bug in the Lock Manager where an expired lock can be
acquired, but the token and metadata is not updated. This means that the
lock cannot be released. Instead it is automatically released when the
TTL expires.

(cherry picked from commit 74e876d12d5a9927d63d41702778cd42e7deec80)
---
 .../src/lock_manager_client.ts                | 10 ++---
 .../distributed_lock_manager.spec.ts          | 40 +++++++++++++++----
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/packages/kbn-lock-manager/src/lock_manager_client.ts b/packages/kbn-lock-manager/src/lock_manager_client.ts
index 0db293a0d0b4d..350711cad4268 100644
--- a/packages/kbn-lock-manager/src/lock_manager_client.ts
+++ b/packages/kbn-lock-manager/src/lock_manager_client.ts
@@ -87,6 +87,8 @@ export class LockManager {
                 def instantNow = Instant.ofEpochMilli(now);
                 ctx._source.createdAt = instantNow.toString();
                 ctx._source.expiresAt = instantNow.plusMillis(params.ttl).toString();
+                ctx._source.metadata = params.metadata;
+                ctx._source.token = params.token;
               } else {
                 ctx.op = 'noop';
               }
@@ -94,13 +96,11 @@ export class LockManager {
             params: {
               ttl,
               token: this.token,
+              metadata,
             },
           },
           // @ts-expect-error
-          upsert: {
-            metadata,
-            token: this.token,
-          },
+          upsert: {},
         },
         {
           retryOnTimeout: true,
@@ -189,7 +189,7 @@ export class LockManager {
         this.logger.debug(`Lock "${this.lockId}" released with token ${this.token}.`);
         return true;
       case 'noop':
-        this.logger.debug(
+        this.logger.warn(
           `Lock "${this.lockId}" with token = ${this.token} could not be released. Token does not match.`
         );
         return false;
diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts
index ef6a9c26cb29b..ad504d546e843 100644
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/distributed_lock_manager/distributed_lock_manager.spec.ts
@@ -236,15 +236,41 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
           expect(lock?.metadata).to.eql({ attempt: 'one' });
         });
 
-        it('allows re-acquisition after expiration', async () => {
-          // Acquire with a very short TTL.
-          const acquired = await manager1.acquire({ ttl: 500, metadata: { attempt: 'one' } });
-          expect(acquired).to.be(true);
+        describe('when a lock by "manager1" expires, and is attempted re-acquired by "manager2"', () => {
+          let expiredLock: LockDocument | undefined;
+          let reacquireResult: boolean;
+          beforeEach(async () => {
+            // Acquire with a very short TTL.
+            const acquired = await manager1.acquire({ ttl: 500, metadata: { attempt: 'one' } });
+            expect(acquired).to.be(true);
+            await sleep(1000); // wait for lock to expire
+            expiredLock = await getLockById(es, LOCK_ID);
+            reacquireResult = await manager2.acquire({ metadata: { attempt: 'two' } });
+          });
+
+          it('can be re-acquired', async () => {
+            expect(reacquireResult).to.be(true);
+          });
+
+          it('updates the token when re-acquired', async () => {
+            const reacquiredLock = await getLockById(es, LOCK_ID);
+            expect(expiredLock?.token).not.to.be(reacquiredLock?.token);
+          });
 
-          await sleep(1000); // wait for lock to expire
+          it('updates the metadata when re-acquired', async () => {
+            const reacquiredLock = await getLockById(es, LOCK_ID);
+            expect(reacquiredLock?.metadata).to.eql({ attempt: 'two' });
+          });
 
-          const reacquired = await manager2.acquire({ metadata: { attempt: 'two' } });
-          expect(reacquired).to.be(true);
+          it('cannot be released by "manager1"', async () => {
+            const res = await manager1.release();
+            expect(res).to.be(false);
+          });
+
+          it('can be released by "manager2"', async () => {
+            const res = await manager2.release();
+            expect(res).to.be(true);
+          });
         });
       });