From 74e876d12d5a9927d63d41702778cd42e7deec80 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.
---
 .../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 0db293a0d0b4..350711cad426 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 ef6a9c26cb29..ad504d546e84 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' } });
+          });
 
-          await sleep(1000); // wait for lock to expire
+          it('can be re-acquired', async () => {
+            expect(reacquireResult).to.be(true);
+          });
 
-          const reacquired = await manager2.acquire({ metadata: { attempt: 'two' } });
-          expect(reacquired).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);
+          });
+
+          it('updates the metadata when re-acquired', async () => {
+            const reacquiredLock = await getLockById(es, LOCK_ID);
+            expect(reacquiredLock?.metadata).to.eql({ attempt: 'two' });
+          });
+
+          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);
+          });
         });
       });