fix: premount loader race condition

Author: ArushKapoorJuspay

src/Payments/PreMountLoader.res

@@ -12,30 +12,36 @@ let sendPromiseData = (promise, key) => {
 
 let useMessageHandler = getPromisesAndHandler => {
   React.useEffect(_ => {
-    let (promises, messageHandler) = getPromisesAndHandler()
+    let messageHandler = getPromisesAndHandler()
+
     let setupMessageListener = _ => {
-      Utils.messageParentWindow([("preMountLoaderIframeMountedCallback", true->JSON.Encode.bool)])
       Window.addEventListener("message", messageHandler)
+      Utils.messageParentWindow([("preMountLoaderIframeMountedCallback", true->JSON.Encode.bool)])
     }
 
     let cleanupMessageListener = _ => {
-      Utils.messageParentWindow([("preMountLoaderIframeUnMount", true->JSON.Encode.bool)])
       Window.removeEventListener("message", messageHandler)
+      Utils.messageParentWindow([("preMountLoaderIframeUnMount", true->JSON.Encode.bool)])
     }
 
-    setupMessageListener()
-
-    let executeAllPromises = async () => {
-      try {
-        let _ = await Promise.all(promises)
-      } catch {
-      | error => Console.error2("Error in message handler:", error)
+    let handleCleanUpEventListener = (ev: Window.event) => {
+      open Utils
+      let dict = ev.data->safeParse->getDictFromJson
+      if dict->Dict.get("cleanUpPreMountLoaderIframe")->Option.isSome {
+        cleanupMessageListener()
       }
-      cleanupMessageListener()
     }
-    executeAllPromises()->ignore
 
-    Some(cleanupMessageListener)
+    Window.addEventListener("message", handleCleanUpEventListener)
+
+    setupMessageListener()
+
+    Some(
+      () => {
+        cleanupMessageListener()
+        Window.removeEventListener("message", handleCleanUpEventListener)
+      },
+    )
   }, [])
 }
 
@@ -87,8 +93,7 @@ module PreMountLoaderForElements = {
         }
       }
 
-      let promises = [paymentMethodsPromise, customerPaymentMethodsPromise, sessionTokensPromise]
-      (promises, messageHandler)
+      messageHandler
     })
 
     React.null
@@ -114,8 +119,7 @@ module PreMountLoaderForPMMElements = {
         }
       }
 
-      let promises = [savedPaymentMethodsPromise]
-      (promises, messageHandler)
+      messageHandler
     })
 
     React.null

src/hyper-loader/Elements.res

@@ -142,12 +142,13 @@ let make = (
       }
     }
 
-    let fetchPaymentsList = (mountedIframeRef, componentType) => {
+    let fetchPaymentsList = (mountedIframeRef, componentType, promiseResolve) => {
       let handlePaymentMethodsLoaded = (event: Types.event) => {
         let json = event.data->anyTypeToJson
         let dict = json->getDictFromJson
         let isPaymentMethodsData = dict->getString("data", "") === "payment_methods"
         if isPaymentMethodsData {
+          promiseResolve()
           isTaxCalculationEnabled.contents =
             dict->getDictFromDict("response")->getBool("is_tax_calculation_enabled", false)
           addSmartEventListener("message", onPlaidCallback(mountedIframeRef), "onPlaidCallback")
@@ -212,6 +213,7 @@ let make = (
       mountedIframeRef,
       disableSavedPaymentMethods,
       componentType,
+      promiseResolve,
     ) => {
       if !disableSavedPaymentMethods {
         let handleCustomerPaymentMethodsLoaded = (event: Types.event) => {
@@ -223,13 +225,16 @@ let make = (
             let json = dict->getJsonFromDict("response", JSON.Encode.null)
             let msg = [("customerPaymentMethods", json)]->Dict.fromArray
             mountedIframeRef->Window.iframePostMessage(msg)
+            promiseResolve()
           }
         }
         addSmartEventListener(
           "message",
           handleCustomerPaymentMethodsLoaded,
           `onCustomerPaymentMethodsLoaded-${componentType}`,
         )
+      } else {
+        promiseResolve()
       }
       let msg =
         [
@@ -809,13 +814,14 @@ let make = (
         addSmartEventListener("message", handleGooglePayThirdPartyFlow, "onGooglePayThirdParty")
         addSmartEventListener("message", handleApplePayThirdPartyFlow, "onApplePayThirdParty")
 
-        let fetchSessionTokens = mountedIframeRef => {
+        let fetchSessionTokens = (mountedIframeRef, promiseResolve) => {
           let handleSessionTokensLoaded = (event: Types.event) => {
             let json = event.data->anyTypeToJson
             let dict = json->getDictFromJson
             let sessionTokensData = dict->getString("data", "") === "session_tokens"
             if sessionTokensData {
               let json = dict->getJsonFromDict("response", JSON.Encode.null)
+              promiseResolve()
 
               {
                 let sessionsArr =
@@ -1288,22 +1294,42 @@ let make = (
         }
         preMountLoaderMountedPromise
         ->then(_ => {
-          fetchPaymentsList(mountedIframeRef, componentType)
+          let paymentMethodsPromise = Promise.make((resolve, _) => {
+            fetchPaymentsList(mountedIframeRef, componentType, resolve)
+          })
           let disableSavedPaymentMethods =
             newOptions
             ->getDictFromJson
             ->getBool("displaySavedPaymentMethods", true) &&
               !(spmComponents->Array.includes(componentType))->not
-          fetchCustomerPaymentMethods(mountedIframeRef, disableSavedPaymentMethods, componentType)
-          fetchSessionTokens(mountedIframeRef)
-          resolve()
+          let customerPaymentMethodsPromise = Promise.make((resolve, _) => {
+            fetchCustomerPaymentMethods(
+              mountedIframeRef,
+              disableSavedPaymentMethods,
+              componentType,
+              resolve,
+            )
+          })
+          let sessionTokensPromise = Promise.make((resolve, _) => {
+            fetchSessionTokens(mountedIframeRef, resolve)
+          })
+          Promise.all([
+            paymentMethodsPromise,
+            customerPaymentMethodsPromise,
+            sessionTokensPromise,
+          ])->then(_ => {
+            let msg = [("cleanUpPreMountLoaderIframe", true->JSON.Encode.bool)]->Dict.fromArray
+            preMountLoaderIframeDiv->Window.iframePostMessage(msg)
+            resolve()
+          })
         })
         ->catch(_ => resolve())
         ->ignore
 
         mountedIframeRef->Window.iframePostMessage(message)
       }
 
+      Console.log("Creating payment element")
       let paymentElement = LoaderPaymentElement.make(
         componentType,
         newOptions,
@@ -1315,6 +1341,7 @@ let make = (
       savedPaymentElement->Dict.set(componentType, paymentElement)
       paymentElement
     }
+    Console.log("Elements initialization completed")
     {
       getElement,
       update,
@@ -1323,6 +1350,7 @@ let make = (
     }
   } catch {
   | e => {
+      Console.log2(`Elements initialization failed: `, e)
       Sentry.captureException(e)
       defaultElement
     }

src/hyper-loader/PaymentMethodsManagementElements.res

@@ -98,13 +98,19 @@ let make = (
       )
     })
 
-    let fetchSavedPaymentMethods = (mountedIframeRef, disableSaveCards, componentType) => {
+    let fetchSavedPaymentMethods = (
+      mountedIframeRef,
+      disableSaveCards,
+      componentType,
+      resolvePromise,
+    ) => {
       if !disableSaveCards {
         let handleSavedPaymentMethodsLoaded = (event: Types.event) => {
           let json = event.data->Identity.anyTypeToJson
           let dict = json->getDictFromJson
           let isSavedPaymentMethodsData = dict->getString("data", "") === "saved_payment_methods"
           if isSavedPaymentMethodsData {
+            resolvePromise()
             let json = dict->getJsonFromDict("response", JSON.Encode.null)
             let msg = [("savedPaymentMethods", json)]->Dict.fromArray
             mountedIframeRef->Window.iframePostMessage(msg)
@@ -115,6 +121,8 @@ let make = (
           handleSavedPaymentMethodsLoaded,
           `onSavedPaymentMethodsLoaded-${componentType}`,
         )
+      } else {
+        resolvePromise()
       }
       let msg =
         [("sendSavedPaymentMethodsResponse", !disableSaveCards->JSON.Encode.bool)]->Dict.fromArray
@@ -208,7 +216,15 @@ let make = (
             disableSavedPaymentMethods &&
             !(expressCheckoutComponents->Array.includes(componentType))
           ) {
-            fetchSavedPaymentMethods(mountedIframeRef, false, componentType)
+            Promise.make((resolve, _) => {
+              fetchSavedPaymentMethods(mountedIframeRef, false, componentType, resolve)
+            })
+            ->then(_ => {
+              let msg = [("cleanUpPreMountLoaderIframe", true->JSON.Encode.bool)]->Dict.fromArray
+              preMountLoaderIframeDiv->Window.iframePostMessage(msg)
+              resolve()
+            })
+            ->ignore
           }
           resolve()
         })