From f70ec13cf02ea6482f7410c9859dd7d7eb1946d7 Mon Sep 17 00:00:00 2001
From: smessie <smessie@smessie.com>
Date: Mon, 26 Aug 2024 09:01:52 +0200
Subject: [PATCH] feat: Allow writer to be ended and error on writing to closed
 writer

---
 src/connectors.ts       |  7 +++++++
 src/connectors/file.ts  |  8 ++++----
 src/connectors/http.ts  |  8 ++++----
 src/connectors/kafka.ts | 13 +++++++------
 src/connectors/ws.ts    | 12 ++++++------
 5 files changed, 28 insertions(+), 20 deletions(-)

diff --git a/src/connectors.ts b/src/connectors.ts
index f47099e..b7ae6f6 100644
--- a/src/connectors.ts
+++ b/src/connectors.ts
@@ -223,6 +223,7 @@ export class ChannelFactory {
 export interface Writer<T> {
     push(item: T): Promise<void>;
     end(): Promise<void>;
+    on(event: "end", listener: Handler<void>): this;
 }
 
 export interface Stream<T> {
@@ -242,6 +243,8 @@ export class SimpleStream<T> implements Stream<T> {
     public readonly disconnect: () => Promise<void>;
     public lastElement?: T | undefined;
 
+    private ended = false;
+
     public constructor(onDisconnect?: () => Promise<void>) {
         this.disconnect = onDisconnect || (async () => {});
     }
@@ -252,6 +255,9 @@ export class SimpleStream<T> implements Stream<T> {
     }
 
     public async push(data: T): Promise<void> {
+        if (this.ended) {
+            throw new Error("Trying to push to a stream that has ended!");
+        }
         this.lastElement = data;
         await Promise.all(this.dataHandlers.map((handler) => handler(data)));
     }
@@ -259,6 +265,7 @@ export class SimpleStream<T> implements Stream<T> {
     public async end(): Promise<void> {
         await this.disconnect();
         await Promise.all(this.endHandlers.map((handler) => handler()));
+        this.ended = true;
     }
 
     public on(event: "data", listener: Handler<T>): this;
diff --git a/src/connectors/file.ts b/src/connectors/file.ts
index 2217f1a..7be4110 100644
--- a/src/connectors/file.ts
+++ b/src/connectors/file.ts
@@ -121,13 +121,15 @@ export const startFileStreamWriter: WriterConstructor<FileWriterConfig> = (
         : `${process.cwd()}/${config.path}`;
     const encoding: BufferEncoding = <BufferEncoding>config.encoding || "utf-8";
 
+    const writer = new SimpleStream<string>();
+
     const init = async () => {
         if (!config.onReplace) {
             await writeFile(path, "", { encoding });
         }
     };
 
-    const push = async (item: string): Promise<void> => {
+    writer.push = async (item: string): Promise<void> => {
         if (config.onReplace) {
             await writeFile(path, item, { encoding });
         } else {
@@ -135,7 +137,5 @@ export const startFileStreamWriter: WriterConstructor<FileWriterConfig> = (
         }
     };
 
-    const end = async (): Promise<void> => {};
-
-    return { writer: { push, end }, init };
+    return { writer, init };
 };
diff --git a/src/connectors/http.ts b/src/connectors/http.ts
index 96b4874..c0a2111 100644
--- a/src/connectors/http.ts
+++ b/src/connectors/http.ts
@@ -101,7 +101,9 @@ export const startHttpStreamWriter: WriterConstructor<HttpWriterConfig> = (
 ) => {
     const requestConfig = <https.RequestOptions>new URL(config.endpoint);
 
-    const push = async (item: string | Buffer): Promise<void> => {
+    const writer = new SimpleStream<string | Buffer>();
+
+    writer.push = async (item: string | Buffer): Promise<void> => {
         await new Promise((resolve) => {
             const options = {
                 hostname: requestConfig.hostname,
@@ -125,7 +127,5 @@ export const startHttpStreamWriter: WriterConstructor<HttpWriterConfig> = (
         });
     };
 
-    const end = async (): Promise<void> => {};
-
-    return { writer: { push, end }, init: async () => {} };
+    return { writer, init: async () => {} };
 };
diff --git a/src/connectors/kafka.ts b/src/connectors/kafka.ts
index 4e2762d..2ac13b4 100644
--- a/src/connectors/kafka.ts
+++ b/src/connectors/kafka.ts
@@ -138,15 +138,16 @@ export const startKafkaStreamWriter: WriterConstructor<KafkaWriterConfig> = (
     const kafka = new Kafka(<KafkaConfig>brokerConfig);
 
     const producer = kafka.producer(config.producer);
+
+    const writer = new SimpleStream<string>(async () => {
+        await producer.disconnect();
+    });
+
     const init = () => producer.connect();
 
-    const push = async (item: string): Promise<void> => {
+    writer.push = async (item: string): Promise<void> => {
         await producer.send({ topic, messages: [{ value: item }] });
     };
 
-    const end = async (): Promise<void> => {
-        await producer.disconnect();
-    };
-
-    return { writer: { push, end }, init };
+    return { writer, init };
 };
diff --git a/src/connectors/ws.ts b/src/connectors/ws.ts
index cd7fbe4..cff627d 100644
--- a/src/connectors/ws.ts
+++ b/src/connectors/ws.ts
@@ -94,14 +94,14 @@ export const startWsStreamWriter: WriterConstructor<WsWriterConfig> = (
         ws.on("open", () => console.log("open"));
     };
 
+    const writer = new SimpleStream<string>(async () => {
+        ws.close();
+    });
+
     /* eslint-disable-next-line @typescript-eslint/no-explicit-any */
-    const push = async (item: any): Promise<void> => {
+    writer.push = async (item: any): Promise<void> => {
         ws.send(item);
     };
 
-    const end = async (): Promise<void> => {
-        ws.close();
-    };
-
-    return { writer: { push, end }, init };
+    return { writer, init };
 };