add a wrapper with a setClient method

src/server/task_manager/client_wrapper.ts

@@ -0,0 +1,40 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { TaskManager } from './task_pool';
+import { TaskInstance } from './task_pool/task';
+
+export class TaskManagerClientWrapper {
+  private client: TaskManager | null;
+
+  constructor() {
+    this.client = null;
+  }
+
+  public setClient(client: TaskManager) {
+    this.client = client;
+  }
+
+  public schedule(task: TaskInstance) {
+    if (this.client == null) {
+      throw new Error('Task Manager Client has not been set properly!');
+    }
+    this.client.schedule(task);
+  }
+}

src/server/task_manager/task_manager_mixin.ts

@@ -18,6 +18,7 @@
  */
 
 import Joi from 'joi';
+import { TaskManagerClientWrapper } from './client_wrapper';
 import {
   Logger,
   TaskDefinition,
@@ -29,16 +30,15 @@ import {
 } from './task_pool';
 import { SanitizedTaskDefinition } from './task_pool/task';
 
-function taskManagerFactory(kbnServer: any, server: any, config: any) {
-  let taskManager: TaskManager;
-  return async function getTaskManager() {
-    if (taskManager) {
-      return taskManager;
-    }
+export async function taskManagerMixin(kbnServer: any, server: any, config: any) {
+  const logger = new Logger((...args) => server.log(...args));
+  const numWorkers = config.get('taskManager.num_workers');
+  const definitions = extractTaskDefinitions(numWorkers, kbnServer.uiExports.taskDefinitions);
+
+  server.decorate('server', 'taskManager', new TaskManagerClientWrapper());
 
-    const logger = new Logger((...args) => server.log(...args));
+  kbnServer.afterPluginsInit(async () => {
     const callCluster = server.plugins.elasticsearch.getCluster('admin').callWithInternalUser;
-    const numWorkers = config.get('taskManager.num_workers');
     const store = new TaskStore({
       index: config.get('taskManager.index'),
       callCluster,
@@ -48,8 +48,6 @@ function taskManagerFactory(kbnServer: any, server: any, config: any) {
     logger.debug('Initializing the task manager index');
     await store.init();
 
-    const definitions = extractTaskDefinitions(numWorkers, kbnServer.uiExports.taskDefinitions);
-
     const pool = new TaskPool({
       logger,
       callCluster,
@@ -62,18 +60,8 @@ function taskManagerFactory(kbnServer: any, server: any, config: any) {
 
     pool.start();
 
-    taskManager = new TaskManager({
-      store,
-      pool,
-    });
-
-    return taskManager;
-  };
-}
-
-export async function taskManagerMixin(kbnServer: any, server: any, config: any) {
-  server.decorate('server', 'taskManager', {
-    getTaskManager: taskManagerFactory(kbnServer, server, config),
+    const client = new TaskManager({ store, pool });
+    server.taskManager.setClient(client);
   });
 }