Container lifecycle transitions

src/windows/wslcsession/DockerEventTracker.cpp

@@ -157,7 +157,11 @@ void DockerEventTracker::OnEvent(const std::string_view& event)
 void DockerEventTracker::OnContainerEvent(const nlohmann::json& parsed, const std::string& action, std::int64_t eventTime)
 {
     static std::map<std::string, ContainerEvent> events{
-        {"start", ContainerEvent::Start}, {"die", ContainerEvent::Stop}, {"destroy", ContainerEvent::Destroy}, {"exec_die", ContainerEvent::ExecDied}};
+        {"start", ContainerEvent::Start},
+        {"die", ContainerEvent::Stop},
+        {"destroy", ContainerEvent::Destroy},
+        {"exec_die", ContainerEvent::ExecDied},
+        {"restart", ContainerEvent::Restart}};
 
     auto actor = parsed.find("Actor");
     THROW_HR_IF_MSG(E_INVALIDARG, actor == parsed.end(), "Missing Actor in container event");
@@ -191,15 +195,23 @@ void DockerEventTracker::OnContainerEvent(const nlohmann::json& parsed, const st
         }
     }
 
-    std::lock_guard lock{m_lock};
-
-    for (const auto& e : m_containerCallbacks)
+    // Snapshot the matching callbacks so that they can be invoked without holding m_lock. Callbacks can register and
+    // unregister callbacks (a container that stops releases its exec processes), which would otherwise mutate the
+    // vector being iterated.
+    std::vector<std::shared_ptr<ContainerCallback>> callbacks;
     {
-        if (e.ContainerId == containerId && (!e.ExecId.has_value() || e.ExecId == execId))
+        std::lock_guard lock{m_lock};
+
+        for (const auto& e : m_containerCallbacks)
         {
-            e.Callback(it->second, exitCode, eventTime);
+            if (e->ContainerId == containerId && (!e->ExecId.has_value() || e->ExecId == execId))
+            {
+                callbacks.emplace_back(e);
+            }
         }
     }
+
+    InvokeCallbacks(callbacks, [&](const ContainerCallback& e) { e.Callback(it->second, exitCode, eventTime); });
 }
 
 void DockerEventTracker::OnVolumeEvent(const nlohmann::json& parsed, const std::string& action, std::int64_t eventTime)
@@ -220,12 +232,13 @@ void DockerEventTracker::OnVolumeEvent(const nlohmann::json& parsed, const std::
 
     auto volumeName = id->get<std::string>();
 
-    std::lock_guard lock{m_lock};
-
-    for (const auto& e : m_volumeCallbacks)
+    std::vector<std::shared_ptr<VolumeCallback>> callbacks;
     {
-        e.Callback(volumeName, it->second, eventTime);
+        std::lock_guard lock{m_lock};
+        callbacks = m_volumeCallbacks;
     }
+
+    InvokeCallbacks(callbacks, [&](const VolumeCallback& e) { e.Callback(volumeName, it->second, eventTime); });
 }
 
 void DockerEventTracker::WaitForObjectCreated(const std::string& ObjectId)
@@ -257,50 +270,70 @@ void DockerEventTracker::WaitForObjectCreated(const std::string& ObjectId)
 DockerEventTracker::EventTrackingReference DockerEventTracker::RegisterContainerStateUpdates(
     const std::string& ContainerId, ContainerStateChangeCallback&& Callback) noexcept
 {
-    std::lock_guard lock{m_lock};
-
     auto id = m_callbackId++;
-    m_containerCallbacks.emplace_back(id, ContainerId, std::optional<std::string>{}, std::move(Callback));
+    auto entry = std::make_shared<ContainerCallback>(id, std::string{ContainerId}, std::optional<std::string>{}, std::move(Callback));
+
+    std::lock_guard lock{m_lock};
+    m_containerCallbacks.emplace_back(std::move(entry));
 
     return EventTrackingReference{this, id};
 }
 
 DockerEventTracker::EventTrackingReference DockerEventTracker::RegisterExecStateUpdates(
     const std::string& ContainerId, const std::string& ExecId, ContainerStateChangeCallback&& Callback) noexcept
 {
-    std::lock_guard lock{m_lock};
-
     auto id = m_callbackId++;
-    m_containerCallbacks.emplace_back(id, ContainerId, ExecId, std::move(Callback));
+    auto entry = std::make_shared<ContainerCallback>(id, std::string{ContainerId}, std::optional<std::string>{ExecId}, std::move(Callback));
+
+    std::lock_guard lock{m_lock};
+    m_containerCallbacks.emplace_back(std::move(entry));
 
     return EventTrackingReference{this, id};
 }
 
 DockerEventTracker::EventTrackingReference DockerEventTracker::RegisterVolumeUpdates(VolumeEventCallback&& Callback) noexcept
 {
-    std::lock_guard lock{m_lock};
-
     auto id = m_callbackId++;
-    m_volumeCallbacks.emplace_back(id, std::move(Callback));
+    auto entry = std::make_shared<VolumeCallback>(id, std::move(Callback));
+
+    std::lock_guard lock{m_lock};
+    m_volumeCallbacks.emplace_back(std::move(entry));
 
     return EventTrackingReference{this, id};
 }
 
 void DockerEventTracker::UnregisterCallback(size_t Id) noexcept
 {
-    std::lock_guard lock{m_lock};
+    std::shared_ptr<CallbackRegistration> registration;
 
-    // Try container callbacks first.
-    auto containerRemove = std::ranges::remove_if(m_containerCallbacks, [Id](auto& entry) { return entry.CallbackId == Id; });
-    if (!containerRemove.empty())
     {
-        WI_ASSERT(containerRemove.size() == 1);
-        m_containerCallbacks.erase(containerRemove.begin(), containerRemove.end());
-        return;
+        std::lock_guard lock{m_lock};
+
+        auto matches = [Id](const auto& e) { return e->CallbackId == Id; };
+
+        // Try container callbacks first, then volume callbacks.
+        if (auto container = std::ranges::find_if(m_containerCallbacks, matches); container != m_containerCallbacks.end())
+        {
+            registration = std::move(*container);
+            m_containerCallbacks.erase(container);
+        }
+        else
+        {
+            auto volume = std::ranges::find_if(m_volumeCallbacks, matches);
+            WI_ASSERT(volume != m_volumeCallbacks.end());
+
+            if (volume != m_volumeCallbacks.end())
+            {
+                registration = std::move(*volume);
+                m_volumeCallbacks.erase(volume);
+            }
+        }
     }
 
-    // Then volume callbacks.
-    auto volumeRemove = std::ranges::remove_if(m_volumeCallbacks, [Id](auto& entry) { return entry.CallbackId == Id; });
-    WI_ASSERT(volumeRemove.size() == 1);
-    m_volumeCallbacks.erase(volumeRemove.begin(), volumeRemove.end());
-}
+    if (registration)
+    {
+        // Wait for any in-flight invocation to complete so the callback can't run once this returns.
+        std::lock_guard invokeLock{registration->InvokeLock};
+        registration->Unregistered = true;
+    }
+}

src/windows/wslcsession/DockerEventTracker.h

@@ -26,6 +26,7 @@ enum class ContainerEvent
 {
     Create,
     Start,
+    Restart,
     Stop,
     Exit,
     Destroy,
@@ -84,28 +85,68 @@ class DockerEventTracker
     void OnContainerEvent(const nlohmann::json& parsed, const std::string& action, std::int64_t eventTime);
     void OnVolumeEvent(const nlohmann::json& parsed, const std::string& action, std::int64_t eventTime);
 
-    struct ContainerCallback
+    // Callbacks are invoked without holding m_lock so that a callback can register or unregister callbacks, and so
+    // that a callback taking its own lock can't invert with a thread that registers a callback under that same lock.
+    struct CallbackRegistration
     {
-        size_t CallbackId;
-        std::string ContainerId;
-        std::optional<std::string> ExecId;
-        ContainerStateChangeCallback Callback;
+        NON_COPYABLE(CallbackRegistration);
+        NON_MOVABLE(CallbackRegistration);
+
+        CallbackRegistration(size_t Id) noexcept : CallbackId(Id)
+        {
+        }
+
+        const size_t CallbackId;
+
+        // Held while the callback runs so it can't be invoked once UnregisterCallback() returned for it.
+        // N.B. Recursive so a running callback can unregister itself.
+        std::recursive_mutex InvokeLock;
+        _Guarded_by_(InvokeLock) bool Unregistered = false;
     };
 
-    struct VolumeCallback
+    struct ContainerCallback : CallbackRegistration
     {
-        size_t CallbackId;
-        VolumeEventCallback Callback;
+        ContainerCallback(size_t Id, std::string&& ContainerId, std::optional<std::string>&& ExecId, ContainerStateChangeCallback&& Callback) :
+            CallbackRegistration(Id), ContainerId(std::move(ContainerId)), ExecId(std::move(ExecId)), Callback(std::move(Callback))
+        {
+        }
+
+        const std::string ContainerId;
+        const std::optional<std::string> ExecId;
+        const ContainerStateChangeCallback Callback;
     };
 
-    std::vector<ContainerCallback> m_containerCallbacks;
-    std::vector<VolumeCallback> m_volumeCallbacks;
+    struct VolumeCallback : CallbackRegistration
+    {
+        VolumeCallback(size_t Id, VolumeEventCallback&& Callback) : CallbackRegistration(Id), Callback(std::move(Callback))
+        {
+        }
+
+        const VolumeEventCallback Callback;
+    };
+
+    _Guarded_by_(m_lock) std::vector<std::shared_ptr<ContainerCallback>> m_containerCallbacks;
+    _Guarded_by_(m_lock) std::vector<std::shared_ptr<VolumeCallback>> m_volumeCallbacks;
+
+    // Invokes a snapshot of callbacks taken under m_lock, skipping registrations that have since been unregistered.
+    template <typename TCallback, typename TInvoke>
+    static void InvokeCallbacks(const std::vector<std::shared_ptr<TCallback>>& Callbacks, const TInvoke& Invoke)
+    {
+        for (const auto& e : Callbacks)
+        {
+            std::lock_guard invokeLock{e->InvokeLock};
+            if (!e->Unregistered)
+            {
+                Invoke(*e);
+            }
+        }
+    }
 
     _Guarded_by_(m_lock) std::unordered_set<std::string> m_createdObjects;
     _Guarded_by_(m_lock) wil::unique_event m_objectCreated { wil::EventOptions::ManualReset };
 
     WSLCSession& m_session;
-    std::recursive_mutex m_lock;
+    std::mutex m_lock;
     std::atomic<size_t> m_callbackId{0};
 };
 } // namespace wsl::windows::service::wslc