[8.x] Onboard elastic owned ECH clusters to use `mget` task…

… claiming (#196757) (#196856) # Backport This will backport the following commits from `main` to `8.x`: - [Onboard elastic owned ECH clusters to use `mget` task claiming (#196757)](#196757)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Mike Côté <[email protected]>
elastic · Oct 18, 2024 · e32e716 · e32e716
1 parent 10a9161
commit e32e716
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 47 deletions.
diff --git a/x-pack/plugins/task_manager/server/lib/set_claim_strategy.test.ts b/x-pack/plugins/task_manager/server/lib/set_claim_strategy.test.ts
@@ -71,61 +71,67 @@ describe('setClaimStrategy', () => {
   });
   for (const isServerless of [true, false]) {
     for (const isCloud of [true, false]) {
-      for (const deploymentId of [undefined, deploymentIdMget, deploymentIdUpdateByQuery]) {
-        for (const configuredStrategy of [CLAIM_STRATEGY_MGET, CLAIM_STRATEGY_UPDATE_BY_QUERY]) {
-          test(`should return config as is when claim strategy is already defined: isServerless=${isServerless}, isCloud=${isCloud}, deploymentId=${deploymentId}`, () => {
-            const config = {
-              ...getConfigWithoutClaimStrategy(),
-              claim_strategy: configuredStrategy,
-            };
-
-            const returnedConfig = setClaimStrategy({
-              config,
-              logger,
-              isCloud,
-              isServerless,
-              deploymentId,
+      for (const isElasticStaffOwned of [true, false]) {
+        for (const deploymentId of [undefined, deploymentIdMget, deploymentIdUpdateByQuery]) {
+          for (const configuredStrategy of [CLAIM_STRATEGY_MGET, CLAIM_STRATEGY_UPDATE_BY_QUERY]) {
+            test(`should return config as is when claim strategy is already defined: isServerless=${isServerless}, isCloud=${isCloud}, isElasticStaffOwned=${isElasticStaffOwned}, deploymentId=${deploymentId}`, () => {
+              const config = {
+                ...getConfigWithoutClaimStrategy(),
+                claim_strategy: configuredStrategy,
+              };
+
+              const returnedConfig = setClaimStrategy({
+                config,
+                logger,
+                isCloud,
+                isServerless,
+                isElasticStaffOwned,
+                deploymentId,
+              });
+
+              expect(returnedConfig).toStrictEqual(config);
+              if (deploymentId) {
+                expect(logger.info).toHaveBeenCalledWith(
+                  `Using claim strategy ${configuredStrategy} as configured for deployment ${deploymentId}`
+                );
+              } else {
+                expect(logger.info).toHaveBeenCalledWith(
+                  `Using claim strategy ${configuredStrategy} as configured`
+                );
+              }
             });
-
-            expect(returnedConfig).toStrictEqual(config);
-            if (deploymentId) {
-              expect(logger.info).toHaveBeenCalledWith(
-                `Using claim strategy ${configuredStrategy} as configured for deployment ${deploymentId}`
-              );
-            } else {
-              expect(logger.info).toHaveBeenCalledWith(
-                `Using claim strategy ${configuredStrategy} as configured`
-              );
-            }
-          });
+          }
         }
       }
     }
   }
 
   for (const isCloud of [true, false]) {
-    for (const deploymentId of [undefined, deploymentIdMget, deploymentIdUpdateByQuery]) {
-      test(`should set claim strategy to mget if in serverless: isCloud=${isCloud}, deploymentId=${deploymentId}`, () => {
-        const config = getConfigWithoutClaimStrategy();
-        const returnedConfig = setClaimStrategy({
-          config,
-          logger,
-          isCloud,
-          isServerless: true,
-          deploymentId,
-        });
+    for (const isElasticStaffOwned of [true, false]) {
+      for (const deploymentId of [undefined, deploymentIdMget, deploymentIdUpdateByQuery]) {
+        test(`should set claim strategy to mget if in serverless: isCloud=${isCloud}, isElasticStaffOwned=${isElasticStaffOwned}, deploymentId=${deploymentId}`, () => {
+          const config = getConfigWithoutClaimStrategy();
+          const returnedConfig = setClaimStrategy({
+            config,
+            logger,
+            isCloud,
+            isServerless: true,
+            isElasticStaffOwned,
+            deploymentId,
+          });
 
-        expect(returnedConfig.claim_strategy).toBe(CLAIM_STRATEGY_MGET);
-        expect(returnedConfig.poll_interval).toBe(MGET_DEFAULT_POLL_INTERVAL);
+          expect(returnedConfig.claim_strategy).toBe(CLAIM_STRATEGY_MGET);
+          expect(returnedConfig.poll_interval).toBe(MGET_DEFAULT_POLL_INTERVAL);
 
-        if (deploymentId) {
-          expect(logger.info).toHaveBeenCalledWith(
-            `Setting claim strategy to mget for serverless deployment ${deploymentId}`
-          );
-        } else {
-          expect(logger.info).toHaveBeenCalledWith(`Setting claim strategy to mget`);
-        }
-      });
+          if (deploymentId) {
+            expect(logger.info).toHaveBeenCalledWith(
+              `Setting claim strategy to mget for serverless deployment ${deploymentId}`
+            );
+          } else {
+            expect(logger.info).toHaveBeenCalledWith(`Setting claim strategy to mget`);
+          }
+        });
+      }
     }
   }
 
@@ -135,6 +141,7 @@ describe('setClaimStrategy', () => {
       config,
       logger,
       isCloud: false,
+      isElasticStaffOwned: false,
       isServerless: false,
     });
 
@@ -150,6 +157,7 @@ describe('setClaimStrategy', () => {
       config,
       logger,
       isCloud: true,
+      isElasticStaffOwned: false,
       isServerless: false,
     });
 
@@ -165,6 +173,7 @@ describe('setClaimStrategy', () => {
       config,
       logger,
       isCloud: true,
+      isElasticStaffOwned: false,
       isServerless: false,
       deploymentId: deploymentIdUpdateByQuery,
     });
@@ -177,12 +186,32 @@ describe('setClaimStrategy', () => {
     );
   });
 
+  test(`should set claim strategy to mget if cloud, deploymentId does not start with a or b, not serverless and isElasticStaffOwned is true`, () => {
+    const config = getConfigWithoutClaimStrategy();
+    const returnedConfig = setClaimStrategy({
+      config,
+      logger,
+      isCloud: true,
+      isElasticStaffOwned: true,
+      isServerless: false,
+      deploymentId: deploymentIdUpdateByQuery,
+    });
+
+    expect(returnedConfig.claim_strategy).toBe(CLAIM_STRATEGY_MGET);
+    expect(returnedConfig.poll_interval).toBe(MGET_DEFAULT_POLL_INTERVAL);
+
+    expect(logger.info).toHaveBeenCalledWith(
+      `Setting claim strategy to mget for deployment ${deploymentIdUpdateByQuery}`
+    );
+  });
+
   test(`should set claim strategy to mget if cloud and not serverless and deploymentId starts with a or b`, () => {
     const config = getConfigWithoutClaimStrategy();
     const returnedConfig = setClaimStrategy({
       config,
       logger,
       isCloud: true,
+      isElasticStaffOwned: false,
       isServerless: false,
       deploymentId: deploymentIdMget,
     });

diff --git a/x-pack/plugins/task_manager/server/lib/set_claim_strategy.ts b/x-pack/plugins/task_manager/server/lib/set_claim_strategy.ts
@@ -19,6 +19,7 @@ interface SetClaimStrategyOpts {
   deploymentId?: string;
   isServerless: boolean;
   isCloud: boolean;
+  isElasticStaffOwned: boolean;
   logger: Logger;
 }
 
@@ -50,7 +51,10 @@ export function setClaimStrategy(opts: SetClaimStrategyOpts): TaskManagerConfig
   let defaultToMget = false;
 
   if (opts.isCloud && !opts.isServerless && opts.deploymentId) {
-    defaultToMget = opts.deploymentId.startsWith('a') || opts.deploymentId.startsWith('b');
+    defaultToMget =
+      opts.deploymentId.startsWith('a') ||
+      opts.deploymentId.startsWith('b') ||
+      opts.isElasticStaffOwned;
     if (defaultToMget) {
       opts.logger.info(`Setting claim strategy to mget for deployment ${opts.deploymentId}`);
     } else {

diff --git a/x-pack/plugins/task_manager/server/plugin.ts b/x-pack/plugins/task_manager/server/plugin.ts
@@ -136,6 +136,7 @@ export class TaskManagerPlugin
       deploymentId: plugins.cloud?.deploymentId,
       isServerless: this.initContext.env.packageInfo.buildFlavor === 'serverless',
       isCloud: plugins.cloud?.isCloudEnabled ?? false,
+      isElasticStaffOwned: plugins.cloud?.isElasticStaffOwned ?? false,
       logger: this.logger,
     });