@@ -78,6 +78,9 @@ class Batcher {
7878 this . contentType = 'application/json'
7979 }
8080
81+ this . batchesSending = 0
82+ this . onBatchesFlushed = ( ) => { }
83+
8184 // If batching is enabled, run the loop
8285 this . options . batching && this . run ( )
8386
@@ -88,6 +91,46 @@ class Batcher {
8891 }
8992 }
9093
94+ /**
95+ * Marks the start of batch submitting.
96+ *
97+ * Must be called right before batcher starts sending logs.
98+ */
99+ batchSending ( ) {
100+ this . batchesSending ++
101+ }
102+
103+ /**
104+ * Marks the end of batch submitting
105+ *
106+ * Must be called after the response from Grafana Loki push endpoint
107+ * is received and completely processed, right before
108+ * resolving/rejecting the promise.
109+ */
110+ batchSent ( ) {
111+ if ( -- this . batchesSending ) return
112+
113+ this . onBatchesFlushed ( )
114+ }
115+
116+ /**
117+ * Returns a promise that resolves after all the logs sent before
118+ * via log(), info(), etc calls are sent to Grafana Loki push endpoint
119+ * and the responses for all of them are received and processed.
120+ *
121+ * @returns {Promise }
122+ */
123+ waitFlushed ( ) {
124+ return new Promise ( ( resolve , reject ) => {
125+ if ( ! this . batchesSending && ! this . batch . streams . length ) { return resolve ( ) }
126+
127+ this . onBatchesFlushed = ( ) => {
128+ this . onBatchesFlushed = ( ) => { }
129+ return resolve ( )
130+ }
131+ } )
132+ }
133+
91134 /**
92135 * Returns a promise that resolves after the given duration.
93136 *
@@ -157,9 +200,11 @@ class Batcher {
157200 * @returns {Promise }
158201 */
159202 sendBatchToLoki ( logEntry ) {
203+ this . batchSending ( )
160204 return new Promise ( ( resolve , reject ) => {
161205 // If the batch is empty, do nothing
162206 if ( this . batch . streams . length === 0 && ! logEntry ) {
207+ this . batchSent ( )
163208 resolve ( )
164209 } else {
165210 let reqBody
@@ -200,6 +245,7 @@ class Batcher {
200245 // Compress the buffer with snappy
201246 reqBody = snappy . compressSync ( buffer )
202247 } catch ( err ) {
248+ this . batchSent ( )
203249 reject ( err )
204250 }
205251 }
@@ -209,6 +255,7 @@ class Batcher {
209255 . then ( ( ) => {
210256 // No need to clear the batch if batching is disabled
211257 logEntry === undefined && this . clearBatch ( )
258+ this . batchSent ( )
212259 resolve ( )
213260 } )
214261 . catch ( err => {
@@ -217,6 +264,7 @@ class Batcher {
217264
218265 this . options . onConnectionError !== undefined && this . options . onConnectionError ( err )
219266
267+ this . batchSent ( )
220268 reject ( err )
221269 } )
222270 }
0 commit comments