feat(session-manager): Manage a pool of sessions for job execution

2026-01-08 04:50:06 +00:00 · 2020-07-17 08:31:27 +01:00
parent a579f481c5
commit a12244cf78
3 changed files with 252 additions and 140 deletions
--- a/src/SessionManager.ts
+++ b/src/SessionManager.ts
@@ -0,0 +1,138 @@
+import { Session, Context, CsrfToken } from "./types";
+import { asyncForEach, makeRequest } from "./utils";
+
+const MAX_SESSION_COUNT = 1;
+
+export class SessionManager {
+  constructor(private serverUrl: string, private contextName: string) {}
+  private sessions: Session[] = [];
+  private currentContext: Context | null = null;
+  private csrfToken: CsrfToken | null = null;
+
+  async getSession(accessToken?: string) {
+    await this.createSessions(accessToken);
+    this.createAndWaitForSession(accessToken);
+    return this.sessions.pop();
+  }
+
+  private async createSessions(accessToken?: string) {
+    if (!this.sessions.length) {
+      if (!this.currentContext) {
+        await this.setCurrentContext(accessToken);
+      }
+      await asyncForEach(new Array(MAX_SESSION_COUNT), async () => {
+        const createdSession = await this.createAndWaitForSession(accessToken);
+        this.sessions.push(createdSession);
+      });
+    }
+  }
+
+  private async createAndWaitForSession(accessToken?: string) {
+    const createSessionRequest = {
+      method: "POST",
+      headers: this.getHeaders(accessToken),
+    };
+    const { result: createdSession, etag } = await this.request<Session>(
+      `${this.serverUrl}/compute/contexts/${this.currentContext!.id}/sessions`,
+      createSessionRequest
+    );
+
+    await this.waitForSession(createdSession, etag);
+    return createdSession;
+  }
+
+  private async setCurrentContext(accessToken?: string) {
+    if (!this.currentContext) {
+      const { result: contexts } = await this.request<{
+        items: Context[];
+      }>(`${this.serverUrl}/compute/contexts`, {
+        headers: this.getHeaders(accessToken),
+      });
+
+      const contextsList =
+        contexts && contexts.items && contexts.items.length
+          ? contexts.items
+          : [];
+
+      const currentContext = contextsList.find(
+        (c: any) => c.name === this.contextName
+      );
+
+      if (!currentContext) {
+        throw new Error(
+          `The context ${this.contextName} was not found on the server ${this.serverUrl}`
+        );
+      }
+
+      this.currentContext = currentContext;
+    }
+  }
+
+  private getHeaders(accessToken?: string) {
+    const headers: any = {
+      "Content-Type": "application/json",
+    };
+    if (accessToken) {
+      headers.Authorization = `Bearer ${accessToken}`;
+    }
+
+    return headers;
+  }
+
+  private async waitForSession(
+    session: Session,
+    etag: string | null,
+    accessToken?: string,
+    silent = false
+  ) {
+    let sessionState = session.state;
+    const headers: any = {
+      ...this.getHeaders(accessToken),
+      "If-None-Match": etag,
+    };
+    const stateLink = session.links.find((l: any) => l.rel === "state");
+    return new Promise(async (resolve, _) => {
+      if (sessionState === "pending") {
+        if (stateLink) {
+          if (!silent) {
+            console.log("Polling session status... \n");
+          }
+          const { result: state } = await this.request<string>(
+            `${this.serverUrl}${stateLink.href}?wait=30`,
+            {
+              headers,
+            },
+            "text"
+          );
+
+          sessionState = state.trim();
+          if (!silent) {
+            console.log(`Current state: ${sessionState}\n`);
+          }
+          resolve(sessionState);
+        }
+      } else {
+        resolve(sessionState);
+      }
+    });
+  }
+
+  private async request<T>(
+    url: string,
+    options: RequestInit,
+    contentType: "text" | "json" = "json"
+  ) {
+    if (this.csrfToken) {
+      options.headers = {
+        ...options.headers,
+        [this.csrfToken.headerName]: this.csrfToken.value,
+      };
+    }
+    return await makeRequest<T>(
+      url,
+      options,
+      (token) => (this.csrfToken = token),
+      contentType
+    );
+  }
+}