diff --git a/src/WinRT.Runtime2/InteropServices/ObjectReference/ContextCallback.cs b/src/WinRT.Runtime2/InteropServices/ObjectReference/ContextCallback.cs
index 01941a99f..b87630265 100644
--- a/src/WinRT.Runtime2/InteropServices/ObjectReference/ContextCallback.cs
+++ b/src/WinRT.Runtime2/InteropServices/ObjectReference/ContextCallback.cs
@@ -11,6 +11,12 @@ namespace WindowsRuntime.InteropServices;
///
internal static unsafe class ContextCallback
{
+ ///
+ /// Storage for the user-provided state for .
+ ///
+ [ThreadStatic]
+ private static object? LocalContextCallbackState;
+
///
/// Calls the given callback in the right context, and returns the result of that invocation.
///
@@ -38,21 +44,10 @@ public static HRESULT CallInContextUnsafe(
return WellKnownErrorCodes.S_OK;
}
- ComCallData comCallData;
- comCallData.dwDispid = 0;
- comCallData.dwReserved = 0;
-
- CallbackData callbackData;
- callbackData.Callback = callback;
- callbackData.State = state;
-
- // We can just store a pointer to the callback to invoke in the context,
- // so we don't need to allocate another closure or anything. The callback
- // will be kept alive automatically, because 'comCallData' is address exposed.
- // In 'InvokeCallback' below, we then invoke the current caller-provided callback.
- comCallData.pUserDefined = &callbackData;
-
- // Stub to invoke on the target context
+ // Native method that invokes the callback on the target context. The state object is guaranteed to be pinned,
+ // so we can access it from a pointer. Note that the object will be stored in a static field, and it will not
+ // be on the stack of the original thread, so it's safe with respect to cross-thread access of managed objects.
+ // See: https://github.com/dotnet/runtime/blob/main/docs/design/specs/Memory-model.md#cross-thread-access-to-local-variables.
[UnmanagedCallersOnly]
static int InvokeCallback(ComCallData* comCallData)
{
@@ -60,7 +55,7 @@ static int InvokeCallback(ComCallData* comCallData)
{
CallbackData* callbackData = (CallbackData*)comCallData->pUserDefined;
- callbackData->Callback(callbackData->State);
+ callbackData->Callback(*callbackData->StatePtr);
return WellKnownErrorCodes.S_OK;
}
@@ -70,10 +65,49 @@ static int InvokeCallback(ComCallData* comCallData)
}
}
+ ref object? localContextCallbackState = ref LocalContextCallbackState;
+
+ // Store the state object in the thread static to pass to the callback.
+ // A thread local is the most efficient solution for this, given that
+ // we need the state to be somewhere on the managed heap to be valid.
+ // The GC doesn't allow cross-thread access to managed stack variables.
+ // We're not using a volatile write, as it wouldn't actually help at all.
+ // Volatile writes disallow the write operations to be moved before memory
+ // operations that precede them, but they have nothing to say with respect
+ // to memory operations after them (whereas here we specifically need this
+ // write to not be reordered before following reads from that static field).
+ if (localContextCallbackState is null)
+ {
+ localContextCallbackState = state;
+ }
+ else
+ {
+ // In case we recursed on this thread, meaning the local storage already holds a state
+ // for some other caller above us in the stack, we can use a throwaway array to store
+ // the current state. This isn't very efficient, but in practice this case should not
+ // ever happen (it was validated in our entire test suite as well as in stress tests
+ // with the Microsoft Store app, and we never detected recursion on this code path).
+ // However just to be extra safe, we still want to keep this branch functional too.
+ object[] objects = [state];
+
+ localContextCallbackState = ref MemoryMarshal.GetArrayDataReference(objects)!;
+ }
+
HRESULT hresult;
+ // Pin the state storage, which we can now safely pass to the target thread
+ fixed (object* statePtr = &localContextCallbackState)
fixed (Guid* riid = &WellKnownWindowsInterfaceIIDs.IID_ICallbackWithNoReentrancyToApplicationSTA)
{
+ CallbackData callbackData;
+ callbackData.Callback = callback;
+ callbackData.StatePtr = statePtr;
+
+ ComCallData comCallData;
+ comCallData.dwDispid = 0;
+ comCallData.dwReserved = 0;
+ comCallData.pUserDefined = &callbackData;
+
// Marshal the supplied callback on the target context
hresult = IContextCallbackVftbl.ContextCallbackUnsafe(
thisPtr: contextCallbackPtr,
@@ -84,6 +118,9 @@ static int InvokeCallback(ComCallData* comCallData)
pUnk: null);
}
+ // Reset the static field to avoid keeping the state alive for longer
+ LocalContextCallbackState = null;
+
return hresult;
}
@@ -98,8 +135,8 @@ private struct CallbackData
public delegate*