Use event instead of termination callback

msipackage/package.wix.in

@@ -306,14 +306,6 @@
                     </RegistryKey>
                 </RegistryKey>
 
-                <!-- ITerminationCallback-->
-                <RegistryKey Root="HKCR" Key="Interface\{7BC4E198-6531-4FA6-ADE2-5EF3D2A04DFE}">
-                    <RegistryValue Value="ITerminationCallback" Type="string" />
-                    <RegistryKey Key="ProxyStubClsid32">
-                        <RegistryValue Value="{4EA0C6DD-E9FF-48E7-994E-13A31D10DC60}" Type="string" />
-                    </RegistryKey>
-                </RegistryKey>
-
                 <!-- ICrashDumpCallback-->
                 <RegistryKey Root="HKCR" Key="Interface\{8C5A7B14-9D26-4FAE-AB31-7E5BC23F4801}">
                     <RegistryValue Value="ICrashDumpCallback" Type="string" />

src/windows/WslcSDK/CMakeLists.txt

@@ -1,7 +1,6 @@
 set(SOURCES
     IOCallback.cpp
     ProgressCallback.cpp
-    TerminationCallback.cpp
     CrashDumpCallback.cpp
     wslcsdk.cpp
     WslcsdkPrivate.cpp
@@ -10,7 +9,6 @@ set(HEADERS
     Defaults.h
     IOCallback.h
     ProgressCallback.h
-    TerminationCallback.h
     CrashDumpCallback.h
     wslcsdk.h
     WslcsdkPrivate.h

src/windows/WslcSDK/WslcsdkPrivate.h

@@ -33,8 +33,6 @@ typedef struct WslcSessionOptionsInternal
 
     WslcVhdRequirements vhdRequirements;
     WslcSessionFeatureFlags featureFlags;
-    WslcSessionTerminationCallback terminationCallback;
-    PVOID terminationCallbackContext;
 } WslcSessionOptionsInternal;
 
 static_assert(sizeof(WslcSessionOptionsInternal) == WSLC_SESSION_OPTIONS_SIZE, "WSLC_SESSION_OPTIONS_INTERNAL size mismatch");
@@ -107,7 +105,6 @@ const WslcContainerOptionsInternal* GetInternalType(const WslcContainerSettings*
 struct WslcSessionImpl
 {
     wil::com_ptr<IWSLCSession> session;
-    wil::com_ptr<ITerminationCallback> terminationCallback;
 };
 
 WslcSessionImpl* GetInternalType(WslcSession handle);

src/windows/WslcSDK/winrt/Session.cpp

@@ -52,12 +52,16 @@ void Session::Start()
         throw winrt::hresult_illegal_method_call(L"Session has already been started");
     }
 
-    winrt::check_hresult(WslcSetSessionSettingsTerminationCallback(GetStructPointer(m_settings), TerminatedCallback, /* context */ this));
-
     wil::unique_cotaskmem_string errorMessage;
     auto hr = WslcCreateSession(GetStructPointer(m_settings), m_session.put(), errorMessage.put());
     THROW_MSG_IF_FAILED(hr, errorMessage);
     m_settings = nullptr;
+
+    winrt::check_hresult(WslcGetSessionTerminationEvent(m_session.get(), m_terminationEvent.put()));
+
+    m_terminationWait.reset(CreateThreadpoolWait(&Session::OnTerminated, this, nullptr));
+    THROW_LAST_ERROR_IF_NULL(m_terminationWait);
+    SetThreadpoolWait(m_terminationWait.get(), m_terminationEvent.get(), nullptr);
 }
 
 void Session::EnsureStarted() const
@@ -300,11 +304,15 @@ WslcSession Session::ToHandle()
     return m_session.get();
 }
 
-void CALLBACK Session::TerminatedCallback(_In_ WslcSessionTerminationReason reason, _In_opt_ PVOID context) noexcept
+void CALLBACK Session::OnTerminated(PTP_CALLBACK_INSTANCE /* instance */, PVOID context, PTP_WAIT /* wait */, TP_WAIT_RESULT /* waitResult */) noexcept
 {
     try
     {
         auto session = static_cast<Session*>(context);
+
+        WslcSessionTerminationReason reason = WSLC_SESSION_TERMINATION_REASON_UNKNOWN;
+        LOG_IF_FAILED(WslcGetSessionTerminationReason(session->m_session.get(), &reason));
+
         session->m_terminatedEvent(static_cast<SessionTerminationReason>(reason));
     }
     CATCH_LOG();

src/windows/WslcSDK/winrt/Session.h

@@ -45,12 +45,15 @@ struct Session : SessionT<Session>
     void EnsureStarted() const;
     winrt::Microsoft::WSL::Containers::SessionSettings m_settings; // Only kept until Start() is called
 
-    static void CALLBACK TerminatedCallback(_In_ WslcSessionTerminationReason reason, _In_opt_ PVOID context) noexcept;
+    // Threadpool callback that raises the Terminated event once the session's termination handle is signaled.
+    static void CALLBACK OnTerminated(PTP_CALLBACK_INSTANCE instance, PVOID context, PTP_WAIT wait, TP_WAIT_RESULT waitResult) noexcept;
 
-    // Releasing the session handle may trigger the termination callback.
-    // Keep these two in this order so that the session handle is released before the termination event is destructed.
     winrt::event<winrt::Microsoft::WSL::Containers::SessionTerminationHandler> m_terminatedEvent;
     wil::unique_any<WslcSession, decltype(&WslcReleaseSession), &WslcReleaseSession> m_session{nullptr};
+
+    // Bridges the one-off termination event surfaced by the SDK to the WinRT Terminated event.
+    wil::unique_handle m_terminationEvent;
+    wil::unique_threadpool_wait m_terminationWait;
 };
 } // namespace winrt::Microsoft::WSL::Containers::implementation
 namespace winrt::Microsoft::WSL::Containers::factory_implementation {

src/windows/WslcSDK/wslcsdk.cpp

@@ -17,7 +17,6 @@ Module Name:
 #include "WslcsdkPrivate.h"
 #include "Defaults.h"
 #include "ProgressCallback.h"
-#include "TerminationCallback.h"
 #include "CrashDumpCallback.h"
 #include "Localization.h"
 #include "WslInstall.h"
@@ -435,12 +434,6 @@ try
     runtimeSettings.MemoryMb = internalType->memoryMb;
     runtimeSettings.BootTimeoutMs = internalType->timeoutMS;
     runtimeSettings.NetworkingMode = WSLCNetworkingModeVirtioProxy;
-    auto terminationCallback = TerminationCallback::CreateIf(internalType);
-    if (terminationCallback)
-    {
-        result->terminationCallback.attach(terminationCallback.as<ITerminationCallback>().detach());
-        runtimeSettings.TerminationCallback = terminationCallback.get();
-    }
     runtimeSettings.FeatureFlags = ConvertFlags(internalType->featureFlags);
     WI_SetFlag(runtimeSettings.FeatureFlags, WslcFeatureFlagsVirtioFs);
     WI_SetFlag(runtimeSettings.FeatureFlags, WslcFeatureFlagsDnsTunneling);
@@ -587,15 +580,39 @@ try
 }
 CATCH_RETURN();
 
-STDAPI WslcSetSessionSettingsTerminationCallback(
-    _In_ WslcSessionSettings* sessionSettings, _In_opt_ WslcSessionTerminationCallback terminationCallback, _In_opt_ PVOID terminationContext)
+STDAPI WslcGetSessionTerminationEvent(_In_ WslcSession session, _Out_ HANDLE* terminationEvent)
 try
 {
-    auto internalType = CheckAndGetInternalType(sessionSettings);
-    RETURN_HR_IF(E_INVALIDARG, terminationCallback == nullptr && terminationContext != nullptr);
+    RETURN_HR_IF_NULL(E_POINTER, terminationEvent);
+    *terminationEvent = nullptr;
+
+    auto internalType = CheckAndGetInternalType(session);
+    RETURN_HR_IF_NULL(HRESULT_FROM_WIN32(ERROR_INVALID_STATE), internalType->session);
+
+    RETURN_HR(internalType->session->GetTerminationEvent(terminationEvent));
+}
+CATCH_RETURN();
+
+STDAPI WslcGetSessionTerminationReason(_In_ WslcSession session, _Out_ WslcSessionTerminationReason* reason)
+try
+{
+    static_assert(
+        WSLC_SESSION_TERMINATION_REASON_UNKNOWN == WSLCVirtualMachineTerminationReasonUnknown &&
+            WSLC_SESSION_TERMINATION_REASON_SHUTDOWN == WSLCVirtualMachineTerminationReasonShutdown &&
+            WSLC_SESSION_TERMINATION_REASON_CRASHED == WSLCVirtualMachineTerminationReasonCrashed,
+        "Termination reason enum values mismatch.");
+
+    RETURN_HR_IF_NULL(E_POINTER, reason);
+    *reason = WSLC_SESSION_TERMINATION_REASON_UNKNOWN;
+
+    auto internalType = CheckAndGetInternalType(session);
+    RETURN_HR_IF_NULL(HRESULT_FROM_WIN32(ERROR_INVALID_STATE), internalType->session);
+
+    WSLCVirtualMachineTerminationReason runtimeReason = WSLCVirtualMachineTerminationReasonUnknown;
+    wil::unique_cotaskmem_string details;
+    RETURN_IF_FAILED(internalType->session->GetTerminationReason(&runtimeReason, &details));
 
-    internalType->terminationCallback = terminationCallback;
-    internalType->terminationCallbackContext = terminationContext;
+    *reason = static_cast<WslcSessionTerminationReason>(runtimeReason);
 
     return S_OK;
 }
@@ -649,10 +666,7 @@ try
 {
     auto internalType = CheckAndGetInternalTypeUniquePointer(session);
 
-    // Drop the session before the termination callback, in case session destruction triggers
-    // the termination callback.
     internalType->session.reset();
-    internalType->terminationCallback.reset();
 
     return S_OK;
 }

src/windows/WslcSDK/wslcsdk.def

@@ -16,7 +16,8 @@ WslcReleaseContainer
 WslcReleaseProcess
 
 WslcSetSessionSettingsFeatureFlags
-WslcSetSessionSettingsTerminationCallback
+WslcGetSessionTerminationEvent
+WslcGetSessionTerminationReason
 WslcSetSessionSettingsCpuCount
 WslcSetSessionSettingsMemory
 WslcSetSessionSettingsTimeout

src/windows/WslcSDK/wslcsdk.h

@@ -42,7 +42,7 @@ EXTERN_C_START
 #define WSLC_E_REGISTRY_BLOCKED_BY_POLICY MAKE_HRESULT(SEVERITY_ERROR, FACILITY_ITF, WSLC_E_BASE + 13) /* 0x8004060D */
 
 // Session values
-#define WSLC_SESSION_OPTIONS_SIZE 88
+#define WSLC_SESSION_OPTIONS_SIZE 72
 #define WSLC_SESSION_OPTIONS_ALIGNMENT 8
 
 typedef struct WslcSessionSettings
@@ -123,8 +123,6 @@ typedef enum WslcSessionTerminationReason
     WSLC_SESSION_TERMINATION_REASON_CRASHED = 2,
 } WslcSessionTerminationReason;
 
-typedef __callback void(CALLBACK* WslcSessionTerminationCallback)(_In_ WslcSessionTerminationReason reason, _In_opt_ PVOID context);
-
 typedef struct WslcSessionCrashDumpInfo
 {
     _Field_z_ PCWSTR dumpPath;
@@ -153,9 +151,8 @@ STDAPI WslcSetSessionSettingsVhd(_In_ WslcSessionSettings* sessionSettings, _In_
 
 STDAPI WslcSetSessionSettingsFeatureFlags(_In_ WslcSessionSettings* sessionSettings, _In_ WslcSessionFeatureFlags flags);
 
-// Pass in Null for callback to clear the termination callback
-STDAPI WslcSetSessionSettingsTerminationCallback(
-    _In_ WslcSessionSettings* sessionSettings, _In_opt_ WslcSessionTerminationCallback terminationCallback, _In_opt_ PVOID terminationContext);
+STDAPI WslcGetSessionTerminationEvent(_In_ WslcSession session, _Out_ HANDLE* terminationEvent);
+STDAPI WslcGetSessionTerminationReason(_In_ WslcSession session, _Out_ WslcSessionTerminationReason* reason);
 
 STDAPI WslcTerminateSession(_In_ WslcSession session);
 STDAPI WslcReleaseSession(_In_ WslcSession session);

src/windows/service/exe/HcsVirtualMachine.cpp

@@ -285,12 +285,6 @@ HcsVirtualMachine::HcsVirtualMachine(_In_ const WSLCSessionSettings* Settings)
         m_guestDeviceManager = std::make_shared<::GuestDeviceManager>(m_vmIdString, m_vmId);
     }
 
-    // Configure termination callback
-    if (Settings->TerminationCallback)
-    {
-        m_terminationCallback = Settings->TerminationCallback;
-    }
-
     hcs::RegisterCallback(m_computeSystem.get(), &HcsVirtualMachine::OnVmExitCallback, this);
 
     // Create a listening socket for mini_init to connect to once the VM is running.
@@ -692,8 +686,6 @@ CATCH_LOG()
 
 void HcsVirtualMachine::OnExit(const HCS_EVENT* Event)
 {
-    m_vmExitEvent.SetEvent();
-
     const auto exitStatus = wsl::shared::FromJson<wsl::windows::common::hcs::SystemExitStatus>(Event->EventData);
 
     auto reason = WSLCVirtualMachineTerminationReasonUnknown;
@@ -715,11 +707,28 @@ void HcsVirtualMachine::OnExit(const HCS_EVENT* Event)
         }
     }
 
-    if (m_terminationCallback)
+    // Cache the termination reason and details before signaling the exit event.
     {
-        LOG_IF_FAILED(m_terminationCallback->OnTermination(reason, Event->EventData));
+        std::lock_guard lock(m_lock);
+        m_terminationReason = reason;
+        m_terminationDetails = Event->EventData;
     }
+
+    m_vmExitEvent.SetEvent();
+}
+
+HRESULT HcsVirtualMachine::GetTerminationReason(_Out_ WSLCVirtualMachineTerminationReason* Reason, _Out_ LPWSTR* Details)
+try
+{
+    RETURN_HR_IF(E_POINTER, Reason == nullptr || Details == nullptr);
+
+    std::lock_guard lock(m_lock);
+    *Reason = m_terminationReason;
+    *Details = wil::make_cotaskmem_string(m_terminationDetails.c_str()).release();
+
+    return S_OK;
 }
+CATCH_RETURN()
 
 void HcsVirtualMachine::OnCrash(const HCS_EVENT* Event)
 {

src/windows/service/exe/HcsVirtualMachine.h

@@ -48,6 +48,7 @@ class HcsVirtualMachine
     IFACEMETHOD(RemoveShare)(_In_ REFGUID ShareId) override;
     IFACEMETHOD(ApplyGuestCapabilities)(_In_ const WSLCGuestCapabilities* Capabilities) override;
     IFACEMETHOD(GetTerminationEvent)(_Out_ HANDLE* Event) override;
+    IFACEMETHOD(GetTerminationReason)(_Out_ WSLCVirtualMachineTerminationReason* Reason, _Out_ LPWSTR* Details) override;
 
 private:
     struct DiskInfo
@@ -104,7 +105,9 @@ class HcsVirtualMachine
     std::atomic<bool> m_vmSavedStateCaptured = false;
     std::atomic<bool> m_crashLogCaptured = false;
 
-    wil::com_ptr<ITerminationCallback> m_terminationCallback;
+    // Termination reason and details, cached when the VM exits (see OnExit). Guarded by m_lock.
+    WSLCVirtualMachineTerminationReason m_terminationReason{WSLCVirtualMachineTerminationReasonUnknown};
+    std::wstring m_terminationDetails;
 };
 
 //