Add batch flushing

index.d.ts

@@ -18,6 +18,7 @@ declare interface LokiTransportOptions extends TransportStream.TransportStreamOp
 declare class LokiTransport extends TransportStream {
 
     constructor(opts: LokiTransportOptions);
+    flush(): Promise<null>;
 }
 
 export = LokiTransport;

index.js

@@ -104,6 +104,21 @@ class LokiTransport extends Transport {
     callback()
   }
 
+  /**
+   * Flush unsent batched logs to Winston transport and return
+   * a promise that resolves after response is received from
+   * the transport. If some (batched or not) logs are being sent
+   * at the time of call, the promise resolves after the transport
+   * responds.
+   *
+   * As a result the promise returned resolves only when the transport
+   * has confirmed receiving all the logs sent via log(), info(), etc
+   * calls preceding the flush() call.
+   */
+  async flush () {
+    return await this.batcher.waitFlushed();
+  }
+
   /**
    * Send batch to loki when clean up
    */

package.json

@@ -1,6 +1,6 @@
 {
   "name": "winston-loki",
-  "version": "6.0.7-rc1",
+  "version": "6.0.7-rc2",
   "description": "A Winston transport for Grafana Loki",
   "keywords": [
     "winston",

src/batcher.js

@@ -78,6 +78,9 @@ class Batcher {
       this.contentType = 'application/json'
     }
 
+    this.batchesSending = 0
+    this.onBatchesFlushed = () => {}
+
     // If batching is enabled, run the loop
     this.options.batching && this.run()
 
@@ -88,6 +91,46 @@ class Batcher {
     }
   }
 
+  /**
+   * Marks the start of batch submitting.
+   *
+   * Must be called right before batcher starts sending logs.
+   */
+  batchSending () {
+    this.batchesSending++
+  }
+
+  /**
+   * Marks the end of batch submitting
+   *
+   * Must be called after the response from Grafana Loki push endpoint
+   * is received and completely processed, right before
+   * resolving/rejecting the promise.
+   */
+  batchSent () {
+    if (--this.batchesSending) return
+
+    this.onBatchesFlushed()
+  }
+
+  /**
+   * Returns a promise that resolves after all the logs sent before
+   * via log(), info(), etc calls are sent to Grafana Loki push endpoint
+   * and the responses for all of them are received and processed.
+   *
+   * @returns {Promise}
+   */
+  waitFlushed () {
+    return new Promise((resolve, reject) => {
+      if (!this.batchesSending && !this.batch.streams.length) { return resolve() }
+
+      this.onBatchesFlushed = () => {
+        this.onBatchesFlushed = () => {}
+        return resolve()
+      }
+    })
+  }
+
   /**
    * Returns a promise that resolves after the given duration.
    *
@@ -157,9 +200,11 @@ class Batcher {
    * @returns {Promise}
    */
   sendBatchToLoki (logEntry) {
+    this.batchSending()
     return new Promise((resolve, reject) => {
       // If the batch is empty, do nothing
       if (this.batch.streams.length === 0 && !logEntry) {
+        this.batchSent()
         resolve()
       } else {
         let reqBody
@@ -200,6 +245,7 @@ class Batcher {
             // Compress the buffer with snappy
             reqBody = snappy.compressSync(buffer)
           } catch (err) {
+            this.batchSent()
             reject(err)
           }
         }
@@ -209,6 +255,7 @@ class Batcher {
           .then(() => {
             // No need to clear the batch if batching is disabled
             logEntry === undefined && this.clearBatch()
+            this.batchSent()
             resolve()
           })
           .catch(err => {
@@ -217,6 +264,7 @@ class Batcher {
 
             this.options.onConnectionError !== undefined && this.options.onConnectionError(err)
 
+            this.batchSent()
             reject(err)
           })
       }