Add WSLAContainer flag to auto delete container after exit

src/windows/wslasession/WSLAContainer.cpp

@@ -267,8 +267,6 @@ std::string SerializeContainerMetadata(const WSLAContainerMetadataV1& metadata)
 
 } // namespace
 
-static constexpr DWORD stopTimeout = 60000; // 60 seconds
-
 WSLAContainerImpl::WSLAContainerImpl(
     WSLAVirtualMachine* parentVM,
     std::string&& Id,
@@ -338,36 +336,7 @@ WSLAContainerImpl::~WSLAContainerImpl()
     }
 
     m_containerEvents.Reset();
-
-    // Disconnect from the COM instance. After this returns, no COM calls can be made to this instance.
-    m_comWrapper->Disconnect();
-
-    // Stop running containers.
-    if (m_state == WslaContainerStateRunning)
-    {
-        try
-        {
-            Stop(WSLASignalSIGKILL, stopTimeout);
-        }
-        CATCH_LOG();
-    }
-
-    // Release port mappings.
-    std::set<uint16_t> allocatedGuestPorts;
-    for (const auto& e : m_mappedPorts)
-    {
-        WI_ASSERT(e.MappedToHost);
-
-        try
-        {
-            m_parentVM->UnmapPort(e.Family, e.HostPort, e.VmPort);
-        }
-        CATCH_LOG();
-
-        allocatedGuestPorts.insert(e.VmPort);
-    }
-
-    m_parentVM->ReleasePorts(allocatedGuestPorts);
+    ReleaseResources();
 }
 
 void WSLAContainerImpl::OnProcessReleased(DockerExecProcessControl* process)
@@ -498,7 +467,6 @@ void WSLAContainerImpl::OnEvent(ContainerEvent event, std::optional<int> exitCod
     if (event == ContainerEvent::Stop)
     {
         THROW_HR_IF(E_UNEXPECTED, !exitCode.has_value());
-
         std::lock_guard<std::recursive_mutex> lock(m_lock);
         m_state = WslaContainerStateExited;
 
@@ -510,6 +478,11 @@ void WSLAContainerImpl::OnEvent(ContainerEvent event, std::optional<int> exitCod
         }
 
         m_processes.clear();
+
+        if (WI_IsFlagSet(m_containerFlags, WSLAContainerFlagsRm))
+        {
+            Delete();
+        }
     }
 
     WSL_LOG(
@@ -547,15 +520,18 @@ void WSLAContainerImpl::Stop(WSLASignal Signal, LONGLONG TimeoutSeconds)
     catch (const DockerHTTPException& e)
     {
         // HTTP 304 is returned when the container is already stopped.
-        if (e.StatusCode() == 304)
+        if (e.StatusCode() != 304)
         {
-            return;
+            THROW_DOCKER_USER_ERROR_MSG(e, "Failed to stop container '%hs'", m_id.c_str());
         }
-
-        THROW_DOCKER_USER_ERROR_MSG(e, "Failed to stop container '%hs'", m_id.c_str());
     }
 
     m_state = WslaContainerStateExited;
+
+    if (WI_IsFlagSet(m_containerFlags, WSLAContainerFlagsRm))
+    {
+        Delete();
+    }
 }
 
 void WSLAContainerImpl::Delete()
@@ -576,7 +552,7 @@ void WSLAContainerImpl::Delete()
     }
     CATCH_AND_THROW_DOCKER_USER_ERROR("Failed to delete container '%hs'", m_id.c_str());
 
-    UnmountVolumes(m_mountedVolumes, *m_parentVM);
+    ReleaseResources();
 
     m_state = WslaContainerStateDeleted;
 }
@@ -896,6 +872,7 @@ std::unique_ptr<WSLAContainerImpl> WSLAContainerImpl::Create(
 
     // Build WSLA metadata to store in a label for recovery on Open().
     WSLAContainerMetadataV1 metadata;
+    metadata.Flags = containerOptions.Flags;
     metadata.InitProcessFlags = containerOptions.InitProcessOptions.Flags;
     metadata.Volumes = volumes;
     metadata.Ports = ports;
@@ -971,7 +948,7 @@ std::unique_ptr<WSLAContainerImpl> WSLAContainerImpl::Open(
         ioRelay,
         DockerStateToWSLAState(dockerContainer.State),
         metadata.InitProcessFlags,
-        WSLAContainerFlagsNone);
+        metadata.Flags);
 
     errorCleanup.release();
     volumeErrorCleanup.release();
@@ -1080,6 +1057,44 @@ std::unique_ptr<RelayedProcessIO> WSLAContainerImpl::CreateRelayedProcessIO(wil:
     return std::make_unique<RelayedProcessIO>(std::move(fds));
 }
 
+void WSLAContainerImpl::ReleaseResources()
+{
+    std::lock_guard<std::recursive_mutex> lock(m_lock);
+
+    // Disconnect the COM wrapper so no new RPC calls can reach this container.
+    if (m_comWrapper)
+    {
+        m_comWrapper->Disconnect();
+        m_comWrapper.Reset();
+    }
+
+    // Unmount volumes.
+    UnmountVolumes(m_mountedVolumes, *m_parentVM);
+    m_mountedVolumes.clear();
+
+    // Unmap and release ports.
+    std::set<uint16_t> allocatedGuestPorts;
+    for (const auto& e : m_mappedPorts)
+    {
+        WI_ASSERT(e.MappedToHost);
+
+        try
+        {
+            m_parentVM->UnmapPort(e.Family, e.HostPort, e.VmPort);
+        }
+        CATCH_LOG();
+
+        allocatedGuestPorts.insert(e.VmPort);
+    }
+
+    if (!allocatedGuestPorts.empty())
+    {
+        m_parentVM->ReleasePorts(allocatedGuestPorts);
+    }
+
+    m_mappedPorts.clear();
+}
+
 WSLAContainer::WSLAContainer(WSLAContainerImpl* impl, std::function<void(const WSLAContainerImpl*)>&& OnDeleted) :
     COMImplClass<WSLAContainerImpl>(impl), m_onDeleted(std::move(OnDeleted))
 {

src/windows/wslasession/WSLAContainer.h

@@ -74,6 +74,18 @@ class WSLAContainerImpl
 
     const std::string& ID() const noexcept;
 
+    // Called when the container stop event is observed so the
+    // implementation can update its internal state and notify
+    // any exec processes.
+    void OnStopped();
+
+    // Returns the container flags used to decide whether to
+    // auto-delete the container on stop.
+    WSLAContainerFlags Flags() const noexcept
+    {
+        return m_containerFlags;
+    }
+
     static std::unique_ptr<WSLAContainerImpl> Create(
         const WSLA_CONTAINER_OPTIONS& Options,
         WSLAVirtualMachine& parentVM,
@@ -93,6 +105,7 @@ class WSLAContainerImpl
 private:
     void OnEvent(ContainerEvent event, std::optional<int> exitCode);
     void WaitForContainerEvent();
+    void ReleaseResources();
     std::unique_ptr<RelayedProcessIO> CreateRelayedProcessIO(wil::unique_handle&& stream, WSLAProcessFlags flags);
 
     wsl::windows::common::wsla_schema::InspectContainer BuildInspectContainer(const wsl::windows::common::docker_schema::InspectContainer& dockerInspect);

src/windows/wslasession/WSLAContainerMetadata.h

@@ -48,11 +48,12 @@ struct WSLAVolumeMount
 
 struct WSLAContainerMetadataV1
 {
+    WSLAContainerFlags Flags{WSLAContainerFlagsNone};
     WSLAProcessFlags InitProcessFlags{WSLAProcessFlagsNone};
     std::vector<WSLAPortMapping> Ports;
     std::vector<WSLAVolumeMount> Volumes;
 
-    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(WSLAContainerMetadataV1, InitProcessFlags, Ports, Volumes);
+    NLOHMANN_DEFINE_TYPE_INTRUSIVE_WITH_DEFAULT(WSLAContainerMetadataV1, Flags, InitProcessFlags, Ports, Volumes);
 };
 
 struct WSLAContainerMetadata

src/windows/wslasession/WSLASession.cpp

@@ -902,6 +902,9 @@ try
 
     // Look for an exact ID match first.
     std::lock_guard lock{m_lock};
+
+    // Purge containers that were auto-deleted via OnEvent (--rm).
+    std::erase_if(m_containers, [](const auto& e) { return e->State() == WslaContainerStateDeleted; });
     auto it = std::ranges::find_if(m_containers, [Id](const auto& e) { return e->ID() == Id; });
 
     // If no match is found, call Inspect() so that partial IDs and names are matched.
@@ -942,6 +945,9 @@ try
 
     std::lock_guard lock{m_lock};
 
+    // Purge containers that were auto-deleted via OnEvent (--rm).
+    std::erase_if(m_containers, [](const auto& e) { return e->State() == WslaContainerStateDeleted; });
+
     auto output = wil::make_unique_cotaskmem<WSLA_CONTAINER[]>(m_containers.size());
 
     size_t index = 0;

src/windows/wslasession/WSLAVirtualMachine.cpp

@@ -747,6 +747,7 @@ try
 CATCH_RETURN();
 
 HRESULT WSLAVirtualMachine::UnmountWindowsFolder(_In_ LPCSTR LinuxPath)
+try
 {
     std::lock_guard lock(m_lock);
 
@@ -767,6 +768,7 @@ HRESULT WSLAVirtualMachine::UnmountWindowsFolder(_In_ LPCSTR LinuxPath)
 
     return S_OK;
 }
+CATCH_RETURN();
 
 void WSLAVirtualMachine::MountGpuLibraries(_In_ LPCSTR LibrariesMountPoint, _In_ LPCSTR DriversMountpoint)
 {

test/windows/WSLATests.cpp

@@ -4191,4 +4191,83 @@ class WSLATests
 
         VERIFY_ARE_EQUAL(result.Output[1], std::format("{}\n", expectedOrder));
     }
+
+    TEST_METHOD(ContainerAutoRemove)
+    {
+        WSL2_TEST_ONLY();
+        SKIP_TEST_ARM64();
+
+        // Test that a container with the Rm flag is automatically deleted on Stop().
+        {
+            WSLAContainerLauncher launcher("debian:latest", "test-auto-remove", {"/bin/cat"}, {}, {}, WSLAProcessFlagsStdin);
+            launcher.SetContainerFlags(WSLAContainerFlagsRm);
+
+            auto container = launcher.Launch(*m_defaultSession);
+
+            VERIFY_ARE_EQUAL(container.State(), WslaContainerStateRunning);
+            VERIFY_SUCCEEDED(container.Get().Stop(WSLASignalSIGKILL, 0));
+
+            VERIFY_ARE_EQUAL(container.Get().Delete(), RPC_E_DISCONNECTED);
+
+            wil::com_ptr<IWSLAContainer> notFound;
+            VERIFY_ARE_EQUAL(m_defaultSession->OpenContainer("test-auto-remove", &notFound), HRESULT_FROM_WIN32(ERROR_NOT_FOUND));
+        }
+
+        // Test that a container with the Rm flag is automatically deleted when the init process is killed.
+        {
+            WSLAContainerLauncher launcher("debian:latest", "test-auto-remove", {"/bin/cat"}, {}, {}, WSLAProcessFlagsStdin);
+            launcher.SetContainerFlags(WSLAContainerFlagsRm);
+
+            // Prevent container from being deleted when handle is closed so we can verify auto-remove behavior.
+            auto container = launcher.Launch(*m_defaultSession);
+            auto process = container.GetInitProcess();
+
+            VERIFY_ARE_EQUAL(container.State(), WslaContainerStateRunning);
+            VERIFY_SUCCEEDED(process.Get().Signal(WSLASignalSIGKILL));
+            process.Wait();
+
+            VERIFY_ARE_EQUAL(container.Get().Delete(), RPC_E_DISCONNECTED);
+
+            wil::com_ptr<IWSLAContainer> notFound;
+            VERIFY_ARE_EQUAL(m_defaultSession->OpenContainer("test-auto-remove", &notFound), HRESULT_FROM_WIN32(ERROR_NOT_FOUND));
+        }
+
+        // Test that the container autoremove flag is applied when the container exits on its own.
+        {
+            WSLAContainerLauncher launcher("debian:latest", "test-hostname", {"/bin/sh", "-c", "echo foo"});
+            launcher.SetContainerFlags(WSLAContainerFlagsRm);
+
+            auto container = launcher.Launch(*m_defaultSession);
+            auto process = container.GetInitProcess();
+            process.Wait();
+
+            VERIFY_ARE_EQUAL(container.Get().Delete(), RPC_E_DISCONNECTED);
+
+            wil::com_ptr<IWSLAContainer> notFound;
+            VERIFY_ARE_EQUAL(m_defaultSession->OpenContainer("test-auto-remove", &notFound), HRESULT_FROM_WIN32(ERROR_NOT_FOUND));
+        }
+
+        // Test that the Rm flag is persisted across wsla sessions.
+        {
+            {
+                WSLAContainerLauncher launcher("debian:latest", "test-auto-remove", {"/bin/cat"}, {}, {}, WSLAProcessFlagsStdin);
+                launcher.SetContainerFlags(WSLAContainerFlagsRm);
+
+                auto container = launcher.Create(*m_defaultSession);
+                container.SetDeleteOnClose(false);
+
+                ResetTestSession();
+            }
+
+            auto container = OpenContainer(m_defaultSession.get(), "test-auto-remove");
+            VERIFY_SUCCEEDED(container.Get().Start(WSLAContainerStartFlagsNone));
+            VERIFY_SUCCEEDED(container.Get().Stop(WSLASignalSIGKILL, 0));
+
+            // verifyContainerDeleted("test-auto-remove");
+            VERIFY_ARE_EQUAL(container.Get().Delete(), RPC_E_DISCONNECTED);
+
+            wil::com_ptr<IWSLAContainer> notFound;
+            VERIFY_ARE_EQUAL(m_defaultSession->OpenContainer("test-auto-remove", &notFound), HRESULT_FROM_WIN32(ERROR_NOT_FOUND));
+        }
+    }
 };