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* Callback; /// - /// The additional argument to supply to . + /// A pointer to the additional argument to supply to . /// - public object State; + public object* StatePtr; } } \ No newline at end of file