Prototype: Redesign CheckoutInstances registry

paymentsheet/src/main/java/com/stripe/android/checkout/Checkout.kt

@@ -409,10 +409,6 @@ class Checkout private constructor(
      */
     val isLoading: StateFlow<Boolean> = _isLoading.asStateFlow()
 
-    init {
-        CheckoutInstances.add(internalState.key, this)
-    }
-
     /**
      * Applies a promotion code to the checkout session.
      *

paymentsheet/src/main/java/com/stripe/android/checkout/CheckoutInstances.kt

@@ -1,57 +1,40 @@
 package com.stripe.android.checkout
 
+import androidx.annotation.MainThread
 import androidx.annotation.VisibleForTesting
-import com.stripe.android.lpmfoundations.paymentmethod.IntegrationMetadata
-import com.stripe.android.lpmfoundations.paymentmethod.PaymentMethodMetadata
 import com.stripe.android.paymentelement.CheckoutSessionPreview
-import java.lang.ref.WeakReference
 
 @OptIn(CheckoutSessionPreview::class)
 internal object CheckoutInstances {
-    private val instanceMap = mutableMapOf<String, MutableList<WeakReference<Checkout>>>()
-
-    operator fun get(key: String): List<Checkout> {
-        val refs = instanceMap[key] ?: return emptyList()
-        val live = refs.mapNotNull { it.get() }
-        if (live.isEmpty()) {
-            instanceMap.remove(key)
-        } else if (live.size != refs.size) {
-            // Prune stale references
-            refs.clear()
-            refs.addAll(live.map { WeakReference(it) })
-        }
-        return live
-    }
 
-    fun ensureNoMutationInFlight(key: String) {
-        this[key].forEach { it.ensureNoMutationInFlight() }
-    }
+    private data class Entry(val checkout: Checkout, val owner: Any)
 
-    fun markIntegrationLaunched(key: String) {
-        this[key].forEach { it.markIntegrationLaunched() }
-    }
+    private val instances = mutableMapOf<String, Entry>()
 
-    fun markIntegrationDismissed(key: String) {
-        this[key].forEach { it.markIntegrationDismissed() }
-    }
+    @MainThread
+    operator fun get(key: String): Checkout? = instances[key]?.checkout
 
-    fun markIntegrationDismissed(paymentMethodMetadata: PaymentMethodMetadata?) {
-        val checkoutSession = paymentMethodMetadata
-            ?.integrationMetadata as? IntegrationMetadata.CheckoutSession ?: return
-        markIntegrationDismissed(checkoutSession.instancesKey)
-    }
-
-    fun add(key: String, checkout: Checkout) {
-        val refs = instanceMap.getOrPut(key) { mutableListOf() }
-        refs.add(WeakReference(checkout))
+    @MainThread
+    fun register(key: String, checkout: Checkout, owner: Any) {
+        val existing = instances[key]
+        if (existing != null) {
+            check(existing.checkout === checkout && existing.owner === owner) {
+                "Checkout already registered by a different owner. " +
+                    "Only one integration can use a Checkout at a time."
+            }
+        }
+        instances[key] = Entry(checkout, owner)
     }
 
-    fun remove(key: String) {
-        instanceMap.remove(key)
+    @MainThread
+    fun unregister(key: String, checkout: Checkout) {
+        if (instances[key]?.checkout === checkout) {
+            instances.remove(key)
+        }
     }
 
     @VisibleForTesting
     fun clear() {
-        instanceMap.clear()
+        instances.clear()
     }
 }

paymentsheet/src/main/java/com/stripe/android/paymentelement/EmbeddedPaymentElement.kt

@@ -18,7 +18,6 @@ import com.stripe.android.ExperimentalAllowsRemovalOfLastSavedPaymentMethodApi
 import com.stripe.android.SharedPaymentTokenSessionPreview
 import com.stripe.android.checkout.Checkout
 import com.stripe.android.checkout.CheckoutConfigurationMerger
-import com.stripe.android.checkout.CheckoutInstances
 import com.stripe.android.common.configuration.ConfigurationDefaults
 import com.stripe.android.common.ui.DelegateDrawable
 import com.stripe.android.core.utils.StatusBarCompat
@@ -27,6 +26,7 @@ import com.stripe.android.model.PaymentIntent
 import com.stripe.android.model.PaymentMethod
 import com.stripe.android.model.SetupIntent
 import com.stripe.android.paymentelement.embedded.EmbeddedSelectionHolder
+import com.stripe.android.paymentelement.embedded.content.EmbeddedCheckoutRegistrar
 import com.stripe.android.paymentelement.embedded.content.EmbeddedConfigurationCoordinator
 import com.stripe.android.paymentelement.embedded.content.EmbeddedConfirmationHelper
 import com.stripe.android.paymentelement.embedded.content.EmbeddedConfirmationStateHolder
@@ -59,6 +59,7 @@ class EmbeddedPaymentElement @Inject internal constructor(
     paymentOptionDisplayDataHolder: PaymentOptionDisplayDataHolder,
     private val configurationCoordinator: EmbeddedConfigurationCoordinator,
     stateHelper: EmbeddedStateHelper,
+    private val checkoutRegistrar: EmbeddedCheckoutRegistrar,
 ) {
 
     /**
@@ -102,7 +103,8 @@ class EmbeddedPaymentElement @Inject internal constructor(
         checkout: Checkout,
         configuration: Configuration,
     ): ConfigureResult {
-        CheckoutInstances.ensureNoMutationInFlight(checkout.internalState.key)
+        checkout.ensureNoMutationInFlight()
+        checkoutRegistrar.register(checkout)
         return configurationCoordinator.configure(
             configuration = CheckoutConfigurationMerger.EmbeddedConfiguration(configuration)
                 .forCheckoutSession(checkout.internalState),

paymentsheet/src/main/java/com/stripe/android/paymentelement/confirmation/gpay/GooglePayDisplayItemsFactory.kt

@@ -14,8 +14,9 @@ internal object GooglePayDisplayItemsFactory {
         val checkoutSessionMetadata = paymentMethodMetadata.integrationMetadata
             as? IntegrationMetadata.CheckoutSession ?: return emptyList()
 
-        val checkout = CheckoutInstances[checkoutSessionMetadata.instancesKey].firstOrNull()
-            ?: return emptyList()
+        val checkout = requireNotNull(CheckoutInstances[checkoutSessionMetadata.instancesKey]) {
+            "Checkout not registered for key '${checkoutSessionMetadata.instancesKey}'."
+        }
 
         val checkoutSession = checkout.checkoutSession.value
         val items = mutableListOf<GooglePayJsonFactory.DisplayItem>()

paymentsheet/src/main/java/com/stripe/android/paymentelement/confirmation/intent/CheckoutSessionConfirmationInterceptor.kt

@@ -117,7 +117,10 @@ internal class CheckoutSessionConfirmationInterceptor @AssistedInject constructo
         params: ConfirmCheckoutSessionParams,
     ): ConfirmationDefinition.Action<Args> {
         try {
-            CheckoutInstances.ensureNoMutationInFlight(integrationMetadata.instancesKey)
+            val checkout = requireNotNull(CheckoutInstances[integrationMetadata.instancesKey]) {
+                "Checkout not registered for key '${integrationMetadata.instancesKey}'."
+            }
+            checkout.ensureNoMutationInFlight()
         } catch (e: IllegalStateException) {
             return ConfirmationDefinition.Action.Fail(
                 cause = e,
@@ -143,9 +146,10 @@ internal class CheckoutSessionConfirmationInterceptor @AssistedInject constructo
     private fun handleConfirmResponse(
         response: CheckoutSessionResponse,
     ): ConfirmationDefinition.Action<Args> {
-        CheckoutInstances[integrationMetadata.instancesKey].forEach { checkout ->
-            checkout.updateWithResponse(response)
+        val checkout = requireNotNull(CheckoutInstances[integrationMetadata.instancesKey]) {
+            "Checkout not registered for key '${integrationMetadata.instancesKey}'."
         }
+        checkout.updateWithResponse(response)
 
         val intent: StripeIntent = response.paymentIntent ?: response.setupIntent
             ?: run {

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/content/DefaultEmbeddedSheetLauncher.kt

@@ -1,3 +1,5 @@
+@file:OptIn(com.stripe.android.paymentelement.CheckoutSessionPreview::class)
+
 package com.stripe.android.paymentelement.embedded.content
 
 import androidx.activity.result.ActivityResultCaller
@@ -119,8 +121,11 @@ internal class DefaultEmbeddedSheetLauncher @Inject constructor(
     ) {
         val checkoutSession = paymentMethodMetadata.integrationMetadata as? IntegrationMetadata.CheckoutSession
         if (checkoutSession != null) {
-            CheckoutInstances.ensureNoMutationInFlight(checkoutSession.instancesKey)
-            CheckoutInstances.markIntegrationLaunched(checkoutSession.instancesKey)
+            val checkout = requireNotNull(CheckoutInstances[checkoutSession.instancesKey]) {
+                "Checkout not registered for key '${checkoutSession.instancesKey}'."
+            }
+            checkout.ensureNoMutationInFlight()
+            checkout.markIntegrationLaunched()
         }
         if (embeddedConfirmationState == null) {
             errorReporter.report(
@@ -156,8 +161,11 @@ internal class DefaultEmbeddedSheetLauncher @Inject constructor(
     ) {
         val checkoutSession = paymentMethodMetadata.integrationMetadata as? IntegrationMetadata.CheckoutSession
         if (checkoutSession != null) {
-            CheckoutInstances.ensureNoMutationInFlight(checkoutSession.instancesKey)
-            CheckoutInstances.markIntegrationLaunched(checkoutSession.instancesKey)
+            val checkout = requireNotNull(CheckoutInstances[checkoutSession.instancesKey]) {
+                "Checkout not registered for key '${checkoutSession.instancesKey}'."
+            }
+            checkout.ensureNoMutationInFlight()
+            checkout.markIntegrationLaunched()
         }
         if (embeddedConfirmationState == null) {
             errorReporter.report(

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/content/EmbeddedCheckoutRegistrar.kt

@@ -0,0 +1,37 @@
+package com.stripe.android.paymentelement.embedded.content
+
+import com.stripe.android.checkout.Checkout
+import com.stripe.android.checkout.CheckoutInstances
+import com.stripe.android.paymentelement.CheckoutSessionPreview
+import javax.inject.Inject
+import javax.inject.Singleton
+
+/**
+ * Holds and manages the [Checkout] registration for the lifetime of one ViewModel-scoped
+ * [EmbeddedPaymentElement].
+ *
+ * Registration happens in [EmbeddedPaymentElement.configure], and unregistration
+ * happens when the owning ViewModel is cleared via [unregisterIfNeeded].
+ */
+@OptIn(CheckoutSessionPreview::class)
+@Singleton
+internal class EmbeddedCheckoutRegistrar @Inject constructor() {
+
+    @Volatile
+    private var registeredCheckout: Checkout? = null
+
+    fun register(checkout: Checkout) {
+        val previous = registeredCheckout
+        if (previous != null && previous !== checkout) {
+            CheckoutInstances.unregister(previous.internalState.key, previous)
+        }
+        CheckoutInstances.register(checkout.internalState.key, checkout, owner = this)
+        registeredCheckout = checkout
+    }
+
+    fun unregisterIfNeeded() {
+        val checkout = registeredCheckout ?: return
+        CheckoutInstances.unregister(checkout.internalState.key, checkout)
+        registeredCheckout = null
+    }
+}

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/content/EmbeddedPaymentElementViewModel.kt

@@ -17,13 +17,15 @@ import kotlin.reflect.KClass
 internal class EmbeddedPaymentElementViewModel @Inject constructor(
     val embeddedPaymentElementSubcomponentFactory: EmbeddedPaymentElementSubcomponent.Factory,
     @ViewModelScope private val customViewModelScope: CoroutineScope,
+    private val checkoutRegistrar: EmbeddedCheckoutRegistrar,
     eventReporter: EventReporter,
 ) : ViewModel() {
     init {
         eventReporter.onInit()
     }
 
     override fun onCleared() {
+        checkoutRegistrar.unregisterIfNeeded()
         customViewModelScope.cancel()
     }
 

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/form/FormActivity.kt

@@ -1,3 +1,5 @@
+@file:OptIn(com.stripe.android.paymentelement.CheckoutSessionPreview::class)
+
 package com.stripe.android.paymentelement.embedded.form
 
 import android.app.Activity
@@ -13,6 +15,7 @@ import androidx.compose.runtime.remember
 import androidx.lifecycle.lifecycleScope
 import com.stripe.android.checkout.CheckoutInstances
 import com.stripe.android.common.ui.BottomSheetScaffold
+import com.stripe.android.lpmfoundations.paymentmethod.IntegrationMetadata
 import com.stripe.android.common.ui.ElementsBottomSheetLayout
 import com.stripe.android.paymentelement.embedded.sheet.EmbeddedNavigator
 import com.stripe.android.paymentelement.embedded.sheet.SheetActivityRegistrar
@@ -145,7 +148,12 @@ internal class FormActivity : AppCompatActivity() {
         super.onDestroy()
 
         if (isFinishing) {
-            CheckoutInstances.markIntegrationDismissed(args?.paymentMethodMetadata)
+            val key = (args?.paymentMethodMetadata?.integrationMetadata
+                as? IntegrationMetadata.CheckoutSession)?.instancesKey
+            if (key != null) {
+                CheckoutInstances[key]?.markIntegrationDismissed()
+                // No unregister: EmbeddedPaymentElement owns the registration lifetime.
+            }
             if (::eventReporter.isInitialized) {
                 eventReporter.onDismiss()
             }

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/sheet/EmbeddedSheetActivity.kt

@@ -1,3 +1,5 @@
+@file:OptIn(com.stripe.android.paymentelement.CheckoutSessionPreview::class)
+
 package com.stripe.android.paymentelement.embedded.sheet
 
 import android.app.Activity
@@ -24,6 +26,7 @@ import androidx.compose.ui.unit.dp
 import androidx.lifecycle.lifecycleScope
 import com.stripe.android.checkout.CheckoutInstances
 import com.stripe.android.common.ui.BottomSheetScaffold
+import com.stripe.android.lpmfoundations.paymentmethod.IntegrationMetadata
 import com.stripe.android.common.ui.ElementsBottomSheetLayout
 import com.stripe.android.paymentelement.embedded.EmbeddedSelectionHolder
 import com.stripe.android.paymentsheet.CustomerStateHolder
@@ -192,7 +195,12 @@ internal class EmbeddedSheetActivity : AppCompatActivity() {
         super.onDestroy()
 
         if (isFinishing) {
-            CheckoutInstances.markIntegrationDismissed(args?.paymentMethodMetadata)
+            val key = (args?.paymentMethodMetadata?.integrationMetadata
+                as? IntegrationMetadata.CheckoutSession)?.instancesKey
+            if (key != null) {
+                CheckoutInstances[key]?.markIntegrationDismissed()
+                // No unregister: EmbeddedPaymentElement owns the registration lifetime.
+            }
             if (::eventReporter.isInitialized) {
                 eventReporter.onDismiss()
             }

paymentsheet/src/main/java/com/stripe/android/paymentsheet/PaymentOptionsActivity.kt

@@ -1,3 +1,5 @@
+@file:OptIn(com.stripe.android.paymentelement.CheckoutSessionPreview::class)
+
 package com.stripe.android.paymentsheet
 
 import android.content.Intent
@@ -11,6 +13,7 @@ import androidx.compose.runtime.getValue
 import androidx.lifecycle.ViewModelProvider
 import com.stripe.android.checkout.CheckoutInstances
 import com.stripe.android.common.ui.ElementsBottomSheetLayout
+import com.stripe.android.lpmfoundations.paymentmethod.IntegrationMetadata
 import com.stripe.android.paymentsheet.ui.BaseSheetActivity
 import com.stripe.android.paymentsheet.ui.PaymentSheetScreen
 import com.stripe.android.paymentsheet.utils.applicationIsTaskOwner
@@ -86,7 +89,12 @@ internal class PaymentOptionsActivity : BaseSheetActivity<PaymentOptionsActivity
     override fun onDestroy() {
         super.onDestroy()
         if (isFinishing) {
-            CheckoutInstances.markIntegrationDismissed(starterArgs?.state?.paymentMethodMetadata)
+            val key = (starterArgs?.state?.paymentMethodMetadata?.integrationMetadata
+                as? IntegrationMetadata.CheckoutSession)?.instancesKey
+            if (key != null) {
+                CheckoutInstances[key]?.markIntegrationDismissed()
+                // No unregister: FlowController owns the registration lifetime.
+            }
         }
     }
 

paymentsheet/src/main/java/com/stripe/android/paymentsheet/PaymentSheet.kt

@@ -1,3 +1,5 @@
+@file:OptIn(com.stripe.android.paymentelement.CheckoutSessionPreview::class)
+
 package com.stripe.android.paymentsheet
 
 import android.content.Context

paymentsheet/src/main/java/com/stripe/android/paymentsheet/PaymentSheetActivity.kt

@@ -1,3 +1,5 @@
+@file:OptIn(com.stripe.android.paymentelement.CheckoutSessionPreview::class)
+
 package com.stripe.android.paymentsheet
 
 import android.app.Activity
@@ -11,7 +13,6 @@ import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.getValue
 import androidx.lifecycle.ViewModelProvider
 import com.stripe.android.PaymentConfiguration
-import com.stripe.android.checkout.CheckoutInstances
 import com.stripe.android.common.model.asCommonConfiguration
 import com.stripe.android.common.ui.ElementsBottomSheetLayout
 import com.stripe.android.paymentsheet.ui.BaseSheetActivity
@@ -95,12 +96,6 @@ internal class PaymentSheetActivity : BaseSheetActivity<PaymentSheetResult>() {
         }
     }
 
-    override fun onDestroy() {
-        super.onDestroy()
-        if (isFinishing && starterArgs != null) {
-            CheckoutInstances.markIntegrationDismissed(viewModel.paymentMethodMetadata.value)
-        }
-    }
 
     override fun setActivityResult(result: PaymentSheetResult) {
         setResult(