suppor multiple pipelines properly and use new response (#71)

Co-authored-by: Jacob Heric <jacob@jacobheric.com:q!>

Author: jacobheric

README.md

@@ -31,7 +31,7 @@ const audience: Audience = {
 
 export const example = async () => {
   const streamdal = new Streamdal(config);
-  const result = await streamdal.processPipeline({
+  const result = await streamdal.process({
     audience,
     data: new TextEncoder().encode(JSON.stringify(exampleData)),
   });

examples/commonjs/src/index.ts

@@ -1,3 +1,5 @@
+import { SDKResponse } from "@streamdal/node-sdk";
+
 const {
   OperationType,
   Streamdal,
@@ -48,16 +50,16 @@ const audience = {
 
 export const example = async () => {
   const streamdal = new Streamdal(config);
-  const result = await streamdal.processPipeline({
+  const result: SDKResponse = await streamdal.process({
     audience,
     data: new TextEncoder().encode(JSON.stringify(exampleData)),
   });
 
   if (result.error) {
-    console.error("Pipeline error", result.message);
+    console.error("Pipeline error", result.errorMessage);
     //
     // Optionally explore more detailed step status information
-    console.dir(result.stepStatuses);
+    console.dir(result.pipelineStatus);
   } else {
     console.info("Pipeline success!");
     //

examples/esm/src/index.ts

@@ -1,9 +1,9 @@
 import {
   Audience,
   OperationType,
+  SDKResponse,
   Streamdal,
   StreamdalConfigs,
-  StreamdalResponse,
 } from "@streamdal/node-sdk";
 
 const exampleData = {
@@ -50,16 +50,16 @@ const audience: Audience = {
 
 export const example = async () => {
   const streamdal = new Streamdal(config);
-  const result: StreamdalResponse = await streamdal.processPipeline({
+  const result: SDKResponse = await streamdal.process({
     audience,
     data: new TextEncoder().encode(JSON.stringify(exampleData)),
   });
 
   if (result.error) {
-    console.error("Pipeline error", result.message);
+    console.error("Pipeline error", result.errorMessage);
     //
     // Optionally explore more detailed step status information
-    console.dir(result.stepStatuses);
+    console.dir(result.pipelineStatus);
   } else {
     console.info("Pipeline success!");
     //

package-lock.json

@@ -43,7 +43,7 @@
     "@grpc/proto-loader": "0.7.7",
     "@protobuf-ts/grpc-transport": "2.9.0",
     "@protobuf-ts/runtime-rpc": "2.9.0",
-    "@streamdal/protos": "^0.0.123",
+    "@streamdal/protos": "^0.0.125",
     "@types/rwlock": "^5.0.3",
     "@types/sinon": "^10.0.16",
     "dotenv": "^16.3.1",

package.json

@@ -97,54 +97,55 @@ const testAttachCommandByServiceResponse = {
 };
 
 describe("pipeline tests", () => {
+  const key = audienceKey(testAudience);
   it("initPipelines should add a given pipeline to internal store and set pipelineInitialized ", async () => {
     sinon
       .stub(testConfigs.grpcClient, "getAttachCommandsByService")
       .resolves(testAttachCommandByServiceResponse as any);
 
     await initPipelines(testConfigs);
 
-    expect(internal.pipelines.has(audienceKey(testAudience))).toEqual(true);
+    expect(internal.pipelines.has(key)).toEqual(true);
     expect(internal.pipelineInitialized).toEqual(true);
   });
 
   describe("processResponse", () => {
-    it("attach command should add pipeline to internal store", () => {
-      processResponse(testAttachCommand);
+    it("attach command should add pipeline to internal store", async () => {
+      await processResponse(testAttachCommand);
+      expect(internal.pipelines.has(key)).toEqual(true);
 
-      expect(internal.pipelines.has(audienceKey(testAudience))).toEqual(true);
+      const pipelines = internal.pipelines.get(key);
+      expect(pipelines?.has(testPipeline.id)).toEqual(true);
     });
 
-    it("detach command should remove pipeline from internal store", () => {
-      internal.pipelines.set(audienceKey(testAudience), {
-        ...testPipeline,
-        paused: false,
-      });
-      processResponse(testDetachCommand);
+    it("detach command should remove pipeline from internal store", async () => {
+      internal.pipelines.set(key, new Map([[testPipeline.id, testPipeline]]));
+      await processResponse(testDetachCommand);
 
-      expect(internal.pipelines.has(audienceKey(testAudience))).toEqual(false);
+      expect(internal.pipelines.get(key)?.has(testPipeline.id)).toEqual(false);
     });
 
-    it("pause command should flag pipeline as paused in internal store", () => {
-      internal.pipelines.set(audienceKey(testAudience), {
-        ...testPipeline,
-        paused: false,
-      });
-      processResponse(testPauseCommand);
+    it("pause command should flag pipeline as paused in internal store", async () => {
+      internal.pipelines.set(
+        key,
+        new Map([[testPipeline.id, { ...testPipeline, paused: false }]])
+      );
+
+      await processResponse(testPauseCommand);
 
-      expect(internal.pipelines.get(audienceKey(testAudience))?.paused).toEqual(
+      expect(internal.pipelines.get(key)?.get(testPipeline.id)?.paused).toEqual(
         true
       );
     });
 
-    it("resume command should flag pipeline as not paused in internal store", () => {
-      internal.pipelines.set(audienceKey(testAudience), {
-        ...testPipeline,
-        paused: false,
-      });
-      processResponse(testResumeCommand);
+    it("resume command should flag pipeline as not paused in internal store", async () => {
+      internal.pipelines.set(
+        key,
+        new Map([[testPipeline.id, { ...testPipeline, paused: true }]])
+      );
+      await processResponse(testResumeCommand);
 
-      expect(internal.pipelines.get(audienceKey(testAudience))?.paused).toEqual(
+      expect(internal.pipelines.get(key)?.get(testPipeline.id)?.paused).toEqual(
         false
       );
     });

src/__tests__/pipeline.test.ts

@@ -4,9 +4,10 @@ import {
   OperationType,
 } from "@streamdal/protos/protos/sp_common";
 import { IInternalClient } from "@streamdal/protos/protos/sp_internal.client";
+import { StepStatus } from "@streamdal/protos/protos/sp_sdk";
 import ReadWriteLock from "rwlock";
 
-import { StepStatus } from "./process.js";
+import { InternalPipeline } from "./pipeline.js";
 
 export const METRIC_INTERVAL = 1000;
 
@@ -18,25 +19,34 @@ export interface MetricsConfigs {
   streamdalToken: string;
 }
 
-export const getStepLabels = (audience: Audience, stepStatus: StepStatus) => ({
+export const getStepLabels = (
+  audience: Audience,
+  pipeline: InternalPipeline
+) => ({
   service: audience.serviceName,
   component: audience.componentName,
   operation: audience.operationName,
-  pipeline_id: stepStatus.pipelineId,
-  pipeline_name: stepStatus.pipelineName,
+  pipeline_id: pipeline.id,
+  pipeline_name: pipeline.name,
 });
 
-export const stepMetrics = async (
-  audience: Audience,
-  stepStatus: StepStatus,
-  payloadSize: number
+export const stepMetrics = async ({
+  audience,
+  stepStatus,
+  pipeline,
+  payloadSize,
+}: {
+  audience: Audience;
+  stepStatus: StepStatus;
+  pipeline: InternalPipeline;
+  payloadSize: number;
   // eslint-disable-next-line @typescript-eslint/require-await
-) => {
+}) => {
   lock.writeLock((release) => {
     const opName =
       audience.operationType === OperationType.CONSUMER ? "consume" : "produce";
 
-    const labels = getStepLabels(audience, stepStatus);
+    const labels = getStepLabels(audience, pipeline);
     const stepErrorKey = `counter_${opName}_errors`;
     const stepProcessedKey = `counter_${opName}_processed`;
     const stepBytesKey = `counter_${opName}_bytes`;

src/internal/metrics.ts

@@ -11,13 +11,12 @@ export type InternalPipeline = Pipeline & {
   paused?: boolean;
 };
 
-export type EnhancedStep = PipelineStep & {
-  pipelineId: string;
-  pipelineName: string;
-};
-
 export const initPipelines = async (configs: Configs) => {
   try {
+    if (internal.pipelineInitialized) {
+      return;
+    }
+
     console.debug("initializing pipelines");
     const { response }: { response: GetAttachCommandsByServiceResponse } =
       await configs.grpcClient.getAttachCommandsByService(
@@ -27,20 +26,20 @@ export const initPipelines = async (configs: Configs) => {
         { meta: { "auth-token": configs.streamdalToken } }
       );
 
-    for (const [k, v] of Object.entries(response.wasmModules)) {
-      void instantiateWasm(k, v.bytes);
+    for await (const [k, v] of Object.entries(response.wasmModules)) {
+      await instantiateWasm(k, v.bytes);
     }
 
-    for (const command of response.active) {
-      processResponse(command);
+    for await (const command of response.active) {
+      await processResponse(command);
     }
     internal.pipelineInitialized = true;
   } catch (e) {
     console.error("Error initializing pipelines", e);
   }
 };
 
-export const processResponse = (response: Command) => {
+export const processResponse = async (response: Command) => {
   if (!response.audience) {
     response.command.oneofKind !== "keepAlive" &&
       console.debug("command response has no audience, ignoring");
@@ -50,10 +49,10 @@ export const processResponse = (response: Command) => {
   switch (response.command.oneofKind) {
     case "attachPipeline":
       response.command.attachPipeline.pipeline &&
-        attachPipeline(
+        (await attachPipeline(
           response.audience,
           response.command.attachPipeline.pipeline
-        );
+        ));
       break;
     case "detachPipeline":
       detachPipeline(
@@ -81,40 +80,50 @@ export const processResponse = (response: Command) => {
   }
 };
 
-export const buildPipeline = (pipeline: Pipeline): Pipeline => {
+export const buildPipeline = async (pipeline: Pipeline): Promise<Pipeline> => {
   return {
     ...pipeline,
-    steps: pipeline.steps.map((step: PipelineStep) => {
-      void instantiateWasm(step.WasmId, step.WasmBytes);
-      return {
-        ...step,
-        WasmBytes: undefined,
-      };
-    }),
+    steps: await Promise.all(
+      pipeline.steps.map(async (step: PipelineStep) => {
+        await instantiateWasm(step.WasmId, step.WasmBytes);
+        return {
+          ...step,
+          WasmBytes: undefined,
+        };
+      })
+    ),
   };
 };
 
-export const attachPipeline = (audience: Audience, pipeline: Pipeline) => {
-  internal.pipelines.set(audienceKey(audience), {
-    ...buildPipeline(pipeline),
-    paused: false,
-  });
-};
-
-export const detachPipeline = (audience: Audience, pipelineId: string) => {
+export const attachPipeline = async (
+  audience: Audience,
+  pipeline: Pipeline
+) => {
   const key = audienceKey(audience);
-  const p = internal.pipelines.get(key);
-  pipelineId === p?.id && internal.pipelines.delete(key);
+  const existing = internal.pipelines.get(key);
+  const built = await buildPipeline(pipeline);
+  internal.pipelines.set(
+    key,
+    existing
+      ? existing.set(pipeline.id, built)
+      : new Map([[pipeline.id, built]])
+  );
 };
 
+export const detachPipeline = (audience: Audience, pipelineId: string) =>
+  internal.pipelines.get(audienceKey(audience))?.delete(pipelineId);
+
 export const togglePausePipeline = (
   audience: Audience,
   pipelineId: string,
   paused: boolean
 ) => {
   const key = audienceKey(audience);
-  const p = internal.pipelines.get(key);
-  pipelineId === p?.id && internal.pipelines.set(key, { ...p, paused });
+  const existing = internal.pipelines.get(key);
+  const p = existing?.get(pipelineId);
+  existing &&
+    p &&
+    internal.pipelines.set(key, existing.set(p.id, { ...p, paused }));
 };
 
 export const tailPipeline = (audience: Audience, { request }: TailCommand) => {