Bound the publish_live_objects graceful flush on stop

Treat a flush write failure or in-loop send failure that races with a
concurrent close() as a benign stop (return success), and time-box the
post-loop flush as secondary insurance against a slow multi-write flush.
Adds flush-fallback, graceful-flush, and no-false-positive tests; the mock
transport now fails writes once closed, mirroring picoquic teardown.

Author: mondain

src/transport/moqt_session.cpp

@@ -4134,6 +4134,11 @@ TransportStatus MoqtSession::publish_live_objects(const openmoq::publisher::Live
                 next->final_in_subgroup,
                 std::span<const std::uint8_t>(next->payload));
             if (!status.ok) {
+                // A send failure that races with a concurrent close() is a
+                // benign stop, not a transport error: break to the flush.
+                if (stop_requested_.load(std::memory_order_acquire)) {
+                    break;
+                }
                 return status;
             }
             record_published_object(next->track_name,
@@ -4181,6 +4186,11 @@ TransportStatus MoqtSession::publish_live_objects(const openmoq::publisher::Live
             next->final_in_subgroup,
             std::span<const std::uint8_t>(next->payload));
         if (!status.ok) {
+            // A send failure that races with a concurrent close() is a benign
+            // stop, not a transport error: break to the flush.
+            if (stop_requested_.load(std::memory_order_acquire)) {
+                break;
+            }
             return status;
         }
         record_published_object(next->track_name,
@@ -4191,15 +4201,35 @@ TransportStatus MoqtSession::publish_live_objects(const openmoq::publisher::Live
         }
     }
 
+    // Bound the post-loop graceful flush. When a stop is in progress, a flush
+    // write can fail (transport torn down by a concurrent close()) or the total
+    // flush can run long; in either case return success rather than surfacing a
+    // stop-driven flush failure. The deadline is secondary insurance for a
+    // multi-write flush that is merely slow.
+    const auto flush_deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
+    const auto flush_budget_exhausted = [&]() {
+        return stop_requested_.load(std::memory_order_acquire) &&
+               std::chrono::steady_clock::now() >= flush_deadline;
+    };
+    const auto stopping = [&]() {
+        return stop_requested_.load(std::memory_order_acquire);
+    };
+
     for (auto& [track_name, sender] : sender_by_track) {
         static_cast<void>(track_name);
+        if (flush_budget_exhausted()) {
+            return TransportStatus::success();
+        }
         status = sender.finish_group(transport_);
         if (!status.ok) {
-            return status;
+            return stopping() ? TransportStatus::success() : status;
         }
     }
 
     for (const auto& [request_id, subscribe] : active_subscriptions) {
+        if (flush_budget_exhausted()) {
+            return TransportStatus::success();
+        }
         const auto stream_it = active_subscription_stream_ids.find(request_id);
         const std::uint64_t response_stream_id =
             uses_request_streams(draft_version) && stream_it != active_subscription_stream_ids.end()
@@ -4210,29 +4240,39 @@ TransportStatus MoqtSession::publish_live_objects(const openmoq::publisher::Live
                                              draft_version, request_id, sender_by_track[subscribe.track_name].stream_count()),
                                          false);
         if (!status.ok) {
-            return status;
+            return stopping() ? TransportStatus::success() : status;
         }
     }
     for (const auto& [track_name, request_id] : subscribe_tracks_publish_request_ids) {
         if (!subscribed_tracks.contains(track_name)) {
             continue;
         }
+        if (flush_budget_exhausted()) {
+            return TransportStatus::success();
+        }
         status = write_publish_done_for_request(transport_,
                                                 draft_version,
                                                 control_stream_id_,
                                                 publish_stream_id_by_request_id_,
                                                 request_id,
                                                 sender_by_track[track_name].stream_count());
         if (!status.ok) {
-            return status;
+            return stopping() ? TransportStatus::success() : status;
         }
     }
 
-    return write_namespace_done_for_request(transport_,
+    if (flush_budget_exhausted()) {
+        return TransportStatus::success();
+    }
+    const auto namespace_done_status = write_namespace_done_for_request(transport_,
                                             draft_version,
                                             control_stream_id_,
                                             namespace_stream_id_,
                                             namespace_message);
+    if (!namespace_done_status.ok && stopping()) {
+        return TransportStatus::success();
+    }
+    return namespace_done_status;
 }
 
 TransportStatus MoqtSession::close(std::uint64_t application_error_code) {

tests/moqt_session_test.cpp

@@ -180,6 +180,16 @@ struct MockTransport final : PublisherTransport {
     TransportStatus write_stream(std::uint64_t stream_id,
                                  std::span<const std::uint8_t> bytes,
                                  bool fin) override {
+        // A blocked write is released either explicitly or by transport close,
+        // mirroring how a real picoquic write unblocks when the connection is
+        // torn down.
+        while (block_writes.load() && !release_writes.load() &&
+               state_ != ConnectionState::kClosed) {
+            std::this_thread::sleep_for(std::chrono::milliseconds(5));
+        }
+        if (state_ == ConnectionState::kClosed) {
+            return TransportStatus::failure("transport closed");
+        }
         writes.push_back({
             .stream_id = stream_id,
             .bytes = std::vector<std::uint8_t>(bytes.begin(), bytes.end()),
@@ -250,6 +260,8 @@ struct MockTransport final : PublisherTransport {
     std::map<std::uint64_t, std::vector<std::vector<std::uint8_t>>> reads;
     std::set<std::uint64_t> accepted_streams;
     std::function<void(MockTransport&, std::uint64_t)> on_read;
+    std::atomic<bool> block_writes{false};
+    std::atomic<bool> release_writes{false};
 };
 
 std::vector<std::size_t> object_write_indices(const MockTransport& transport) {
@@ -2712,5 +2724,126 @@ int main() {
         }
     }
 
+    {
+        // Flush fallback: if post-stop flush writes block, publish_live_objects
+        // must still return promptly. close() tears down the transport, which
+        // makes the blocked write fail; the stop-swallow path returns success.
+        MockTransport transport;
+        transport.reads[0].push_back(encode_server_setup_message({
+            .draft = DraftVersion::kDraft14,
+            .max_request_id = 8,
+        }));
+        transport.reads[0].push_back(encode_publish_namespace_ok_message(DraftVersion::kDraft14, 0));
+        MoqtSession session(transport, std::string(kTestTrackNamespace),
+                            /*auto_forward=*/true, /*publish_catalog=*/false,
+                            /*paced=*/false, std::chrono::seconds(1));
+        ok &= expect(session.connect(endpoint, tls).ok,
+                     "expected flush-fallback session connect to succeed");
+
+        LiveObjectSource source;
+        source.tracks = {LiveTrack{.track_name = "events"}};
+        std::atomic<int> object_calls{0};
+        source.next_object = [&object_calls]() -> std::optional<LiveObject> {
+            const int n = object_calls.fetch_add(1);
+            return LiveObject{.track_name = "events", .group_id = 1,
+                              .object_id = static_cast<std::size_t>(n),
+                              .payload = {'O', 'K'}};
+        };
+
+        std::promise<TransportStatus> result_promise;
+        auto result_future = result_promise.get_future();
+        std::thread worker([&]() {
+            result_promise.set_value(session.publish_live_objects(source, DraftVersion::kDraft14));
+        });
+        std::this_thread::sleep_for(std::chrono::milliseconds(50));
+        transport.block_writes.store(true);  // any flush write now blocks
+        session.close(0);                     // tears down transport -> unblocks via failure
+
+        const auto wait_status = result_future.wait_for(std::chrono::seconds(4));
+        ok &= expect(wait_status == std::future_status::ready,
+                     "expected publish_live_objects to return promptly despite blocked flush");
+        transport.release_writes.store(true);  // belt-and-suspenders unblock
+        if (wait_status == std::future_status::ready) {
+            worker.join();
+            ok &= expect(result_future.get().ok,
+                         "expected stop with blocked flush to still return success");
+        } else {
+            worker.detach();
+        }
+    }
+
+    {
+        // Graceful flush: with a healthy transport, a stop still returns success
+        // and the loop served at least one object.
+        MockTransport transport;
+        transport.reads[0].push_back(encode_server_setup_message({
+            .draft = DraftVersion::kDraft14,
+            .max_request_id = 8,
+        }));
+        transport.reads[0].push_back(encode_publish_namespace_ok_message(DraftVersion::kDraft14, 0));
+        MoqtSession session(transport, std::string(kTestTrackNamespace),
+                            /*auto_forward=*/true, /*publish_catalog=*/false,
+                            /*paced=*/false, std::chrono::seconds(1));
+        ok &= expect(session.connect(endpoint, tls).ok,
+                     "expected graceful-flush session connect to succeed");
+
+        LiveObjectSource source;
+        source.tracks = {LiveTrack{.track_name = "events"}};
+        std::atomic<int> object_calls{0};
+        source.next_object = [&object_calls]() -> std::optional<LiveObject> {
+            const int n = object_calls.fetch_add(1);
+            return LiveObject{.track_name = "events", .group_id = 1,
+                              .object_id = static_cast<std::size_t>(n),
+                              .payload = {'O', 'K'}};
+        };
+
+        std::promise<TransportStatus> result_promise;
+        auto result_future = result_promise.get_future();
+        std::thread worker([&]() {
+            result_promise.set_value(session.publish_live_objects(source, DraftVersion::kDraft14));
+        });
+        std::this_thread::sleep_for(std::chrono::milliseconds(50));
+        session.close(0);
+        const auto wait_status = result_future.wait_for(std::chrono::seconds(3));
+        ok &= expect(wait_status == std::future_status::ready,
+                     "expected graceful-flush publish to return after close()");
+        if (wait_status == std::future_status::ready) {
+            worker.join();
+            ok &= expect(result_future.get().ok, "expected graceful-flush return success");
+            ok &= expect(object_calls.load() > 0,
+                         "expected graceful flush after serving objects");
+        } else {
+            worker.detach();
+        }
+    }
+
+    {
+        // No false positive: a transport error WITHOUT a stop request must still
+        // return failure. next_object returns an object for an unknown track,
+        // which the loop rejects with failure at runtime; close() is never called.
+        MockTransport transport;
+        transport.reads[0].push_back(encode_server_setup_message({
+            .draft = DraftVersion::kDraft14,
+            .max_request_id = 8,
+        }));
+        transport.reads[0].push_back(encode_publish_namespace_ok_message(DraftVersion::kDraft14, 0));
+        MoqtSession session(transport, std::string(kTestTrackNamespace),
+                            /*auto_forward=*/true, /*publish_catalog=*/false,
+                            /*paced=*/false, std::chrono::seconds(1));
+        ok &= expect(session.connect(endpoint, tls).ok,
+                     "expected no-false-positive session connect to succeed");
+
+        LiveObjectSource source;
+        source.tracks = {LiveTrack{.track_name = "events"}};
+        source.next_object = []() -> std::optional<LiveObject> {
+            return LiveObject{.track_name = "not_a_declared_track", .group_id = 1,
+                              .object_id = 0, .payload = {'X'}};
+        };
+
+        const TransportStatus result = session.publish_live_objects(source, DraftVersion::kDraft14);
+        ok &= expect(!result.ok,
+                     "expected an unknown-track error without stop to return failure");
+    }
+
     return ok ? 0 : 1;
 }