diff --git a/include/trossen_sdk/runtime/session_manager.hpp b/include/trossen_sdk/runtime/session_manager.hpp index 3ebcc041..f6823c5c 100644 --- a/include/trossen_sdk/runtime/session_manager.hpp +++ b/include/trossen_sdk/runtime/session_manager.hpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -23,6 +24,7 @@ #include #include "trossen_sdk/hw/producer_base.hpp" +#include "trossen_sdk/hw/session_control/session_control_capable.hpp" #include "trossen_sdk/io/backends/lerobot_v2/lerobot_v2_backend.hpp" #include "trossen_sdk/io/backends/trossen_mcap/trossen_mcap_backend.hpp" #include "trossen_sdk/io/sink.hpp" @@ -197,6 +199,33 @@ class SessionManager { */ void request_rerecord(); + /** + * @brief Post a session-control event from any thread. + * + * Thread-safe. The event is queued and applied on the orchestrating thread + * by monitor_episode() / wait_for_reset(), so a hardware callback never + * touches the single-threaded session state directly. + */ + void post_event(hw::session_control::SessionControlEvent event); + + /** + * @brief Attach a push-based session-control source (VR buttons, keyboard, …). + * + * Installs callbacks that forward intents through post_event(), starts the + * source, and remembers it so shutdown() stops it (joining its thread) + * before the session is torn down. Every source drives the session through + * the same thread-safe queue, so keyboard and VR are interchangeable. + * + * @note The manager holds a shared_ptr to `source`, so the source is + * guaranteed to outlive the SessionManager's use of it (until shutdown() / + * destruction releases it). This removes the dangling-pointer risk a + * non-owning pointer would carry if the caller dropped the source early. + * @note Events are consumed by the synchronous monitor_episode() / + * wait_for_reset() loop, so drive the session with those rather than + * start_async_monitoring(). + */ + void attach_control(std::shared_ptr source); + /** * @brief Check if an episode is currently running * @@ -517,6 +546,34 @@ class SessionManager { /// @brief Flag indicating that re-record was requested std::atomic rerecord_requested_{false}; + /// @brief Phase in which queued control events are interpreted. + enum class ControlPhase { kRecording, kReset }; + + /// @brief Apply queued session-control events on the orchestrating thread, + /// translating each into the loop's control signals for the given phase. + void drain_control_events(ControlPhase phase); + + // Grants the control-integration unit tests access to drain_control_events() + // and the resulting control flags. + friend class SessionManagerControlTest; + + /// @brief Thread-safe queue of session-control events posted by sources. + std::deque control_events_; + + /// @brief Guards control_events_. + std::mutex control_events_mutex_; + + /// @brief Attached control sources; stopped in shutdown() before teardown. + /// Held by shared_ptr so a source cannot be destroyed while the manager may + /// still stop it. + std::vector> control_sources_; + + /// @brief Set by drain when a control event asks the current episode to stop. + std::atomic episode_stop_requested_{false}; + + /// @brief Set by drain when kStart during recording asks to skip the reset. + std::atomic skip_reset_{false}; + /// @brief Flag indicating that episode final statistics were emitted mutable std::atomic stats_emitted_{false}; diff --git a/src/runtime/session_manager.cpp b/src/runtime/session_manager.cpp index 5f10f116..8e44723e 100644 --- a/src/runtime/session_manager.cpp +++ b/src/runtime/session_manager.cpp @@ -93,6 +93,13 @@ SessionManager::~SessionManager() { } void SessionManager::shutdown() { + // Stop attached control sources first (joins their reader threads) so no + // queued event can drive the session while it is being torn down. + for (auto& source : control_sources_) { + if (source) source->stop(); + } + control_sources_.clear(); + // stop_episode() runs first so episode-ended callbacks fire against a still-non- // terminal SessionManager; shutdown_called_ latches afterwards and gates start_episode(). stop_episode(); @@ -891,6 +898,12 @@ UserAction SessionManager::wait_for_reset() { reset_signaled_.store(false); rerecord_requested_.store(false); + // A kStart received while recording asked to skip the reset wait entirely. + if (skip_reset_.exchange(false)) { + return trossen::utils::g_stop_requested ? UserAction::kStop + : UserAction::kContinue; + } + // Enable raw terminal input to detect keypresses without Enter trossen::utils::RawModeGuard raw_mode; @@ -898,12 +911,22 @@ UserAction SessionManager::wait_for_reset() { auto wait_or_signal = [this]() { std::unique_lock lk(reset_mutex_); reset_cv_.wait_for(lk, std::chrono::milliseconds(100), [this]() { - return reset_signaled_.load() || rerecord_requested_.load() || - trossen::utils::g_stop_requested; + if (reset_signaled_.load() || rerecord_requested_.load() || + trossen::utils::g_stop_requested) { + return true; + } + // Also wake when a control-source event is queued so poll_keys() drains + // it promptly instead of waiting out the interval. + std::lock_guard qlk(control_events_mutex_); + return !control_events_.empty(); }); }; auto poll_keys = [this]() -> UserAction { + // Apply any queued control-source events (VR buttons, etc.) on this thread. + drain_control_events(ControlPhase::kReset); + if (rerecord_requested_.load()) return UserAction::kReRecord; + auto key = trossen::utils::poll_keypress(); if (key == trossen::utils::KeyPress::kRightArrow) { reset_signaled_.store(true); @@ -968,12 +991,81 @@ void SessionManager::signal_reset_complete() { reset_cv_.notify_all(); } +void SessionManager::post_event(hw::session_control::SessionControlEvent event) { + using Event = hw::session_control::SessionControlEvent; + if (event == Event::kNone) return; + { + std::lock_guard lock(control_events_mutex_); + control_events_.push_back(event); + } + // Wake wait_for_reset() if it is blocked so the queued event is drained on + // the next poll pass instead of waiting out the ~100 ms interval. + reset_cv_.notify_all(); +} + +void SessionManager::attach_control( + std::shared_ptr source) +{ + if (!source) return; + using Event = hw::session_control::SessionControlEvent; + source->set_callbacks( + [this](Event e) { post_event(e); }, + [this] { + std::cerr << "\n[session] control source disconnected; ending session." + << std::endl; + post_event(Event::kStopSession); + }); + source->start(); + control_sources_.push_back(std::move(source)); +} + +void SessionManager::drain_control_events(ControlPhase phase) { + using Event = hw::session_control::SessionControlEvent; + + std::deque events; + { + std::lock_guard lock(control_events_mutex_); + events.swap(control_events_); + } + + for (const Event e : events) { + switch (e) { + case Event::kStart: + // "Advance": during recording, end the episode and skip the reset + // wait; during reset, resume immediately. + if (phase == ControlPhase::kRecording) { + skip_reset_.store(true); + episode_stop_requested_.store(true); + } else { + reset_signaled_.store(true); + } + break; + case Event::kStopEarly: + // Stop the current episode (a normal reset follows). Only meaningful + // while recording. + if (phase == ControlPhase::kRecording) { + episode_stop_requested_.store(true); + } + break; + case Event::kRerecord: + rerecord_requested_.store(true); + break; + case Event::kStopSession: + trossen::utils::g_stop_requested = true; + break; + case Event::kNone: + break; + } + } +} + UserAction SessionManager::monitor_episode( std::chrono::duration update_interval, std::chrono::duration sleep_interval, bool print_stats) { rerecord_requested_.store(false); + episode_stop_requested_.store(false); // Enable raw terminal input to detect keypresses during recording trossen::utils::RawModeGuard raw_mode; @@ -981,11 +1073,26 @@ UserAction SessionManager::monitor_episode( auto last_update = std::chrono::steady_clock::now(); while (!are_final_stats_emitted()) { + // Apply any queued control-source events (VR buttons, etc.) on this thread. + drain_control_events(ControlPhase::kRecording); + + // End-session takes priority over any episode-level action queued in the + // same drain. + if (trossen::utils::g_stop_requested) { + return UserAction::kStop; + } + // Check for re-record request if (rerecord_requested_.load()) { return UserAction::kReRecord; } + // A control source asked to end the episode (kStart advances, kStopEarly + // just stops); return so the main loop stops it on this thread. + if (episode_stop_requested_.load()) { + return UserAction::kContinue; + } + // Poll for arrow keys: left = re-record, right = early exit to next episode auto key = trossen::utils::poll_keypress(); if (key == trossen::utils::KeyPress::kLeftArrow) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e42e1330..7851915c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -65,6 +65,9 @@ target_link_libraries(test_session_manager PRIVATE trossen_sdk GTest::gtest_main add_executable(test_session_manager_discard test_session_manager_discard.cpp) target_link_libraries(test_session_manager_discard PRIVATE trossen_sdk GTest::gtest_main) +add_executable(test_session_manager_control test_session_manager_control.cpp) +target_link_libraries(test_session_manager_control PRIVATE trossen_sdk GTest::gtest_main) + add_executable(test_keyboard_input_utils test_keyboard_input_utils.cpp) target_link_libraries(test_keyboard_input_utils PRIVATE trossen_sdk GTest::gtest_main) @@ -119,6 +122,7 @@ gtest_discover_tests(test_sink) gtest_discover_tests(test_scheduler) gtest_discover_tests(test_session_manager) gtest_discover_tests(test_session_manager_discard) +gtest_discover_tests(test_session_manager_control) gtest_discover_tests(test_keyboard_input_utils) gtest_discover_tests(test_announce) gtest_discover_tests(test_teleop_space_helpers) diff --git a/tests/test_session_manager_control.cpp b/tests/test_session_manager_control.cpp new file mode 100644 index 00000000..e5e03074 --- /dev/null +++ b/tests/test_session_manager_control.cpp @@ -0,0 +1,178 @@ +/** + * @file test_session_manager_control.cpp + * @brief Unit tests for SessionManager's session-control integration. + * + * Covers the thread-safe event queue (post_event), the attach/detach lifecycle + * (attach_control + shutdown stops sources), and the phase-aware translation of + * queued events into the loop's control signals (drain_control_events). No + * hardware and no terminal are required. + */ + +#include +#include +#include +#include + +#include "gtest/gtest.h" +#include "nlohmann/json.hpp" + +#include "trossen_sdk/configuration/global_config.hpp" +#include "trossen_sdk/hw/session_control/session_control_capable.hpp" +#include "trossen_sdk/runtime/session_manager.hpp" +#include "trossen_sdk/utils/app_utils.hpp" + +namespace trossen::runtime { + +using hw::session_control::SessionControlCapable; +using hw::session_control::SessionControlEvent; + +// A source with no real input; tests drive it by calling the protected base +// helpers directly and observe start()/stop() lifecycle. +class FakeControlSource : public SessionControlCapable { + public: + void start() override { started = true; } + void stop() override { stopped = true; } + + void fireEvent(SessionControlEvent e) { emit_event(e); } + void fireDisconnect() { signal_disconnect(); } + void reconnect() { arm_disconnect(); } + + bool started = false; + bool stopped = false; +}; + +// Friend fixture: reaches into SessionManager's private drain + control flags. +class SessionManagerControlTest : public ::testing::Test { + protected: + static void SetUpTestSuite() { + nlohmann::json config = { + {"session_manager", { + {"type", "session_manager"}, + {"max_episodes", 100}, + {"backend_type", "null"} + }} + }; + trossen::configuration::GlobalConfig::instance().load_from_json(config); + } + + void SetUp() override { trossen::utils::g_stop_requested = false; } + void TearDown() override { trossen::utils::g_stop_requested = false; } + + void drainRecording() { + sm_.drain_control_events(SessionManager::ControlPhase::kRecording); + } + void drainReset() { + sm_.drain_control_events(SessionManager::ControlPhase::kReset); + } + + bool episodeStop() const { return sm_.episode_stop_requested_.load(); } + bool skipReset() const { return sm_.skip_reset_.load(); } + bool resetSignaled() const { return sm_.reset_signaled_.load(); } + bool rerecord() const { return sm_.rerecord_requested_.load(); } + std::size_t pendingCount() { + std::lock_guard lock(sm_.control_events_mutex_); + return sm_.control_events_.size(); + } + + SessionManager sm_; +}; + +// ── attach / shutdown lifecycle ───────────────────────────────────────────── + +TEST_F(SessionManagerControlTest, AttachStartsSourceAndShutdownStopsIt) { + auto fake = std::make_shared(); + sm_.attach_control(fake); + EXPECT_TRUE(fake->started); + + sm_.shutdown(); + EXPECT_TRUE(fake->stopped); +} + +TEST_F(SessionManagerControlTest, SourceEventForwardsThroughAttach) { + auto fake = std::make_shared(); + sm_.attach_control(fake); + + fake->fireEvent(SessionControlEvent::kRerecord); + drainRecording(); + EXPECT_TRUE(rerecord()); + + sm_.shutdown(); +} + +// ── post_event queueing ───────────────────────────────────────────────────── + +TEST_F(SessionManagerControlTest, PostEventQueuesAndDrainClears) { + sm_.post_event(SessionControlEvent::kStart); + sm_.post_event(SessionControlEvent::kStopEarly); + EXPECT_EQ(pendingCount(), 2u); + + drainRecording(); + EXPECT_EQ(pendingCount(), 0u); +} + +TEST_F(SessionManagerControlTest, PostEventIgnoresNone) { + sm_.post_event(SessionControlEvent::kNone); + EXPECT_EQ(pendingCount(), 0u); +} + +// ── phase-aware translation ───────────────────────────────────────────────── + +TEST_F(SessionManagerControlTest, StartWhileRecordingStopsAndSkipsReset) { + sm_.post_event(SessionControlEvent::kStart); + drainRecording(); + EXPECT_TRUE(episodeStop()); + EXPECT_TRUE(skipReset()); +} + +TEST_F(SessionManagerControlTest, StartWhileResettingSignalsResume) { + sm_.post_event(SessionControlEvent::kStart); + drainReset(); + EXPECT_TRUE(resetSignaled()); + EXPECT_FALSE(skipReset()); +} + +TEST_F(SessionManagerControlTest, StopEarlyWhileRecordingStopsWithoutSkip) { + sm_.post_event(SessionControlEvent::kStopEarly); + drainRecording(); + EXPECT_TRUE(episodeStop()); + EXPECT_FALSE(skipReset()); +} + +TEST_F(SessionManagerControlTest, StopEarlyWhileResettingIsIgnored) { + sm_.post_event(SessionControlEvent::kStopEarly); + drainReset(); + EXPECT_FALSE(episodeStop()); + EXPECT_FALSE(resetSignaled()); +} + +TEST_F(SessionManagerControlTest, RerecordSetsFlagInBothPhases) { + sm_.post_event(SessionControlEvent::kRerecord); + drainRecording(); + EXPECT_TRUE(rerecord()); +} + +TEST_F(SessionManagerControlTest, StopSessionSetsGlobalStop) { + sm_.post_event(SessionControlEvent::kStopSession); + drainReset(); + EXPECT_TRUE(trossen::utils::g_stop_requested.load()); +} + +// ── disconnect composes with the base fire-once guarantee ─────────────────── + +TEST_F(SessionManagerControlTest, DisconnectPostsStopSessionExactlyOnce) { + auto fake = std::make_shared(); + sm_.attach_control(fake); + + // The base guarantees a single disconnect until re-armed, so two drops + // queue only one kStopSession. + fake->fireDisconnect(); + fake->fireDisconnect(); + EXPECT_EQ(pendingCount(), 1u); + + drainReset(); + EXPECT_TRUE(trossen::utils::g_stop_requested.load()); + + sm_.shutdown(); +} + +} // namespace trossen::runtime