From b491aa28f22cedc635e2312dab7ac162fbdf9cb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafa=C5=82=20Sumis=C5=82awski?=
 <rafal.sumislawski@gmail.com>
Date: Wed, 4 Oct 2023 14:59:08 +0200
Subject: [PATCH] Await for flushInProgress when flushingAll (#3)

* Await for flushInProgress when flushingAll. This is important to ensure that all exports are complete before freezing lambda instance.

* fix lint error
---
 .../src/export/BatchSpanProcessorBase.ts         | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/packages/opentelemetry-sdk-trace-base/src/export/BatchSpanProcessorBase.ts b/packages/opentelemetry-sdk-trace-base/src/export/BatchSpanProcessorBase.ts
index 5cb29e33e4..e4eb9b367d 100644
--- a/packages/opentelemetry-sdk-trace-base/src/export/BatchSpanProcessorBase.ts
+++ b/packages/opentelemetry-sdk-trace-base/src/export/BatchSpanProcessorBase.ts
@@ -45,7 +45,7 @@ export abstract class BatchSpanProcessorBase<T extends BufferConfig>
   private _timer: NodeJS.Timeout | undefined;
   private _shutdownOnce: BindOnceFuture<void>;
   private _droppedSpansCount: number = 0;
-  private _isFlushInProgress: boolean = false;
+  private _flushInProgress: Promise<void> | null = null;
 
   constructor(private readonly _exporter: SpanExporter, config?: T) {
     const env = getEnv();
@@ -162,6 +162,9 @@ export abstract class BatchSpanProcessorBase<T extends BufferConfig>
         // run exports in parallel ignoring _isFlushInProgress
         promises.push(this._export(spans));
       }
+      if (this._flushInProgress) {
+        promises.push(this._flushInProgress);
+      }
       Promise.all(promises)
         .then(() => {
           resolve();
@@ -230,16 +233,15 @@ export abstract class BatchSpanProcessorBase<T extends BufferConfig>
   }
 
   private _flush() {
-    if (this._isFlushInProgress) {
+    if (this._flushInProgress) {
       return;
     }
-    this._isFlushInProgress = true;
     this._clearTimer();
-
     const spans = this._finishedSpans.splice(0, this._maxExportBatchSize);
-    this._export(spans)
+    this._flushInProgress = this._export(spans);
+    this._flushInProgress
       .then(() => {
-        this._isFlushInProgress = false;
+        this._flushInProgress = null;
         if (this._finishedSpans.length >= this._maxExportBatchSize) {
           this._flush();
         } else if (this._finishedSpans.length > 0) {
@@ -247,7 +249,7 @@ export abstract class BatchSpanProcessorBase<T extends BufferConfig>
         }
       })
       .catch(e => {
-        this._isFlushInProgress = false;
+        this._flushInProgress = null;
         globalErrorHandler(e);
         if (this._finishedSpans.length >= this._maxExportBatchSize) {
           this._flush();