From 501b7e1b25b47bca55ca7b9ffea4683be21850df Mon Sep 17 00:00:00 2001
From: Gert Hengeveld <info@ghengeveld.nl>
Date: Wed, 10 Aug 2022 09:46:02 +0200
Subject: [PATCH] Always return chained promise so state will be restored
 before execution continues

---
 code/lib/instrumenter/src/instrumenter.ts | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/code/lib/instrumenter/src/instrumenter.ts b/code/lib/instrumenter/src/instrumenter.ts
index 1768e4486798..6488be15096a 100644
--- a/code/lib/instrumenter/src/instrumenter.ts
+++ b/code/lib/instrumenter/src/instrumenter.ts
@@ -502,9 +502,8 @@ export class Instrumenter {
           const res = arg(...args);
 
           // Reset cursor and ancestors to their original values before we entered the callback.
-          if (res instanceof Promise) res.then(restore, restore);
-          else restore();
-
+          if (res instanceof Promise) return res.then(restore, restore);
+          restore();
           return res;
         };
       });