@rdfc/js-runner 2.0.0 → 3.0.1

package/src/reader.ts

@@ -1,6 +1,11 @@
 import { ClientReadableStream } from '@grpc/grpc-js'
-import { DataChunk, Message, RunnerClient, StreamMessage } from '@rdfc/proto'
-import winston from 'winston'
+import {
+  DataChunk,
+  ReceivingMessage,
+  ReceivingStreamMessage,
+  RunnerClient,
+} from '@rdfc/proto'
+import { Logger } from 'winston'
 import {
   AnyConvertor,
   Convertor,
@@ -8,17 +13,19 @@ import {
   StreamConvertor,
   StringConvertor,
 } from './convertor'
+import { Writable } from './runner'
+import { promisify } from 'util'
 
 export type Any =
   | {
-      string: string
-    }
+    string: string
+  }
   | {
-      stream: AsyncGenerator<Uint8Array>
-    }
+    stream: AsyncGenerator<Uint8Array>
+  }
   | {
-      buffer: Uint8Array
-    }
+    buffer: Uint8Array
+  }
 
 export interface Reader {
   readonly uri: string
@@ -28,41 +35,47 @@ export interface Reader {
   anys(): AsyncIterable<Any>
 }
 
+type Todo<T> = {
+  item: T
+  onComplete: () => void
+}
+
 class MyIter<T> implements AsyncIterable<T> {
   private convertor: Convertor<T>
-  private queue: (T | undefined)[] = []
+  private queue: Todo<T | undefined>[] = []
   private resolveNext: ((value: undefined) => void) | null = null
 
   constructor(convertor: Convertor<T>) {
     this.convertor = convertor
   }
 
-  push(buffer: Uint8Array) {
+  push(buffer: Uint8Array, onComplete: () => void) {
     const item = this.convertor.from(buffer)
-    this.queue.push(item)
+    this.queue.push({ item, onComplete })
     if (this.resolveNext) {
       this.resolveNext(undefined)
       this.resolveNext = null
     }
   }
 
-  close() {
-    this.queue.push(undefined)
+  close(onComplete: () => void) {
+    this.queue.push({ item: undefined, onComplete })
     if (this.resolveNext) {
       this.resolveNext(undefined)
       this.resolveNext = null
     }
   }
 
-  async pushStream(chunks: ClientReadableStream<DataChunk>) {
-    const stream = (async function* (stream) {
-      for await (const c of stream) {
-        const chunk: DataChunk = c
+  async pushStream(chunks: AsyncIterable<DataChunk>, onComplete: () => void) {
+    // This is an async generator that transforms DataChunks to Buffers
+    const stream = (async function*(stream) {
+      for await (const chunk of stream) {
         yield chunk.data
       }
     })(chunks)
+
     const item = await this.convertor.fromStream(stream)
-    this.queue.push(item)
+    this.queue.push({ item, onComplete })
     if (this.resolveNext) {
       this.resolveNext(undefined)
       this.resolveNext = null
@@ -72,9 +85,15 @@ class MyIter<T> implements AsyncIterable<T> {
   async *[Symbol.asyncIterator]() {
     while (true) {
       if (this.queue.length > 0) {
-        const item = this.queue.shift()!
-        if (item === undefined) break
+        const { item, onComplete } = this.queue.shift()!
+        if (item === undefined) {
+          onComplete()
+          break
+        }
         yield item
+        // Note: execution pauses at `yield` until the consumer calls `.next()` again.
+        // We call onComplete *after* resuming, so the producer knows the item was actually consumed.
+        onComplete()
       } else {
         await new Promise<undefined>((resolve) => (this.resolveNext = resolve))
       }
@@ -85,58 +104,217 @@ class MyIter<T> implements AsyncIterable<T> {
 export class ReaderInstance implements Reader {
   private client: RunnerClient
   readonly uri: string
-  private logger: winston.Logger
+  private logger: Logger
+  private readonly notifyOrchestrator: Writable
 
-  private iterators: MyIter<unknown>[] = []
+  private consumers: MyIter<unknown>[] = []
 
-  constructor(uri: string, client: RunnerClient, logger: winston.Logger) {
+  constructor(
+    uri: string,
+    client: RunnerClient,
+    notifyOrchestrator: Writable,
+    logger: Logger,
+  ) {
     this.uri = uri
     this.client = client
     this.logger = logger
+    this.notifyOrchestrator = notifyOrchestrator
   }
 
   anys(): AsyncIterable<Any> {
     const iter = new MyIter(AnyConvertor)
-    this.iterators.push(iter)
+    this.consumers.push(iter)
     return iter
   }
 
   strings(): AsyncIterable<string> {
     const iter = new MyIter(StringConvertor)
-    this.iterators.push(iter)
+    this.consumers.push(iter)
     return iter
   }
 
   buffers(): AsyncIterable<Uint8Array> {
     const iter = new MyIter(NoConvertor)
-    this.iterators.push(iter)
+    this.consumers.push(iter)
     return iter
   }
 
   streams(): AsyncIterable<AsyncGenerator<Uint8Array>> {
     const iter = new MyIter(StreamConvertor)
-    this.iterators.push(iter)
+    this.consumers.push(iter)
     return iter
   }
 
-  handleMsg(msg: Message) {
+  handleMsg(msg: ReceivingMessage) {
     this.logger.debug(`${this.uri} handling message`)
-    for (const iter of this.iterators) {
-      iter.push(msg.data)
+
+    const promises = []
+    for (const iter of this.consumers) {
+      promises.push(new Promise((res) => iter.push(msg.data, () => res(null))))
     }
+
+    Promise.all(promises).then(() =>
+      this.notifyOrchestrator({
+        processed: {
+          globalSequenceNumber: msg.globalSequenceNumber,
+          channel: this.uri,
+        },
+      }),
+    )
   }
 
   close() {
-    for (const iter of this.iterators) {
-      iter.close()
+    for (const iter of this.consumers) {
+      iter.close(() => { })
     }
   }
 
-  handleStreamingMessage(msg: StreamMessage) {
+  // There is a stream message available for this reader
+  async handleStreamingMessage({
+    channel,
+    globalSequenceNumber,
+  }: ReceivingStreamMessage) {
     this.logger.debug(`${this.uri} handling streaming message`)
-    const chunks = this.client.receiveStreamMessage(msg.id!)
-    for (const iter of this.iterators) {
-      iter.pushStream(chunks)
+
+    const chunks = this.client.receiveStreamMessage()
+    const writeControlMessage = promisify(chunks.write.bind(chunks))
+    const consumersConsumed = []
+
+    // After each chunk is handled by all consumer, emit a processed message
+    let idx = 0
+    const messageIterators = fanoutStream(
+      chunks,
+      this.consumers.length,
+      async () => {
+        await writeControlMessage({ streamSequenceNumber: idx++ })
+      },
+    )
+
+    for (const consumer of this.consumers) {
+      consumersConsumed.push(
+        new Promise((res) =>
+          consumer.pushStream(messageIterators.pop()!, () => res(null)),
+        ),
+      )
+    }
+
+    await writeControlMessage({ globalSequenceNumber })
+
+    Promise.all(consumersConsumed).then(() => {
+      console.log('Writing processed for streaming message')
+      this.notifyOrchestrator({ processed: { globalSequenceNumber, channel } })
+    })
+  }
+}
+
+/**
+ * Helper function to tee a stream `numConsumers` times
+ * When each tee'd stream has handled a chunk, call {@link onAllHandled}
+ */
+function fanoutStream<T>(
+  stream: ClientReadableStream<T>,
+  numConsumers: number,
+  onAllHandled: () => void | Promise<void>,
+): AsyncIterable<T>[] {
+  type Waiter = (value: IteratorResult<T>) => void
+
+  let ended = false
+  const buffer: T[] = []
+  const pending: Waiter[] = []
+  let activeConsumers = numConsumers
+
+  // consumer bookkeeping
+  let awaitingAck = 0
+
+  function pushChunk(chunk: T) {
+    buffer.push(chunk)
+    flush()
+  }
+
+  function flush() {
+    while (buffer.length > 0 && pending.length > 0) {
+      const chunk = buffer[0] // keep until all consumers ack
+      const waiter = pending.shift()!
+      waiter({ value: chunk, done: false })
+      awaitingAck++
+    }
+  }
+
+  function end() {
+    ended = true
+    while (pending.length > 0) {
+      const waiter = pending.shift()!
+      waiter({ value: undefined, done: true })
+    }
+  }
+
+  stream.on('data', (chunk: T) => {
+    pushChunk(chunk)
+  })
+
+  stream.on('end', () => {
+    end()
+  })
+
+  stream.on('error', (err) => {
+    while (pending.length > 0) {
+      const waiter = pending.shift()!
+      waiter({ value: undefined, done: true })
+    }
+    throw err
+  })
+
+  function makeIterable(): AsyncIterable<T> {
+    return {
+      [Symbol.asyncIterator]() {
+        return {
+          next(): Promise<IteratorResult<T>> {
+            if (buffer.length > 0) {
+              const chunk = buffer[0]
+              awaitingAck++
+              return Promise.resolve({ value: chunk, done: false })
+            }
+            if (ended) {
+              return Promise.resolve({ value: undefined, done: true })
+            }
+            return new Promise((resolve) => {
+              pending.push(resolve)
+            })
+          },
+          async return() {
+            activeConsumers--
+            if (activeConsumers === 0) {
+              end()
+            }
+            return { value: undefined, done: true }
+          },
+        }
+      },
     }
   }
+
+  async function ack() {
+    awaitingAck--
+    if (awaitingAck === 0) {
+      // all consumers done with the current chunk
+      buffer.shift() // drop it
+      await onAllHandled()
+      flush() // continue with next chunk
+    }
+  }
+
+  // wrap consumer so they *must* call ack() after processing
+  function wrap(iterable: AsyncIterable<T>): AsyncIterable<T> {
+    return {
+      async *[Symbol.asyncIterator]() {
+        for await (const item of iterable) {
+          yield item
+          await ack()
+        }
+      },
+    }
+  }
+
+  const rawIterables = Array.from({ length: numConsumers }, makeIterable)
+  return rawIterables.map(wrap)
 }

package/src/runner.ts

@@ -1,15 +1,18 @@
 import {
-  OrchestratorMessage,
+  Close,
+  FromRunner,
   Processor,
   RunnerClient,
-  RunnerMessage,
+  ToRunner,
+  LocalAck,
+  ReceivingMessage,
+  ReceivingStreamMessage,
 } from '@rdfc/proto'
 import { Reader, ReaderInstance } from './reader'
 import { Writer, WriterInstance } from './writer'
 import { Processor as Proc } from './processor'
-import { Logger } from 'winston'
+import { createLogger, Logger } from 'winston'
 
-import winston from 'winston'
 import { RpcTransport } from './logger'
 import { Cont, empty, extractShapes, Shapes } from 'rdf-lens'
 import { NamedNode, Parser } from 'n3'
@@ -36,7 +39,7 @@ const RDFC = createNamespace(
   'Writer',
 )
 
-export type Writable = (msg: OrchestratorMessage) => Promise<unknown>
+export type Writable = (msg: FromRunner) => Promise<unknown>
 
 type ProcessorConfig = {
   location: string
@@ -44,13 +47,12 @@ type ProcessorConfig = {
   clazz?: string
 }
 
-export type FullProc<C extends Proc<unknown>> =
-  C extends Proc<infer T> ? T & C : unknown
+export type FullProc<C> = C extends Proc<infer T> ? T & C : never
 export class Runner {
-  private readonly readers: { [uri: string]: ReaderInstance[] } = {}
-  private readonly writers: { [uri: string]: WriterInstance[] } = {}
+  private readonly readers: { [uri: string]: ReaderInstance } = {}
+  private readonly writers: { [uri: string]: WriterInstance } = {}
   private readonly client: RunnerClient
-  private readonly write: Writable
+  private readonly notifyOrchestrator: Writable
   private readonly logger: Logger
   private shapes: Shapes
   private quads: Quad[] = []
@@ -62,12 +64,12 @@ export class Runner {
 
   constructor(
     client: RunnerClient,
-    write: Writable,
+    notifyOrchestrator: Writable,
     uri: string,
     logger: Logger,
   ) {
     this.client = client
-    this.write = write
+    this.notifyOrchestrator = notifyOrchestrator
     this.uri = uri
     this.logger = logger
   }
@@ -75,7 +77,7 @@ export class Runner {
   async addProcessor<P extends Proc<unknown>>(
     proc: Processor,
   ): Promise<FullProc<P>> {
-    const procLogger = winston.createLogger({
+    const procLogger = createLogger({
       transports: [
         new RpcTransport({
           entities: [proc.uri, this.uri],
@@ -84,14 +86,13 @@ export class Runner {
       ],
     })
 
-    const ty = JSON.stringify(
-      this.quads
-        .filter(
-          (x) =>
-            x.subject.value === proc.uri && x.predicate.equals(RDF.terms.type),
-        )
-        .map((x) => x.object.value),
-    )
+    const ty = this.quads
+      .filter(
+        (x) =>
+          x.subject.value === proc.uri && x.predicate.equals(RDF.terms.type),
+      )
+      .map((x) => x.object.value)
+
     this.logger.info('parsing ' + proc.uri + ' type ' + ty)
     const args = this.shapes.lenses[RDFL.TypedExtract].execute({
       id: new NamedNode(proc.uri),
@@ -99,8 +100,6 @@ export class Runner {
     })
 
     const config: ProcessorConfig = JSON.parse(proc.config)
-    // const url = new URL(config.location)
-    // process.chdir(url.pathname)
     const jsProgram = await import(config.file)
     const clazz = jsProgram[config.clazz || 'default']
     const instance: Proc<unknown> = new clazz(args, procLogger)
@@ -111,7 +110,7 @@ export class Runner {
     this.processors.push(instance)
     this.processorTransforms.push(instance.transform())
 
-    await this.write({ init: { uri: proc.uri } })
+    await this.notifyOrchestrator({ initialized: { uri: proc.uri } })
 
     return <FullProc<P>>instance
   }
@@ -126,72 +125,135 @@ export class Runner {
   }
 
   createWriter(uri: Term): Writer {
-    const ids = uri.value
+    const id = uri.value
 
-    if (this.writers[ids] === undefined) {
-      this.writers[ids] = []
+    if (this.writers[id] !== undefined) {
+      return this.writers[id]
     }
-    const writer = new WriterInstance(ids, this.client, this.write, this.logger)
-    this.writers[ids].push(writer)
+    const writer = new WriterInstance(
+      id,
+      this.client,
+      this.notifyOrchestrator,
+      this.uri,
+      this.logger,
+    )
+    this.writers[id] = writer
     return writer
   }
 
   createReader(uri: Term): Reader {
     const ids = uri.value
 
-    if (this.readers[ids] === undefined) {
-      this.readers[ids] = []
+    if (this.readers[ids] !== undefined) {
+      return this.readers[ids]
     }
-    const reader = new ReaderInstance(ids, this.client, this.logger)
-    this.readers[ids].push(reader)
+    const reader = new ReaderInstance(
+      ids,
+      this.client,
+      this.notifyOrchestrator,
+      this.logger,
+    )
+    this.readers[ids] = reader
     return reader
   }
 
-  async handleOrchMessage(msg: RunnerMessage) {
+  async handleOrchMessage(msg: ToRunner) {
     if (msg.msg) {
-      this.logger.debug('Handling data msg for ' + msg.msg.channel)
-      for (const reader of this.readers[msg.msg.channel] || []) {
-        reader.handleMsg(msg.msg)
-      }
+      this.handleMsg(msg.msg)
     }
 
     if (msg.streamMsg) {
-      for (const reader of this.readers[msg.streamMsg.channel] || []) {
-        reader.handleStreamingMessage(msg.streamMsg)
-      }
+      await this.handleStreamMsg(msg.streamMsg)
     }
 
     if (msg.close) {
-      const uri = msg.close.channel
+      await this.handleClose(msg.close)
+    }
 
-      for (const reader of this.readers[uri] || []) {
-        reader.close()
-      }
+    if (msg.pipeline) {
+      this.handlePipeline(msg.pipeline)
+    }
 
-      for (const writer of this.writers[uri] || []) {
-        await writer.close(true)
-      }
+    if (msg.processed) {
+      this.handleProcessed(msg.processed)
     }
+  }
 
-    if (msg.pipeline) {
-      try {
-        // here
-        const quads = new Parser().parse(msg.pipeline)
-        this.shapes = extractShapes(
-          quads,
-          {
-            [RDFC.Reader]: (x: Cont) => this.createReader(x.id),
-            [RDFC.Writer]: (x: Cont) => this.createWriter(x.id),
-          },
-          {
-            [RDFC.Reader]: empty<Cont>(),
-            [RDFC.Writer]: empty<Cont>(),
-          },
-        )
-        this.quads = quads
-      } catch (ex: unknown) {
-        this.logger.error('Pipeline failed: ' + JSON.stringify(ex))
-      }
+  private async handleClose(close: Close) {
+    const uri = close.channel
+    const r = this.readers[uri]
+
+    let closed = false
+    if (r) {
+      r.close()
+      closed = true
+    }
+    const w = this.writers[uri]
+    if (w) {
+      closed = true
+      await w.close(true)
+    }
+
+    if (!closed) {
+      this.logger.error(
+        `Received a close event for channel ${uri}, but neither reader nor writer is present.`,
+      )
+    }
+  }
+
+  private handlePipeline(pipeline: string) {
+    try {
+      const quads = new Parser().parse(pipeline)
+      this.shapes = extractShapes(
+        quads,
+        {
+          [RDFC.Reader]: (x: Cont) => this.createReader(x.id),
+          [RDFC.Writer]: (x: Cont) => this.createWriter(x.id),
+        },
+        {
+          [RDFC.Reader]: empty<Cont>(),
+          [RDFC.Writer]: empty<Cont>(),
+        },
+      )
+      this.quads = quads
+    } catch (ex: unknown) {
+      this.logger.error('Pipeline failed: ' + JSON.stringify(ex))
+    }
+  }
+
+  private handleMsg(msg: ReceivingMessage) {
+    this.logger.debug('Handling data msg for ' + msg.channel)
+    const r = this.readers[msg.channel]
+
+    if (r) {
+      r.handleMsg(msg)
+    } else {
+      this.logger.error(
+        `Received message for channel ${msg.channel}, but no reader was present.`,
+      )
+    }
+  }
+
+  private async handleStreamMsg(streamMsg: ReceivingStreamMessage) {
+    const r = this.readers[streamMsg.channel]
+
+    if (r) {
+      await r.handleStreamingMessage(streamMsg)
+    } else {
+      this.logger.error(
+        `Received stream message for channel ${streamMsg.channel}, but no reader was present.`,
+      )
+    }
+  }
+
+  private handleProcessed(processed: LocalAck) {
+    const writer = this.writers[processed.channel]
+    if (writer) {
+      writer.handled()
+    } else {
+      this.logger.error(
+        `Received processed message for channel ${processed.channel}, but no writer was present.`,
+      )
     }
   }
 }