From 3babdd7c37433c821f1d126e91a77e69a3b944e5 Mon Sep 17 00:00:00 2001 From: denys-gif Date: Mon, 12 Oct 2026 19:04:29 +0100 Subject: [PATCH 1/5] feat: client feature flags --- .../docs/NATS_INTEGRATION_TESTS.md | 212 +++++++++ clients/openframe-client/src/cli.rs | 7 + clients/openframe-client/src/cli/restart.rs | 22 + .../src/config/update_config.rs | 5 + clients/openframe-client/src/lib.rs | 62 ++- .../src/listener/command/core_nats.rs | 130 ++++++ .../src/listener/command/delivery.rs | 126 ++++++ .../src/listener/command/handlers.rs | 65 +++ .../src/listener/command/handlers_tests.rs | 119 +++++ .../src/listener/command/integration_tests.rs | 425 ++++++++++++++++++ .../src/listener/command/mod.rs | 127 ++++++ .../src/listener/command/restart.rs | 73 +++ .../src/listener/command/restart_tests.rs | 118 +++++ .../src/listener/execution_listener.rs | 5 +- clients/openframe-client/src/listener/mod.rs | 2 + .../openframe-client/src/models/delivery.rs | 34 ++ .../openframe-client/src/models/execution.rs | 4 +- .../src/models/feature_flags.rs | 77 ++++ clients/openframe-client/src/models/mod.rs | 9 +- .../src/platform/detached_process.rs | 68 +++ clients/openframe-client/src/platform/mod.rs | 1 + .../src/platform/uninstall.rs | 63 +-- .../src/services/feature_flags_service.rs | 126 ++++++ .../services/feature_flags_service_tests.rs | 236 ++++++++++ clients/openframe-client/src/services/mod.rs | 1 + .../src/services/nats_connection_manager.rs | 5 + .../src/services/result_outbox_run_manager.rs | 4 +- clients/openframe-client/src/test_support.rs | 68 +++ 28 files changed, 2125 insertions(+), 69 deletions(-) create mode 100644 clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md create mode 100644 clients/openframe-client/src/cli/restart.rs create mode 100644 clients/openframe-client/src/listener/command/core_nats.rs create mode 100644 clients/openframe-client/src/listener/command/delivery.rs create mode 100644 clients/openframe-client/src/listener/command/handlers.rs create mode 100644 clients/openframe-client/src/listener/command/handlers_tests.rs create mode 100644 clients/openframe-client/src/listener/command/integration_tests.rs create mode 100644 clients/openframe-client/src/listener/command/mod.rs create mode 100644 clients/openframe-client/src/listener/command/restart.rs create mode 100644 clients/openframe-client/src/listener/command/restart_tests.rs create mode 100644 clients/openframe-client/src/models/delivery.rs create mode 100644 clients/openframe-client/src/models/feature_flags.rs create mode 100644 clients/openframe-client/src/platform/detached_process.rs create mode 100644 clients/openframe-client/src/services/feature_flags_service.rs create mode 100644 clients/openframe-client/src/services/feature_flags_service_tests.rs create mode 100644 clients/openframe-client/src/test_support.rs diff --git a/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md b/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md new file mode 100644 index 0000000000..702e78e778 --- /dev/null +++ b/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md @@ -0,0 +1,212 @@ +# NATS integration tests + +Tests that run real client code against a real NATS server on your machine — +subjects, JSON payloads, JetStream publishes and acks all go over the wire. +No mocks of NATS. + +## How it works + +``` +test + ├─ NatsServer::start() spawns `nats-server -a 127.0.0.1 -p -js -sd ` + │ waits until the port accepts connections, kills it on drop + ├─ builds the agent side real services on a temp DirectoryManager, + │ NatsConnectionManager given a plain nats:// client + ├─ "tester" async-nats client plays the backend: publishes commands, subscribes to results + └─ asserts on what comes back ACKED / DONE / FAILED, files on disk, signals +``` + +- Every test starts its **own** server on a free port with its own temp storage, + so tests run in parallel and leave nothing behind. +- JetStream is on (`-js`), so code that publishes to a stream (`publish_acked`) + works. The test creates the streams it needs. +- The agent connects without TLS and without a token: the test hands + `NatsConnectionManager` a ready client instead of calling `connect()`. + +### Pieces + +| What | Where | Use | +|---|---|---| +| `NatsServer` | `src/test_support.rs` (`cfg(test)`) | `NatsServer::start() -> Option`, `.url()` | +| `use_client_for_tests` | `src/services/nats_connection_manager.rs` (`cfg(test)`) | puts a connected `async_nats::Client` into the manager | +| Example suite | `src/listener/command/integration_tests.rs` | the feature-flags listener end to end | + +## Running + +Install the server once (single ~1 MB binary): + +```bash +brew install nats-server +``` + +On Windows or Linux, download a release from + and put it on `PATH`. + +Then the usual: + +```bash +make test +``` + +runs them together with the unit tests. + +Only the integration suite, failing if the server is missing: + +```bash +REQUIRE_NATS_SERVER=1 OPENFRAME_VERSION=0.0.0-dev cargo test --lib -- integration_tests +``` + +| Variable | Effect | +|---|---| +| *(none)* | no `nats-server` on `PATH` → each test returns early and **passes without running** | +| `REQUIRE_NATS_SERVER=1` | a missing binary is a test failure — use in CI so the suite can't silently skip | +| `NATS_SERVER_BIN=/path/to/nats-server` | use a binary that isn't on `PATH` | + +A skipped test prints ``skipping: `nats-server` is not installed (brew install nats-server)`` to stderr +(visible with `-- --nocapture`). + +### CI + +CI does not install `nats-server` yet, so there the suite skips and nothing +fails. To make it run: add an install step before `make test` in the Rust job +(`brew install nats-server` on macOS, a release download or `choco install nats-server` +on Windows) and set `REQUIRE_NATS_SERVER: "1"` on that step. + +## Writing a test + +### 1. Start a server and skip if there is none + +```rust +use crate::test_support::NatsServer; + +#[tokio::test] +async fn my_command_is_acked_and_done() { + let Some(server) = NatsServer::start() else { + return; + }; + // ... +} +``` + +Always use the `let … else { return }` form: that is what makes the test skip +on machines without the binary. + +### 2. Build the agent side with real services + +Everything is built on a temp directory, the same way `Client::new` does it: + +```rust +let dir = tempfile::tempdir().unwrap(); +let directories = DirectoryManager::with_custom_dirs( + dir.path().join("logs"), + dir.path().join("app"), + dir.path().join("secured"), +); +let config = AgentConfigurationService::new(directories.clone()).unwrap(); +config + .save_registration_data("it-machine".into(), "client".into(), "secret".into()) + .await + .unwrap(); + +let manager = NatsConnectionManager::new(/* config, initial config, auth, tls, ... */); +manager.use_client_for_tests(async_nats::connect(server.url()).await.unwrap()); +``` + +The machine id saved in the config is what listeners put into their subjects +(`machine..…`). See `Agent::start()` in `listener/command/integration_tests.rs` +for the full list of services and copy it. + +### 3. Create the streams your code publishes to + +`nats-server` starts empty. If the code under test publishes with an ack, create +the stream first, or the publish fails with "no responders": + +```rust +let tester = async_nats::connect(server.url()).await.unwrap(); +async_nats::jetstream::new(tester.clone()) + .create_stream(async_nats::jetstream::stream::Config { + name: "DELIVERY_RESULT".into(), + subjects: vec!["machine.*.delivery.result".into()], + ..Default::default() + }) + .await + .unwrap(); +``` + +### 4. Subscribe to what you assert on before sending + +```rust +let mut results = tester + .subscribe("machine.it-machine.delivery.result") + .await + .unwrap(); +tester.flush().await.unwrap(); +``` + +Core NATS has no backlog: anything published before the subscription exists is +lost. The same applies to the agent's listener — see the next step. + +### 5. Wait until the listener is subscribed + +A listener subscribes asynchronously after `start()`. Don't `sleep`; send a probe +command and repeat until it is answered: + +```rust +timeout(BOUND, async { + loop { + send(probe()).await; + if timeout(Duration::from_millis(200), next_result("probe")).await.is_ok() { + return; + } + } +}) +.await +.expect("listener never answered"); +``` + +Pick a probe that has no side effects (the flags suite uses a stale `sequence`). +Later assertions must ignore extra probe answers — filter results by `dispatchId`. + +### 6. Assert with bounded waits + +- **Something must happen:** wait for it with `timeout(BOUND, …)` (10 s in the + example suite) and fail with a message that says what was expected. +- **Something must not happen:** there is no event to wait for, so wait a short + fixed window (`QUIET`, 500 ms) and assert nothing arrived. Keep these few. +- **Results arrive out of order or interleaved:** don't drop messages you don't + want yet — buffer them (`Agent::next_result` keeps an `unclaimed` list). + Dropping them is the classic flaky-test bug: a later `expect` waits for a + message the helper already threw away. + +### 7. Assert on real side effects + +Read the files the code wrote (`feature_flags.json`), wait on the signals it +raises (`CommandContext::restart_requested()`), check what reached the "backend" +subscription. Prefer these over inspecting internal state. + +## Reusing the example harness + +`listener/command/integration_tests.rs` wraps all of the above in a small +`Agent` struct: + +| Helper | Does | +|---|---| +| `Agent::start()` | server + agent + streams + result subscription + listener up, or `None` to skip | +| `agent.launch()` | starts another listener on the same agent state — "the next process" after a restart | +| `agent.send(payload)` | publishes a command to the listener's subject | +| `agent.expect(dispatch_id, result)` | waits for `ACKED` / `DONE` / `FAILED` for that dispatch | +| `agent.expect_silence(dispatch_id)` | asserts nothing arrives for that dispatch within `QUIET` | +| `agent.persisted()` | reads the state file the listener wrote | +| `command(dispatch_id, sequence, body)` | builds a payload with a `delivery` block | + +For a new delivery command, copy the struct, swap the handler in `launch()` +and the subject in `send()`, and keep the helpers. + +## Limits + +- What the server enforces in production but the local server doesn't: auth, + TLS, and the per-user subject permissions. A `Permissions Violation` on a real + environment is not caught here — check the NATS config when you add a subject. +- Process-level effects (the service manager restarting the client, `exit`) are + out of reach; tests stop at the signal the code raises. +- Each test starts its own server; the 7 tests of the example suite take about a second in total. Fine for dozens of tests, not thousands. diff --git a/clients/openframe-client/src/cli.rs b/clients/openframe-client/src/cli.rs index 3666a4d836..aa0e92a6ba 100644 --- a/clients/openframe-client/src/cli.rs +++ b/clients/openframe-client/src/cli.rs @@ -3,6 +3,7 @@ mod auth; mod healthcheck; mod install; mod permissions; +mod restart; mod uninstall; use crate::doctor::healing::{HealOutcome, HealResult}; @@ -14,6 +15,8 @@ use std::process; use tokio::runtime::Runtime; +pub const RESTART_SERVICE_SUBCOMMAND: &str = "restart-service"; + #[derive(Parser)] #[command(name = "openframe-client", author, version = env!("OPENFRAME_VERSION"), about, long_about = None)] struct Cli { @@ -74,6 +77,9 @@ enum Commands { CheckPermissions, /// Run environment health check (reads config from installed agent) Doctor, + /// Stop and start the client service (run detached by the client to apply feature flags) + #[command(hide = true, name = RESTART_SERVICE_SUBCOMMAND)] + RestartService, } /// Parse CLI arguments and run the requested OpenFrame client command. @@ -91,6 +97,7 @@ pub fn run() -> Result<()> { Some(Commands::Run) => agent::run_direct(&rt), Some(Commands::RunAsService) => agent::run_as_service(), Some(Commands::CheckPermissions) => permissions::run(), + Some(Commands::RestartService) => restart::run(&rt), None => agent::run_legacy(&rt), } diff --git a/clients/openframe-client/src/cli/restart.rs b/clients/openframe-client/src/cli/restart.rs new file mode 100644 index 0000000000..f6e7f93843 --- /dev/null +++ b/clients/openframe-client/src/cli/restart.rs @@ -0,0 +1,22 @@ +use std::process; + +use tokio::runtime::Runtime; +use tracing::{error, info}; + +use crate::service::Service; + +pub(super) fn run(rt: &Runtime) -> ! { + super::init_file_logging(); + info!("Restarting the OpenFrame client service"); + + match rt.block_on(Service::nudge_restart()) { + Ok(()) => { + info!("OpenFrame client service restarted"); + process::exit(0); + } + Err(e) => { + error!("Failed to restart the OpenFrame client service: {:#}", e); + process::exit(1); + } + } +} diff --git a/clients/openframe-client/src/config/update_config.rs b/clients/openframe-client/src/config/update_config.rs index 1932368a0e..2b92afaaa1 100644 --- a/clients/openframe-client/src/config/update_config.rs +++ b/clients/openframe-client/src/config/update_config.rs @@ -17,6 +17,11 @@ pub const CONSUMER_CYCLE_PAUSE_MS: u64 = 30000; // 30 seconds pause between retr // Reconnection pub const RECONNECTION_DELAY_MS: u64 = 5000; // 5 seconds +pub const SELF_RESTART_EXIT_CODE: i32 = 75; +pub const SELF_RESTART_IDLE_WAIT_SECS: u64 = 600; +pub const SELF_RESTART_IDLE_POLL_SECS: u64 = 5; +pub const SELF_RESTART_HANDOFF_SECS: u64 = 120; + // Durable result outbox pub const OUTBOX_MAX_ENTRIES: usize = 500; pub const OUTBOX_FLUSH_INTERVAL_SECS: u64 = 30; diff --git a/clients/openframe-client/src/lib.rs b/clients/openframe-client/src/lib.rs index 152c07126e..2ea4fc54f5 100644 --- a/clients/openframe-client/src/lib.rs +++ b/clients/openframe-client/src/lib.rs @@ -3,7 +3,7 @@ use serde::{Deserialize, Serialize}; use std::sync::Arc; use tokio::sync::RwLock; use tokio::time::Duration; -use tracing::{error, info}; +use tracing::{error, info, warn}; pub mod clients; mod config; @@ -29,6 +29,9 @@ pub mod service_adapter; pub mod system; pub mod utils; +#[cfg(test)] +mod test_support; + pub mod cli; pub use cli::run; @@ -38,8 +41,12 @@ pub mod executor; use crate::clients::build_agent_http_client; use crate::clients::tool_agent_file_client::ToolAgentFileClient; use crate::clients::{AuthClient, RegistrationClient, ToolApiClient}; -use crate::config::update_config::{DOWNLOAD_CLIENT_TIMEOUT_SECS, HTTP_CLIENT_TIMEOUT_SECS}; +use crate::config::update_config::{ + DOWNLOAD_CLIENT_TIMEOUT_SECS, HTTP_CLIENT_TIMEOUT_SECS, SELF_RESTART_EXIT_CODE, + SELF_RESTART_HANDOFF_SECS, +}; use crate::listener::client_uninstall_message_listener::ClientUninstallMessageListener; +use crate::listener::command::{CommandContext, CommandListener, FeatureFlagsHandler}; use crate::listener::execution_listener::ExecutionListener; use crate::listener::openframe_client_update_listener::OpenFrameClientUpdateListener; use crate::listener::remote_access_message_listener::RemoteAccessMessageListener; @@ -59,6 +66,7 @@ use crate::services::deactivation_service::DeactivationService; use crate::services::device_data_fetcher::DeviceDataFetcher; use crate::services::encryption_service::EncryptionService; use crate::services::execution_service::ExecutionService; +use crate::services::feature_flags_service::FeatureFlagsService; use crate::services::github_download_service::GithubDownloadService; use crate::services::hostname_report_publisher::HostnameReportPublisher; use crate::services::initial_authentication_processor::InitialAuthenticationProcessor; @@ -172,6 +180,8 @@ pub struct Client { tool_agent_update_listener: ToolAgentUpdateListener, client_uninstall_message_listener: ClientUninstallMessageListener, remote_access_message_listener: RemoteAccessMessageListener, + feature_flags_listener: CommandListener, + command_context: CommandContext, command_execution_listener: ExecutionListener, script_execution_listener: ExecutionListener, script_bootstrap_execution_listener: ExecutionListener, @@ -366,6 +376,10 @@ impl Client { let tool_ops = ToolOps::default(); + let feature_flags_service = Arc::new(FeatureFlagsService::load( + directory_manager.secured_dir().join("feature_flags.json"), + )); + // Initialize tool run manager let tool_run_manager = ToolRunManager::new( installed_tools_service.clone(), @@ -553,6 +567,20 @@ impl Client { ); let result_store_for_recovery = result_store.clone(); + let command_context = CommandContext::new( + nats_connection_manager.clone(), + config_service.clone(), + tool_ops.clone(), + tool_run_manager.clone(), + nats_message_publisher.clone(), + result_store.clone(), + flush_notify.clone(), + ); + let feature_flags_listener = CommandListener::new( + FeatureFlagsHandler::new(feature_flags_service), + command_context.clone(), + ); + let command_execution_listener = ExecutionListener::::new( nats_connection_manager.clone(), nats_message_publisher.clone(), @@ -635,6 +663,8 @@ impl Client { tool_agent_update_listener, client_uninstall_message_listener, remote_access_message_listener, + feature_flags_listener, + command_context, command_execution_listener, script_execution_listener, script_bootstrap_execution_listener, @@ -713,6 +743,34 @@ impl Client { self.nats_connection_manager.start_connection_watchdog(); + self.feature_flags_listener.start().await?; + let command_context = self.command_context.clone(); + tokio::spawn(async move { + command_context.restart_requested().await; + match crate::platform::detached_process::spawn_detached( + &crate::service::Service::get_install_location(), + cli::RESTART_SERVICE_SUBCOMMAND, + ) { + Ok(pid) => { + info!(pid, "Service restart handed to a detached process"); + tokio::time::sleep(Duration::from_secs(SELF_RESTART_HANDOFF_SECS)).await; + warn!( + "Service was not restarted within {}s, exiting instead", + SELF_RESTART_HANDOFF_SECS + ); + } + Err(e) => warn!( + "Cannot hand the service restart off, exiting instead: {:#}", + e + ), + } + info!( + code = SELF_RESTART_EXIT_CODE, + "Exiting so the service manager restarts the client" + ); + std::process::exit(SELF_RESTART_EXIT_CODE); + }); + // Handle any pending update from previous run (after NATS is connected) if let Err(e) = self.update_handler_service.handle_pending_update().await { error!("Failed to handle pending update: {:#}", e); diff --git a/clients/openframe-client/src/listener/command/core_nats.rs b/clients/openframe-client/src/listener/command/core_nats.rs new file mode 100644 index 0000000000..93639cf605 --- /dev/null +++ b/clients/openframe-client/src/listener/command/core_nats.rs @@ -0,0 +1,130 @@ +use anyhow::{anyhow, bail, Result}; +use async_nats::Message; +use futures::StreamExt; +use serde::Deserialize; +use tracing::{error, info, warn}; + +use super::{CommandHandler, CommandListener, Outcome}; +use crate::models::DeliveryRef; + +#[derive(Deserialize)] +struct Envelope { + delivery: DeliveryRef, +} + +pub(super) async fn report_restart_completion(listener: &CommandListener) { + let Some(delivery) = listener.handler.awaiting_restart() else { + return; + }; + let machine_id = match listener.context.config_service.get_machine_id() { + Ok(machine_id) => machine_id, + Err(e) => { + error!(kind = H::KIND, dispatch_id = %delivery.dispatch_id, "Cannot resolve machine id to report the restart: {:#}", e); + return; + } + }; + info!(kind = H::KIND, dispatch_id = %delivery.dispatch_id, "Applied after restart, reporting done"); + if !listener + .context + .delivery + .settle(&machine_id, delivery, Outcome::Done) + .await + { + return; + } + if let Err(e) = listener.handler.restart_reported() { + warn!( + kind = H::KIND, + "Failed to clear the restart marker, done is re-sent on next start: {:#}", e + ); + } +} + +pub(super) async fn listen(listener: &CommandListener) -> Result<()> { + let nats = &listener.context.nats_connection_manager; + let mut client_rx = nats.on_client_replaced(); + let client = nats.get_client().await?; + let machine_id = listener.context.config_service.get_machine_id()?; + + let subject = H::subject(&machine_id); + let mut subscriber = client + .subscribe(subject.clone()) + .await + .map_err(|e| anyhow!("failed to subscribe to {}: {}", subject, e))?; + info!(kind = H::KIND, subject = %subject, "Command listener active"); + + loop { + tokio::select! { + message = subscriber.next() => match message { + Some(message) => handle(listener, message, &machine_id).await, + None => bail!("subscription to {} ended", subject), + }, + _ = client_rx.changed() => { + info!(kind = H::KIND, "NATS client replaced, resubscribing"); + return Ok(()); + } + } + } +} + +async fn handle( + listener: &CommandListener, + message: Message, + machine_id: &str, +) { + let payload = String::from_utf8_lossy(&message.payload); + info!(kind = H::KIND, "Received command message: {:?}", payload); + + let delivery = match serde_json::from_slice::(&message.payload) { + Ok(envelope) => envelope.delivery, + Err(e) => { + error!( + kind = H::KIND, + "Command message without a valid delivery block, dropping: {:#}", e + ); + return; + } + }; + let context = &listener.context; + let dispatch_id = delivery.dispatch_id.clone(); + + let Some(admitted) = context.restart.admit() else { + info!(kind = H::KIND, dispatch_id = %dispatch_id, "Restart pending, leaving the command to the backend retry"); + return; + }; + + let command = match serde_json::from_slice::(&message.payload) { + Ok(command) => command, + Err(e) => { + error!(kind = H::KIND, dispatch_id = %dispatch_id, "Failed to parse command message: {:#}", e); + context + .delivery + .settle( + machine_id, + delivery, + Outcome::Failed(format!("malformed payload: {e}")), + ) + .await; + return; + } + }; + + context.delivery.acked(machine_id, &delivery).await; + let outcome = listener.handler.run(machine_id, &delivery, command).await; + match &outcome { + Outcome::Done => info!(kind = H::KIND, dispatch_id = %dispatch_id, "Command done"), + Outcome::RestartRequired => { + info!(kind = H::KIND, dispatch_id = %dispatch_id, "Command applied, restart required") + } + Outcome::Failed(reason) => { + error!(kind = H::KIND, dispatch_id = %dispatch_id, "Command failed: {}", reason) + } + } + + if outcome != Outcome::RestartRequired { + context.delivery.settle(machine_id, delivery, outcome).await; + return; + } + drop(admitted); + context.restart.drain_and_request().await; +} diff --git a/clients/openframe-client/src/listener/command/delivery.rs b/clients/openframe-client/src/listener/command/delivery.rs new file mode 100644 index 0000000000..2a7b0b7e7f --- /dev/null +++ b/clients/openframe-client/src/listener/command/delivery.rs @@ -0,0 +1,126 @@ +use std::sync::Arc; + +use tokio::sync::Notify; +use tokio::time::{timeout, Duration}; +use tracing::{error, info, warn}; + +use super::Outcome; +use crate::config::update_config::FLUSH_PUBLISH_TIMEOUT_SECS; +use crate::models::{delivery_result_subject, DeliveryRef, DeliveryResult, DeliveryResultMessage}; +use crate::services::nats_message_publisher::NatsMessagePublisher; +use crate::services::result_store::ResultStore; + +#[derive(Clone)] +pub(super) struct Delivery { + publisher: NatsMessagePublisher, + store: Arc, + flush_notify: Arc, +} + +impl Delivery { + pub(super) fn new( + publisher: NatsMessagePublisher, + store: Arc, + flush_notify: Arc, + ) -> Self { + Self { + publisher, + store, + flush_notify, + } + } + + pub(super) async fn acked(&self, machine_id: &str, delivery: &DeliveryRef) { + let message = DeliveryResultMessage { + delivery: delivery.clone(), + result: DeliveryResult::Acked, + error: None, + }; + let subject = delivery_result_subject(machine_id); + let published = match serde_json::to_vec(&message) { + Ok(bytes) => timeout( + Duration::from_secs(FLUSH_PUBLISH_TIMEOUT_SECS), + self.publisher.publish_raw(&subject, &bytes), + ) + .await + .unwrap_or_else(|_| Err(anyhow::anyhow!("publish timed out"))), + Err(e) => Err(e.into()), + }; + match published { + Ok(()) => info!(dispatch_id = %delivery.dispatch_id, "Delivery ACKED"), + Err(e) => warn!( + dispatch_id = %delivery.dispatch_id, + "Failed to report ACKED, the backend retry re-sends the command: {:#}", e + ), + } + } + + pub(super) async fn settle( + &self, + machine_id: &str, + delivery: DeliveryRef, + outcome: Outcome, + ) -> bool { + let dispatch_id = delivery.dispatch_id.clone(); + let (result, error) = match outcome { + Outcome::Done | Outcome::RestartRequired => (DeliveryResult::Done, None), + Outcome::Failed(reason) => (DeliveryResult::Failed, Some(reason)), + }; + let message = DeliveryResultMessage { + delivery, + result, + error, + }; + let subject = delivery_result_subject(machine_id); + + if self.store.enabled() { + let enqueued = match serde_json::to_vec(&message) { + Ok(bytes) => { + self.store + .enqueue(format!("delivery:{}", dispatch_id), subject, bytes) + .await + } + Err(e) => Err(e.into()), + }; + return match enqueued { + Ok(()) => { + self.flush_notify.notify_one(); + info!(dispatch_id = %dispatch_id, ?result, "Delivery result queued"); + true + } + Err(e) => { + error!( + dispatch_id = %dispatch_id, + ?result, + "Failed to persist delivery result, the backend closes it on timeout: {:#}", e + ); + false + } + }; + } + + let published = match serde_json::to_vec(&message) { + Ok(bytes) => timeout( + Duration::from_secs(FLUSH_PUBLISH_TIMEOUT_SECS), + self.publisher.publish_acked(&subject, &bytes), + ) + .await + .unwrap_or_else(|_| Err(anyhow::anyhow!("publish timed out"))), + Err(e) => Err(e.into()), + }; + match published { + Ok(()) => { + info!(dispatch_id = %dispatch_id, ?result, "Delivery result published"); + true + } + Err(e) => { + error!( + dispatch_id = %dispatch_id, + ?result, + "Failed to publish delivery result without a durable store, the backend closes it on timeout: {:#}", e + ); + false + } + } + } +} diff --git a/clients/openframe-client/src/listener/command/handlers.rs b/clients/openframe-client/src/listener/command/handlers.rs new file mode 100644 index 0000000000..3c730de1f7 --- /dev/null +++ b/clients/openframe-client/src/listener/command/handlers.rs @@ -0,0 +1,65 @@ +use std::sync::Arc; + +use anyhow::Result; +use tracing::info; + +use super::{CommandHandler, Outcome}; +use crate::models::{DeliveryRef, FeatureFlagsMessage}; +use crate::services::feature_flags_service::{FeatureFlagsService, FlagsApplied}; + +#[derive(Clone)] +pub struct FeatureFlagsHandler { + service: Arc, +} + +impl FeatureFlagsHandler { + pub fn new(service: Arc) -> Self { + Self { service } + } +} + +impl CommandHandler for FeatureFlagsHandler { + type Command = FeatureFlagsMessage; + + const KIND: &'static str = "feature-flags"; + + fn subject(machine_id: &str) -> String { + format!("machine.{}.feature-flags", machine_id) + } + + async fn run( + &self, + machine_id: &str, + delivery: &DeliveryRef, + command: Self::Command, + ) -> Outcome { + let Some(sequence) = delivery.sequence else { + return Outcome::Failed("feature flags without a sequence".to_string()); + }; + let applied_sequence = self.service.applied_sequence(machine_id); + if sequence <= applied_sequence { + info!( + sequence, + applied_sequence, "Stale feature flags, already applied a newer set" + ); + return Outcome::Done; + } + match self.service.apply(machine_id, delivery, sequence, command) { + Ok(FlagsApplied::Unchanged) => Outcome::Done, + Ok(FlagsApplied::RestartRequired) => Outcome::RestartRequired, + Err(e) => Outcome::Failed(format!("{e:#}")), + } + } + + fn awaiting_restart(&self) -> Option { + self.service.awaiting_restart() + } + + fn restart_reported(&self) -> Result<()> { + self.service.restart_reported() + } +} + +#[cfg(test)] +#[path = "handlers_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/listener/command/handlers_tests.rs b/clients/openframe-client/src/listener/command/handlers_tests.rs new file mode 100644 index 0000000000..d0b2d80b13 --- /dev/null +++ b/clients/openframe-client/src/listener/command/handlers_tests.rs @@ -0,0 +1,119 @@ +use super::*; +use tempfile::TempDir; + +const MACHINE: &str = "machine-1"; + +fn handler() -> (FeatureFlagsHandler, TempDir) { + let dir = tempfile::tempdir().unwrap(); + let service = FeatureFlagsService::load(dir.path().join("feature_flags.json")); + (FeatureFlagsHandler::new(Arc::new(service)), dir) +} + +fn delivery(sequence: Option) -> DeliveryRef { + DeliveryRef { + delivery_type: "CLIENT_FEATURE_FLAGS".to_string(), + target_id: "tenant".to_string(), + dispatch_id: "dispatch".to_string(), + sequence, + } +} + +fn message(flags: serde_json::Value) -> FeatureFlagsMessage { + serde_json::from_value(serde_json::json!({ "flags": flags })).unwrap() +} + +#[test] +fn subject_is_machine_scoped() { + assert_eq!( + FeatureFlagsHandler::subject("abc123"), + "machine.abc123.feature-flags" + ); +} + +#[tokio::test] +async fn flags_without_a_sequence_fail() { + let (handler, _dir) = handler(); + + let outcome = handler + .run( + MACHINE, + &delivery(None), + message(serde_json::json!({ "delivery-tool-installation": true })), + ) + .await; + + assert!(matches!(outcome, Outcome::Failed(_))); + assert_eq!(handler.awaiting_restart(), None); +} + +#[tokio::test] +async fn a_changed_set_requires_a_restart_and_an_unchanged_one_does_not() { + let (handler, _dir) = handler(); + + let unchanged = handler + .run( + MACHINE, + &delivery(Some(1)), + message(serde_json::json!({ "delivery-tool-installation": false })), + ) + .await; + let changed = handler + .run( + MACHINE, + &delivery(Some(2)), + message(serde_json::json!({ "delivery-tool-installation": true })), + ) + .await; + + assert_eq!(unchanged, Outcome::Done); + assert_eq!(changed, Outcome::RestartRequired); +} + +#[tokio::test] +async fn stale_or_repeated_sequence_is_done_without_applying() { + let (handler, _dir) = handler(); + handler + .run( + MACHINE, + &delivery(Some(5)), + message(serde_json::json!({ "delivery-tool-installation": true })), + ) + .await; + + for stale in [5, 4] { + let outcome = handler + .run( + MACHINE, + &delivery(Some(stale)), + message(serde_json::json!({})), + ) + .await; + + assert_eq!(outcome, Outcome::Done, "sequence {stale}"); + assert_eq!(handler.service.applied_sequence(MACHINE), 5); + assert!(handler.awaiting_restart().is_some()); + } +} + +#[tokio::test] +async fn a_reregistered_machine_applies_a_lower_sequence() { + let (handler, _dir) = handler(); + handler + .run( + MACHINE, + &delivery(Some(50)), + message(serde_json::json!({ "delivery-tool-installation": true })), + ) + .await; + + let outcome = handler + .run( + "machine-2", + &delivery(Some(3)), + message(serde_json::json!({ "delivery-tool-installation": true })), + ) + .await; + + assert_eq!(outcome, Outcome::RestartRequired); + assert_eq!(handler.service.applied_sequence("machine-2"), 3); +} diff --git a/clients/openframe-client/src/listener/command/integration_tests.rs b/clients/openframe-client/src/listener/command/integration_tests.rs new file mode 100644 index 0000000000..b821d2fe70 --- /dev/null +++ b/clients/openframe-client/src/listener/command/integration_tests.rs @@ -0,0 +1,425 @@ +use std::path::PathBuf; +use std::sync::Arc; + +use futures::StreamExt; +use serde_json::json; +use tempfile::TempDir; +use tokio::sync::Notify; +use tokio::time::{timeout, Duration}; + +use super::{CommandContext, CommandListener, FeatureFlagsHandler}; +use crate::clients::AuthClient; +use crate::models::{DeliveryResult, DeliveryResultMessage, FeatureFlagsState}; +use crate::platform::DirectoryManager; +use crate::services::encryption_service::EncryptionService; +use crate::services::feature_flags_service::FeatureFlagsService; +use crate::services::local_tls_config_provider::LocalTlsConfigProvider; +use crate::services::nats_connection_manager::NatsConnectionManager; +use crate::services::nats_message_publisher::NatsMessagePublisher; +use crate::services::result_outbox_run_manager::ResultOutboxRunManager; +use crate::services::result_store::ResultStore; +use crate::services::shared_token_service::SharedTokenService; +use crate::services::tool_run_manager::ToolRunManager; +use crate::services::{ + AgentAuthService, AgentConfigurationService, DeactivationService, InitialConfigurationService, + InstalledToolsService, MachineIdService, ToolCommandParamsResolver, ToolKillService, ToolOps, +}; +use crate::test_support::NatsServer; + +const MACHINE: &str = "it-machine"; +const BOUND: Duration = Duration::from_secs(10); +const QUIET: Duration = Duration::from_millis(500); +const PROBE: &str = "probe"; + +struct Agent { + _server: NatsServer, + _dir: TempDir, + tester: async_nats::Client, + results: async_nats::Subscriber, + unclaimed: Vec, + manager: NatsConnectionManager, + config: AgentConfigurationService, + tool_run_manager: ToolRunManager, + publisher: NatsMessagePublisher, + store: Arc, + flush: Arc, + flags_path: PathBuf, +} + +impl Agent { + async fn start() -> Option<(Self, CommandContext)> { + let server = NatsServer::start()?; + let dir = tempfile::tempdir().unwrap(); + let directories = DirectoryManager::with_custom_dirs( + dir.path().join("logs"), + dir.path().join("app"), + dir.path().join("secured"), + ); + let config = AgentConfigurationService::new(directories.clone()).unwrap(); + config + .save_registration_data(MACHINE.into(), "client".into(), "secret".into()) + .await + .unwrap(); + let initial = InitialConfigurationService::new(directories.clone()).unwrap(); + let tool_run_manager = ToolRunManager::new( + InstalledToolsService::new(directories.clone()).unwrap(), + ToolCommandParamsResolver::new(directories.clone(), initial.clone(), config.clone()), + ToolKillService::new(), + ToolOps::default(), + ); + let deactivation = DeactivationService::new(&directories); + let auth = AgentAuthService::new( + AuthClient::new( + "http://127.0.0.1:9".into(), + reqwest::Client::new(), + deactivation.clone(), + ), + config.clone(), + SharedTokenService::new(directories.clone(), EncryptionService::new()), + ); + let manager = NatsConnectionManager::new( + server.url(), + config.clone(), + initial.clone(), + auth, + LocalTlsConfigProvider::new(initial), + deactivation, + MachineIdService::new(&directories), + ); + manager.use_client_for_tests(async_nats::connect(server.url()).await.unwrap()); + + let tester = async_nats::connect(server.url()).await.unwrap(); + async_nats::jetstream::new(tester.clone()) + .create_stream(async_nats::jetstream::stream::Config { + name: "DELIVERY_RESULT".into(), + subjects: vec!["machine.*.delivery.result".into()], + ..Default::default() + }) + .await + .unwrap(); + let results = tester + .subscribe(format!("machine.{MACHINE}.delivery.result")) + .await + .unwrap(); + tester.flush().await.unwrap(); + + let publisher = NatsMessagePublisher::new(manager.clone()); + let store = Arc::new(ResultStore::open_or_degrade( + directories.secured_dir().join("scheduled_outbox.redb"), + )); + let flush = Arc::new(Notify::new()); + ResultOutboxRunManager::new(store.clone(), Arc::new(publisher.clone()), flush.clone()) + .start(); + + let mut agent = Self { + _server: server, + flags_path: directories.secured_dir().join("feature_flags.json"), + _dir: dir, + tester, + results, + unclaimed: Vec::new(), + manager, + config, + tool_run_manager, + publisher, + store, + flush, + }; + let context = agent.launch().await; + Some((agent, context)) + } + + async fn launch(&mut self) -> CommandContext { + let context = CommandContext::new( + self.manager.clone(), + self.config.clone(), + ToolOps::default(), + self.tool_run_manager.clone(), + self.publisher.clone(), + self.store.clone(), + self.flush.clone(), + ); + let service = Arc::new(FeatureFlagsService::load(self.flags_path.clone())); + CommandListener::new(FeatureFlagsHandler::new(service), context.clone()) + .start() + .await + .unwrap(); + self.until_listening().await; + context + } + + async fn until_listening(&mut self) { + timeout(BOUND, async { + loop { + self.send(command(PROBE, Some(0), json!({}))).await; + if timeout(Duration::from_millis(200), self.next_result(PROBE)) + .await + .is_ok() + { + return; + } + } + }) + .await + .expect("the flags listener never answered"); + } + + async fn send(&self, payload: serde_json::Value) { + self.tester + .publish( + format!("machine.{MACHINE}.feature-flags"), + serde_json::to_vec(&payload).unwrap().into(), + ) + .await + .unwrap(); + self.tester.flush().await.unwrap(); + } + + async fn next_result(&mut self, dispatch_id: &str) -> DeliveryResultMessage { + if let Some(index) = self + .unclaimed + .iter() + .position(|result| result.delivery.dispatch_id == dispatch_id) + { + return self.unclaimed.remove(index); + } + loop { + let message = self + .results + .next() + .await + .expect("result subscription ended"); + let result: DeliveryResultMessage = serde_json::from_slice(&message.payload).unwrap(); + if result.delivery.dispatch_id == dispatch_id { + return result; + } + self.unclaimed.push(result); + } + } + + async fn expect( + &mut self, + dispatch_id: &str, + expected: DeliveryResult, + ) -> DeliveryResultMessage { + let result = timeout(BOUND, self.next_result(dispatch_id)) + .await + .unwrap_or_else(|_| panic!("no {expected:?} for {dispatch_id}")); + assert_eq!(result.result, expected, "{dispatch_id}: {result:?}"); + result + } + + async fn expect_silence(&mut self, dispatch_id: &str) { + let heard = timeout(QUIET, self.next_result(dispatch_id)).await; + assert!( + heard.is_err(), + "unexpected result for {dispatch_id}: {heard:?}" + ); + } + + fn persisted(&self) -> FeatureFlagsState { + serde_json::from_slice(&std::fs::read(&self.flags_path).unwrap()).unwrap() + } +} + +fn command( + dispatch_id: &str, + sequence: Option, + flags: serde_json::Value, +) -> serde_json::Value { + let mut delivery = json!({ + "type": "CLIENT_FEATURE_FLAGS", + "targetId": "tenant", + "dispatchId": dispatch_id, + }); + if let Some(sequence) = sequence { + delivery["sequence"] = json!(sequence); + } + json!({ "delivery": delivery, "flags": flags }) +} + +#[tokio::test] +async fn unchanged_flags_are_acked_done_and_persisted() { + let Some((mut agent, _context)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d1", + Some(1), + json!({ "delivery-tool-installation": false }), + )) + .await; + + agent.expect("d1", DeliveryResult::Acked).await; + agent.expect("d1", DeliveryResult::Done).await; + let persisted = agent.persisted(); + assert_eq!(persisted.sequence, 1); + assert_eq!(persisted.awaiting_restart, None); +} + +#[tokio::test] +async fn changed_flags_restart_and_the_next_process_reports_done() { + let Some((mut agent, context)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d2", + Some(2), + json!({ "delivery-tool-installation": true }), + )) + .await; + + agent.expect("d2", DeliveryResult::Acked).await; + timeout(BOUND, context.restart_requested()) + .await + .expect("restart was never requested"); + agent.expect_silence("d2").await; + assert_eq!( + agent + .persisted() + .awaiting_restart + .map(|delivery| delivery.dispatch_id), + Some("d2".to_string()) + ); + + agent.launch().await; + + agent.expect("d2", DeliveryResult::Done).await; + timeout(BOUND, async { + while agent.persisted().awaiting_restart.is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the restart marker was never cleared"); +} + +#[tokio::test] +async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { + let Some((mut agent, context)) = Agent::start().await else { + return; + }; + agent + .send(command( + "d3", + Some(3), + json!({ "delivery-tool-installation": true }), + )) + .await; + timeout(BOUND, context.restart_requested()) + .await + .expect("restart was never requested"); + + agent + .send(command( + "d4", + Some(4), + json!({ "delivery-tool-installation": false }), + )) + .await; + + agent.expect_silence("d4").await; + assert_eq!(agent.persisted().sequence, 3); +} + +#[tokio::test] +async fn stale_sequence_is_done_without_applying() { + let Some((mut agent, _context)) = Agent::start().await else { + return; + }; + agent.send(command("d5", Some(5), json!({}))).await; + agent.expect("d5", DeliveryResult::Acked).await; + agent.expect("d5", DeliveryResult::Done).await; + + agent + .send(command( + "d6", + Some(4), + json!({ "delivery-tool-installation": false }), + )) + .await; + + agent.expect("d6", DeliveryResult::Acked).await; + agent.expect("d6", DeliveryResult::Done).await; + assert_eq!(agent.persisted().sequence, 5); +} + +#[tokio::test] +async fn malformed_flags_fail_with_a_reason() { + let Some((mut agent, _context)) = Agent::start().await else { + return; + }; + + for (dispatch_id, payload) in [ + ("d7", command("d7", Some(7), json!("not-a-map"))), + ("d7-null", command("d7-null", Some(7), json!(null))), + ( + "d7-missing", + json!({ "delivery": { "type": "CLIENT_FEATURE_FLAGS", "targetId": "tenant", "dispatchId": "d7-missing", "sequence": 7 } }), + ), + ] { + agent.send(payload).await; + + let failed = agent.expect(dispatch_id, DeliveryResult::Failed).await; + assert!( + failed.error.unwrap().contains("malformed payload"), + "{dispatch_id}" + ); + } + assert!(!agent.flags_path.exists()); +} + +#[tokio::test] +async fn a_message_without_delivery_gets_no_answer() { + let Some((mut agent, _context)) = Agent::start().await else { + return; + }; + + agent + .send(json!({ "flags": { "delivery-tool-installation": true } })) + .await; + agent + .send(command( + "d8", + Some(8), + json!({ "delivery-tool-installation": false }), + )) + .await; + + let first = timeout(BOUND, async { + loop { + let message = agent.results.next().await.unwrap(); + let result: DeliveryResultMessage = serde_json::from_slice(&message.payload).unwrap(); + if result.delivery.dispatch_id != PROBE { + return result; + } + } + }) + .await + .unwrap(); + assert_eq!(first.delivery.dispatch_id, "d8"); + assert_eq!(agent.persisted().sequence, 8); +} + +#[tokio::test] +async fn flags_without_a_sequence_fail() { + let Some((mut agent, _context)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d9", + None, + json!({ "delivery-tool-installation": true }), + )) + .await; + + agent.expect("d9", DeliveryResult::Acked).await; + let failed = agent.expect("d9", DeliveryResult::Failed).await; + assert!(failed.error.unwrap().contains("sequence")); +} diff --git a/clients/openframe-client/src/listener/command/mod.rs b/clients/openframe-client/src/listener/command/mod.rs new file mode 100644 index 0000000000..1d69ea2e19 --- /dev/null +++ b/clients/openframe-client/src/listener/command/mod.rs @@ -0,0 +1,127 @@ +mod core_nats; +mod delivery; +mod handlers; +mod restart; + +use std::future::Future; +use std::sync::Arc; + +use anyhow::Result; +use serde::de::DeserializeOwned; +use tokio::sync::Notify; +use tokio::time::Duration; +use tracing::{error, info}; + +use crate::config::update_config::RECONNECTION_DELAY_MS; +use crate::models::DeliveryRef; +use crate::services::nats_connection_manager::NatsConnectionManager; +use crate::services::nats_message_publisher::NatsMessagePublisher; +use crate::services::result_store::ResultStore; +use crate::services::tool_ops::ToolOps; +use crate::services::tool_run_manager::ToolRunManager; +use crate::services::AgentConfigurationService; + +use delivery::Delivery; +use restart::RestartGate; + +pub use handlers::FeatureFlagsHandler; + +pub trait CommandHandler: Clone + Send + Sync + 'static { + type Command: DeserializeOwned + Send + 'static; + + const KIND: &'static str; + + fn subject(machine_id: &str) -> String; + + fn run( + &self, + machine_id: &str, + delivery: &DeliveryRef, + command: Self::Command, + ) -> impl Future + Send; + + fn awaiting_restart(&self) -> Option { + None + } + + fn restart_reported(&self) -> Result<()> { + Ok(()) + } +} + +#[derive(Debug, PartialEq)] +pub enum Outcome { + Done, + RestartRequired, + Failed(String), +} + +#[derive(Clone)] +pub struct CommandContext { + nats_connection_manager: NatsConnectionManager, + config_service: AgentConfigurationService, + restart: RestartGate, + delivery: Delivery, +} + +impl CommandContext { + pub fn new( + nats_connection_manager: NatsConnectionManager, + config_service: AgentConfigurationService, + tool_ops: ToolOps, + tool_run_manager: ToolRunManager, + publisher: NatsMessagePublisher, + result_store: Arc, + flush_notify: Arc, + ) -> Self { + Self { + nats_connection_manager, + config_service, + restart: RestartGate::new(tool_ops, tool_run_manager), + delivery: Delivery::new(publisher, result_store, flush_notify), + } + } +} + +impl CommandContext { + pub async fn restart_requested(&self) { + self.restart.requested().await; + } +} + +#[derive(Clone)] +pub struct CommandListener { + handler: H, + context: CommandContext, +} + +impl CommandListener { + pub fn new(handler: H, context: CommandContext) -> Self { + Self { handler, context } + } + + pub async fn start(&self) -> Result> { + let listener = self.clone(); + let handle = tokio::spawn(async move { + core_nats::report_restart_completion(&listener).await; + loop { + info!(kind = H::KIND, "Starting command listener..."); + match core_nats::listen(&listener).await { + Ok(()) => info!(kind = H::KIND, "Command listener rebinding"), + Err(e) => error!(kind = H::KIND, "Command listener error: {:#}", e), + } + info!( + kind = H::KIND, + "Reconnecting command listener in {} seconds...", + RECONNECTION_DELAY_MS / 1000 + ); + tokio::time::sleep(Duration::from_millis(RECONNECTION_DELAY_MS)).await; + } + }); + Ok(handle) + } +} + +#[cfg(test)] +#[path = "integration_tests.rs"] +mod integration_tests; diff --git a/clients/openframe-client/src/listener/command/restart.rs b/clients/openframe-client/src/listener/command/restart.rs new file mode 100644 index 0000000000..2072df2054 --- /dev/null +++ b/clients/openframe-client/src/listener/command/restart.rs @@ -0,0 +1,73 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; + +use tokio::sync::{Notify, OwnedRwLockReadGuard, RwLock}; +use tokio::time::{sleep, timeout_at, Duration, Instant}; +use tracing::{info, warn}; + +use crate::config::update_config::{SELF_RESTART_IDLE_POLL_SECS, SELF_RESTART_IDLE_WAIT_SECS}; +use crate::services::tool_ops::ToolOps; +use crate::services::tool_run_manager::ToolRunManager; + +#[derive(Clone)] +pub(super) struct RestartGate { + closing: Arc, + work: Arc>, + requested: Arc, + tool_ops: ToolOps, + tool_run_manager: ToolRunManager, +} + +impl RestartGate { + pub(super) fn new(tool_ops: ToolOps, tool_run_manager: ToolRunManager) -> Self { + Self { + closing: Arc::new(AtomicBool::new(false)), + work: Arc::new(RwLock::new(())), + requested: Arc::new(Notify::new()), + tool_ops, + tool_run_manager, + } + } + + pub(super) fn admit(&self) -> Option> { + if self.closing.load(Ordering::Acquire) { + return None; + } + self.work.clone().try_read_owned().ok() + } + + pub(super) async fn drain_and_request(&self) { + self.closing.store(true, Ordering::Release); + let deadline = Instant::now() + Duration::from_secs(SELF_RESTART_IDLE_WAIT_SECS); + info!("Restart scheduled: no new commands are taken, waiting for running ones to finish"); + + if timeout_at(deadline, self.work.write()).await.is_err() { + warn!( + "Commands still running after {}s, restarting anyway", + SELF_RESTART_IDLE_WAIT_SECS + ); + } + while self.busy().await && Instant::now() < deadline { + sleep(Duration::from_secs(SELF_RESTART_IDLE_POLL_SECS)).await; + } + if self.busy().await { + warn!( + "Tool operation or client update still running after {}s, restarting anyway", + SELF_RESTART_IDLE_WAIT_SECS + ); + } + self.requested.notify_one(); + } + + async fn busy(&self) -> bool { + self.tool_ops.any_busy() || self.tool_run_manager.is_client_update_pending().await + } + + pub(super) async fn requested(&self) { + self.requested.notified().await; + } +} + +#[cfg(test)] +#[path = "restart_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/listener/command/restart_tests.rs b/clients/openframe-client/src/listener/command/restart_tests.rs new file mode 100644 index 0000000000..12a60677d7 --- /dev/null +++ b/clients/openframe-client/src/listener/command/restart_tests.rs @@ -0,0 +1,118 @@ +use super::*; +use crate::platform::DirectoryManager; +use crate::services::{ + AgentConfigurationService, InitialConfigurationService, InstalledToolsService, + ToolCommandParamsResolver, ToolKillService, +}; +use tempfile::TempDir; +use tokio::time::timeout; + +const BOUND: Duration = Duration::from_secs(10); + +async fn until_closed(gate: &RestartGate) { + timeout(BOUND, async { + while gate.admit().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the gate never closed"); +} + +fn start_drain(gate: &RestartGate) -> tokio::task::JoinHandle<()> { + let gate = gate.clone(); + tokio::spawn(async move { gate.drain_and_request().await }) +} + +fn gate_with(tool_ops: ToolOps) -> (RestartGate, ToolRunManager, TempDir) { + let dir = TempDir::new().unwrap(); + let directories = DirectoryManager::with_custom_dirs( + dir.path().join("logs"), + dir.path().join("app_support"), + dir.path().join("secured"), + ); + let tool_run_manager = ToolRunManager::new( + InstalledToolsService::new(directories.clone()).unwrap(), + ToolCommandParamsResolver::new( + directories.clone(), + InitialConfigurationService::new(directories.clone()).unwrap(), + AgentConfigurationService::new(directories).unwrap(), + ), + ToolKillService::new(), + tool_ops.clone(), + ); + ( + RestartGate::new(tool_ops, tool_run_manager.clone()), + tool_run_manager, + dir, + ) +} + +async fn until_requested(gate: &RestartGate) { + timeout(BOUND, gate.requested()) + .await + .expect("restart was never requested"); +} + +#[tokio::test] +async fn admits_concurrent_commands_while_no_restart_is_pending() { + let (gate, _manager, _dir) = gate_with(ToolOps::default()); + + let first = gate.admit(); + let second = gate.admit(); + + assert!(first.is_some()); + assert!(second.is_some()); +} + +#[tokio::test] +async fn drain_refuses_new_commands_and_waits_for_the_admitted_one() { + let (gate, _manager, _dir) = gate_with(ToolOps::default()); + let running = gate.admit().expect("admitted before the restart"); + + let draining = start_drain(&gate); + until_closed(&gate).await; + assert!(!draining.is_finished()); + + drop(running); + until_requested(&gate).await; + + assert!(gate.admit().is_none()); +} + +#[tokio::test] +async fn drain_waits_for_a_busy_tool() { + let tool_ops = ToolOps::default(); + let (gate, _manager, _dir) = gate_with(tool_ops.clone()); + let op = tool_ops.try_lock("fleetmdm-agent").unwrap().mark_busy(); + + let draining = start_drain(&gate); + until_closed(&gate).await; + assert!(!draining.is_finished()); + + drop(op); + until_requested(&gate).await; +} + +#[tokio::test] +async fn the_gate_stays_closed_after_the_restart_is_requested() { + let (gate, _manager, _dir) = gate_with(ToolOps::default()); + + start_drain(&gate); + until_requested(&gate).await; + + assert!(gate.admit().is_none()); +} + +#[tokio::test] +async fn drain_waits_for_a_pending_client_update() { + let (gate, tool_run_manager, _dir) = gate_with(ToolOps::default()); + tool_run_manager.mark_client_update_pending().await; + + let draining = start_drain(&gate); + until_closed(&gate).await; + assert!(!draining.is_finished()); + + tool_run_manager.clear_client_update_pending().await; + until_requested(&gate).await; +} diff --git a/clients/openframe-client/src/listener/execution_listener.rs b/clients/openframe-client/src/listener/execution_listener.rs index 6e4e5a9af6..958d562527 100644 --- a/clients/openframe-client/src/listener/execution_listener.rs +++ b/clients/openframe-client/src/listener/execution_listener.rs @@ -13,7 +13,8 @@ use crate::config::update_config::{ FALLBACK_PUBLISH_MAX_RETRY_DELAY_MS, RECONNECTION_DELAY_MS, }; use crate::models::{ - is_ack_subject, ExecutionAck, ExecutionMessage, ExecutionRequest, RmmResult, EXECUTION_ACK_KIND, + requires_stream_ack, ExecutionAck, ExecutionMessage, ExecutionRequest, RmmResult, + EXECUTION_ACK_KIND, }; use crate::services::execution_service::ExecutionService; use crate::services::nats_connection_manager::NatsConnectionManager; @@ -254,7 +255,7 @@ impl ExecutionListener { fn spawn_fallback_publish(&self, subject: String, bytes: Vec) { let publisher = self.nats_message_publisher.clone(); - let acked = is_ack_subject(&subject); + let acked = requires_stream_ack(&subject); tokio::spawn(async move { let max_backoff = Duration::from_millis(FALLBACK_PUBLISH_MAX_RETRY_DELAY_MS); let mut backoff = Duration::from_millis(FALLBACK_PUBLISH_INITIAL_RETRY_DELAY_MS); diff --git a/clients/openframe-client/src/listener/mod.rs b/clients/openframe-client/src/listener/mod.rs index bbf2880fe4..b8cac5820c 100644 --- a/clients/openframe-client/src/listener/mod.rs +++ b/clients/openframe-client/src/listener/mod.rs @@ -1,5 +1,6 @@ pub mod client_uninstall_message_listener; pub mod client_update_gate; +pub mod command; pub mod execution_listener; pub mod openframe_client_update_listener; pub mod remote_access_message_listener; @@ -9,6 +10,7 @@ pub mod tool_restart_message_listener; pub mod tool_uninstall_message_listener; pub use client_uninstall_message_listener::ClientUninstallMessageListener; +pub use command::CommandListener; pub use execution_listener::ExecutionListener; pub use openframe_client_update_listener::OpenFrameClientUpdateListener; pub use remote_access_message_listener::RemoteAccessMessageListener; diff --git a/clients/openframe-client/src/models/delivery.rs b/clients/openframe-client/src/models/delivery.rs new file mode 100644 index 0000000000..a4cc814dd7 --- /dev/null +++ b/clients/openframe-client/src/models/delivery.rs @@ -0,0 +1,34 @@ +use serde::{Deserialize, Serialize}; + +pub const DELIVERY_RESULT_KIND: &str = "delivery.result"; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct DeliveryRef { + #[serde(rename = "type")] + pub delivery_type: String, + pub target_id: String, + pub dispatch_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub sequence: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "SCREAMING_SNAKE_CASE")] +pub enum DeliveryResult { + Acked, + Done, + Failed, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct DeliveryResultMessage { + pub delivery: DeliveryRef, + pub result: DeliveryResult, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +pub fn delivery_result_subject(machine_id: &str) -> String { + format!("machine.{}.{}", machine_id, DELIVERY_RESULT_KIND) +} diff --git a/clients/openframe-client/src/models/execution.rs b/clients/openframe-client/src/models/execution.rs index 8214e1118c..43505f6fab 100644 --- a/clients/openframe-client/src/models/execution.rs +++ b/clients/openframe-client/src/models/execution.rs @@ -172,8 +172,8 @@ pub struct ExecutionAck { pub const EXECUTION_ACK_KIND: &str = "execution.acknowledge"; -pub fn is_ack_subject(subject: &str) -> bool { - subject.ends_with(EXECUTION_ACK_KIND) +pub fn requires_stream_ack(subject: &str) -> bool { + subject.ends_with(EXECUTION_ACK_KIND) || subject.ends_with(super::DELIVERY_RESULT_KIND) } pub struct ExecutionRequest<'a> { diff --git a/clients/openframe-client/src/models/feature_flags.rs b/clients/openframe-client/src/models/feature_flags.rs new file mode 100644 index 0000000000..66b44aa76d --- /dev/null +++ b/clients/openframe-client/src/models/feature_flags.rs @@ -0,0 +1,77 @@ +use std::collections::{BTreeMap, HashSet}; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::DeliveryRef; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Feature { + DeliveryToolInstallation, + DeliveryClientUninstall, + DeliveryClientUpdate, +} + +impl Feature { + pub const ALL: [Feature; 3] = [ + Feature::DeliveryToolInstallation, + Feature::DeliveryClientUninstall, + Feature::DeliveryClientUpdate, + ]; + + pub fn name(self) -> &'static str { + match self { + Feature::DeliveryToolInstallation => "delivery-tool-installation", + Feature::DeliveryClientUninstall => "delivery-client-uninstall", + Feature::DeliveryClientUpdate => "delivery-client-update", + } + } + + pub fn default_enabled(self) -> bool { + match self { + Feature::DeliveryToolInstallation + | Feature::DeliveryClientUninstall + | Feature::DeliveryClientUpdate => false, + } + } +} + +#[derive(Debug, Deserialize)] +pub struct FeatureFlagsMessage { + pub flags: BTreeMap, +} + +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct FeatureFlagsState { + #[serde(default)] + pub machine_id: String, + pub sequence: i64, + pub flags: BTreeMap, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub awaiting_restart: Option, +} + +#[derive(Debug, Clone, Default, PartialEq)] +pub struct FeatureFlags { + enabled: HashSet, +} + +impl FeatureFlags { + pub fn from_flags(flags: &BTreeMap) -> Self { + Self { + enabled: Feature::ALL + .into_iter() + .filter(|feature| { + flags + .get(feature.name()) + .copied() + .unwrap_or(feature.default_enabled()) + }) + .collect(), + } + } + + pub fn is_enabled(&self, feature: Feature) -> bool { + self.enabled.contains(&feature) + } +} diff --git a/clients/openframe-client/src/models/mod.rs b/clients/openframe-client/src/models/mod.rs index 28ce7f07a9..fc68f4c9c7 100644 --- a/clients/openframe-client/src/models/mod.rs +++ b/clients/openframe-client/src/models/mod.rs @@ -2,9 +2,11 @@ pub mod agent_configuration; pub mod agent_registration_request; pub mod agent_registration_response; pub mod agent_token_response; +pub mod delivery; pub mod device_tag; pub mod download_configuration; pub mod execution; +pub mod feature_flags; pub mod hostname_report_message; pub mod initial_configuration; pub mod installed_agent_message; @@ -24,7 +26,7 @@ pub mod tool_version_overrides; pub mod update_state; pub use execution::{ - is_ack_subject, BootstrapRoute, BootstrapScriptMessage, CommandMessage, ExecutionAck, + requires_stream_ack, BootstrapRoute, BootstrapScriptMessage, CommandMessage, ExecutionAck, ExecutionMessage, ExecutionRequest, PrivilegeLevel, PrivilegePolicy, RmmResult, RoutedScriptMessage, ScriptEnvVar, ScriptMessage, ScriptRoute, ScriptScheduleExecutionItem, ScriptScheduleExecutionMessage, ScriptShell, ScriptSpec, SoftwareRoute, SoftwareScriptMessage, @@ -35,8 +37,13 @@ pub use agent_configuration::AgentConfiguration; pub use agent_registration_request::AgentRegistrationRequest; pub use agent_registration_response::AgentRegistrationResponse; pub use agent_token_response::AgentTokenResponse; +pub use delivery::{ + delivery_result_subject, DeliveryRef, DeliveryResult, DeliveryResultMessage, + DELIVERY_RESULT_KIND, +}; pub use device_tag::DeviceTag; pub use download_configuration::{DownloadConfiguration, InstallationType}; +pub use feature_flags::{Feature, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState}; pub use hostname_report_message::HostnameReportMessage; pub use initial_configuration::InitialConfiguration; pub use installed_agent_message::InstalledAgentMessage; diff --git a/clients/openframe-client/src/platform/detached_process.rs b/clients/openframe-client/src/platform/detached_process.rs new file mode 100644 index 0000000000..31af1932b0 --- /dev/null +++ b/clients/openframe-client/src/platform/detached_process.rs @@ -0,0 +1,68 @@ +use anyhow::Result; +use std::path::Path; + +/// Spawn ` ` detached, so it survives the service manager stopping +/// our own `com.openframe.client` service. macOS + Windows only. +pub fn spawn_detached(install_path: &Path, subcommand: &str) -> Result { + #[cfg(target_os = "macos")] + { + spawn_detached_macos(install_path, subcommand) + } + #[cfg(target_os = "windows")] + { + spawn_detached_windows(install_path, subcommand) + } + #[cfg(not(any(target_os = "macos", target_os = "windows")))] + { + let _ = (install_path, subcommand); + anyhow::bail!("detached client subcommands are not supported on this platform") + } +} + +/// Detached child in its own session (setsid) so it outlives launchd stopping our own +/// `com.openframe.client` service — no launchd job/plist needed. +#[cfg(target_os = "macos")] +fn spawn_detached_macos(install_path: &Path, subcommand: &str) -> Result { + use anyhow::Context; + use std::os::unix::process::CommandExt; + use std::process::{Command, Stdio}; + + let child = unsafe { + Command::new(install_path) + .arg(subcommand) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + // SAFETY: setsid() is async-signal-safe; it moves the child into a new session so + // launchd stopping our service can't signal it away mid-run. + .pre_exec(|| { + if libc::setsid() == -1 { + return Err(std::io::Error::last_os_error()); + } + Ok(()) + }) + .spawn() + } + .with_context(|| format!("Failed to spawn detached `{subcommand}` process"))?; + + Ok(child.id()) +} + +/// Detached child (inherits LocalSystem) so it survives the SCM stopping our service. +#[cfg(target_os = "windows")] +fn spawn_detached_windows(install_path: &Path, subcommand: &str) -> Result { + use anyhow::Context; + use std::os::windows::process::CommandExt; + use std::process::Command; + + const CREATE_NO_WINDOW: u32 = 0x0800_0000; + const DETACHED_PROCESS: u32 = 0x0000_0008; + + let child = Command::new(install_path) + .arg(subcommand) + .creation_flags(CREATE_NO_WINDOW | DETACHED_PROCESS) + .spawn() + .with_context(|| format!("Failed to spawn detached `{subcommand}` process"))?; + + Ok(child.id()) +} diff --git a/clients/openframe-client/src/platform/mod.rs b/clients/openframe-client/src/platform/mod.rs index 8d50da66d8..038b8a207a 100644 --- a/clients/openframe-client/src/platform/mod.rs +++ b/clients/openframe-client/src/platform/mod.rs @@ -1,5 +1,6 @@ pub mod binary_writer; pub mod console; +pub mod detached_process; pub mod directories; pub mod dmg_extractor; pub mod file_acl; diff --git a/clients/openframe-client/src/platform/uninstall.rs b/clients/openframe-client/src/platform/uninstall.rs index 545ab2e963..9285e1c92d 100644 --- a/clients/openframe-client/src/platform/uninstall.rs +++ b/clients/openframe-client/src/platform/uninstall.rs @@ -93,66 +93,9 @@ fn remove_update_temp_files(install_path: &Path) { /// Self-uninstall stops the `com.openframe.client` service (our own process), so it must run /// out-of-process. macOS + Windows only — Linux client self-uninstall is unsupported. pub fn spawn_detached_uninstall(install_path: &Path) -> Result<()> { - #[cfg(target_os = "macos")] - { - spawn_detached_uninstall_macos(install_path) - } - #[cfg(target_os = "windows")] - { - spawn_detached_uninstall_windows(install_path) - } - #[cfg(not(any(target_os = "macos", target_os = "windows")))] - { - let _ = install_path; - anyhow::bail!("client self-uninstall is not supported on this platform") - } -} - -/// Detached child in its own session (setsid) so it outlives launchd stopping our own -/// `com.openframe.client` service — no launchd job/plist needed. -#[cfg(target_os = "macos")] -fn spawn_detached_uninstall_macos(install_path: &Path) -> Result<()> { - use std::os::unix::process::CommandExt; - use std::process::{Command, Stdio}; - - let child = unsafe { - Command::new(install_path) - .arg(UNINSTALL_SUBCOMMAND) - .stdin(Stdio::null()) - .stdout(Stdio::null()) - .stderr(Stdio::null()) - // SAFETY: setsid() is async-signal-safe; it moves the child into a new session so - // launchd stopping our service can't signal it away mid-uninstall. - .pre_exec(|| { - if libc::setsid() == -1 { - return Err(std::io::Error::last_os_error()); - } - Ok(()) - }) - .spawn() - } - .context("Failed to spawn detached self-uninstall process")?; - - info!("Self-uninstall process launched (PID: {})", child.id()); - Ok(()) -} - -/// Detached child (inherits LocalSystem) so it survives the SCM stopping our service. -#[cfg(target_os = "windows")] -fn spawn_detached_uninstall_windows(install_path: &Path) -> Result<()> { - use std::os::windows::process::CommandExt; - use std::process::Command; - - const CREATE_NO_WINDOW: u32 = 0x0800_0000; - const DETACHED_PROCESS: u32 = 0x0000_0008; - - let child = Command::new(install_path) - .arg(UNINSTALL_SUBCOMMAND) - .creation_flags(CREATE_NO_WINDOW | DETACHED_PROCESS) - .spawn() - .context("Failed to spawn detached self-uninstall process")?; - - info!("Self-uninstall process launched (PID: {})", child.id()); + let pid = + crate::platform::detached_process::spawn_detached(install_path, UNINSTALL_SUBCOMMAND)?; + info!("Self-uninstall process launched (PID: {})", pid); Ok(()) } diff --git a/clients/openframe-client/src/services/feature_flags_service.rs b/clients/openframe-client/src/services/feature_flags_service.rs new file mode 100644 index 0000000000..3b6f21ed48 --- /dev/null +++ b/clients/openframe-client/src/services/feature_flags_service.rs @@ -0,0 +1,126 @@ +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::{Mutex, PoisonError}; + +use anyhow::{Context, Result}; +use serde_json::Value; +use tracing::{info, warn}; + +use crate::models::{DeliveryRef, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState}; +use crate::utils::fs::atomic_write; + +pub struct FeatureFlagsService { + path: PathBuf, + running: FeatureFlags, + state: Mutex, +} + +#[derive(Debug, PartialEq)] +pub enum FlagsApplied { + Unchanged, + RestartRequired, +} + +impl FeatureFlagsService { + pub fn load(path: PathBuf) -> Self { + let state = match std::fs::read(&path) { + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(state) => state, + Err(e) => { + warn!(path = %path.display(), "Feature flags file unreadable, all features off: {:#}", e); + FeatureFlagsState::default() + } + }, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => FeatureFlagsState::default(), + Err(e) => { + warn!(path = %path.display(), "Failed to read feature flags file, all features off: {:#}", e); + FeatureFlagsState::default() + } + }; + let running = FeatureFlags::from_flags(&state.flags); + info!(sequence = state.sequence, flags = ?state.flags, "Feature flags loaded"); + Self { + path, + running, + state: Mutex::new(state), + } + } + + pub fn flags(&self) -> &FeatureFlags { + &self.running + } + + pub fn applied_sequence(&self, machine_id: &str) -> i64 { + let state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + if state.machine_id == machine_id { + return state.sequence; + } + info!( + applied_for = %state.machine_id, + machine_id, + "Feature flags were applied for another registration, accepting any sequence" + ); + 0 + } + + pub fn apply( + &self, + machine_id: &str, + delivery: &DeliveryRef, + sequence: i64, + message: FeatureFlagsMessage, + ) -> Result { + let mut flags = BTreeMap::new(); + for (name, value) in message.flags { + match value { + Value::Bool(enabled) => { + flags.insert(name, enabled); + } + other => { + warn!(flag = %name, value = %other, "Ignoring feature flag with a non-boolean value") + } + } + } + let applied = if FeatureFlags::from_flags(&flags) == self.running { + FlagsApplied::Unchanged + } else { + FlagsApplied::RestartRequired + }; + let state = FeatureFlagsState { + machine_id: machine_id.to_string(), + sequence, + flags, + awaiting_restart: (applied == FlagsApplied::RestartRequired).then(|| delivery.clone()), + }; + self.persist(&state)?; + info!(sequence, flags = ?state.flags, ?applied, "Feature flags persisted"); + *self.state.lock().unwrap_or_else(PoisonError::into_inner) = state; + Ok(applied) + } + + pub fn awaiting_restart(&self) -> Option { + self.state + .lock() + .unwrap_or_else(PoisonError::into_inner) + .awaiting_restart + .clone() + } + + pub fn restart_reported(&self) -> Result<()> { + let mut state = self.state.lock().unwrap_or_else(PoisonError::into_inner); + let mut reported = state.clone(); + reported.awaiting_restart = None; + self.persist(&reported)?; + *state = reported; + Ok(()) + } + + fn persist(&self, state: &FeatureFlagsState) -> Result<()> { + atomic_write(&self.path, serde_json::to_vec_pretty(state)?) + .with_context(|| format!("failed to persist feature flags to {}", self.path.display())) + } +} + +#[cfg(test)] +#[path = "feature_flags_service_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/services/feature_flags_service_tests.rs b/clients/openframe-client/src/services/feature_flags_service_tests.rs new file mode 100644 index 0000000000..a3f749b1d5 --- /dev/null +++ b/clients/openframe-client/src/services/feature_flags_service_tests.rs @@ -0,0 +1,236 @@ +use super::*; +use crate::models::Feature; +use tempfile::TempDir; + +const MACHINE: &str = "machine-1"; + +fn flags_path() -> (PathBuf, TempDir) { + let dir = tempfile::tempdir().unwrap(); + (dir.path().join("feature_flags.json"), dir) +} + +fn delivery(sequence: i64) -> DeliveryRef { + DeliveryRef { + delivery_type: "CLIENT_FEATURE_FLAGS".to_string(), + target_id: "tenant".to_string(), + dispatch_id: format!("dispatch-{sequence}"), + sequence: Some(sequence), + } +} + +fn message(flags: serde_json::Value) -> FeatureFlagsMessage { + serde_json::from_value(serde_json::json!({ "flags": flags })).unwrap() +} + +fn apply(service: &FeatureFlagsService, sequence: i64, flags: serde_json::Value) -> FlagsApplied { + service + .apply(MACHINE, &delivery(sequence), sequence, message(flags)) + .unwrap() +} + +fn tool_installation_delivery(service: &FeatureFlagsService) -> bool { + service + .flags() + .is_enabled(Feature::DeliveryToolInstallation) +} + +#[test] +fn missing_file_means_every_feature_off() { + let (path, _dir) = flags_path(); + + let service = FeatureFlagsService::load(path); + + assert!(!tool_installation_delivery(&service)); + assert_eq!(service.applied_sequence(MACHINE), 0); + assert_eq!(service.awaiting_restart(), None); +} + +#[test] +fn corrupt_file_means_every_feature_off() { + let (path, _dir) = flags_path(); + std::fs::write(&path, b"{ not json").unwrap(); + + let service = FeatureFlagsService::load(path); + + assert!(!tool_installation_delivery(&service)); + assert_eq!(service.applied_sequence(MACHINE), 0); +} + +#[test] +fn enabling_a_feature_needs_a_restart_and_takes_effect_on_next_start() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path.clone()); + + let applied = apply( + &service, + 7, + serde_json::json!({ "delivery-tool-installation": true }), + ); + + assert_eq!(applied, FlagsApplied::RestartRequired); + assert!(!tool_installation_delivery(&service)); + assert_eq!(service.awaiting_restart(), Some(delivery(7))); + + let restarted = FeatureFlagsService::load(path); + assert!(tool_installation_delivery(&restarted)); + assert_eq!(restarted.applied_sequence(MACHINE), 7); + assert_eq!(restarted.awaiting_restart(), Some(delivery(7))); +} + +#[test] +fn reported_restart_is_not_reported_again() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path.clone()); + apply( + &service, + 7, + serde_json::json!({ "delivery-tool-installation": true }), + ); + + let restarted = FeatureFlagsService::load(path.clone()); + restarted.restart_reported().unwrap(); + + assert_eq!(restarted.awaiting_restart(), None); + let next_start = FeatureFlagsService::load(path); + assert_eq!(next_start.awaiting_restart(), None); + assert!(tool_installation_delivery(&next_start)); +} + +#[test] +fn absent_false_and_non_boolean_all_resolve_to_the_default() { + for flags in [ + serde_json::json!({}), + serde_json::json!({ "delivery-tool-installation": false }), + serde_json::json!({ "delivery-tool-installation": "true" }), + serde_json::json!({ "delivery-tool-installation": 1 }), + serde_json::json!({ "delivery-tool-installation": null }), + ] { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path); + + let applied = apply(&service, 1, flags.clone()); + + assert_eq!(applied, FlagsApplied::Unchanged, "flags: {flags}"); + assert_eq!(service.awaiting_restart(), None, "flags: {flags}"); + } +} + +#[test] +fn unknown_flags_are_kept_on_disk_but_change_nothing() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path.clone()); + + let applied = apply(&service, 3, serde_json::json!({ "FUTURE_FEATURE": true })); + + assert_eq!(applied, FlagsApplied::Unchanged); + let persisted: FeatureFlagsState = + serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!(persisted.flags.get("FUTURE_FEATURE"), Some(&true)); +} + +#[test] +fn flipping_back_before_the_restart_cancels_it() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path); + + apply( + &service, + 1, + serde_json::json!({ "delivery-tool-installation": true }), + ); + let applied = apply( + &service, + 2, + serde_json::json!({ "delivery-tool-installation": false }), + ); + + assert_eq!(applied, FlagsApplied::Unchanged); + assert_eq!(service.awaiting_restart(), None); + assert_eq!(service.applied_sequence(MACHINE), 2); +} + +#[test] +fn persisted_state_matches_the_documented_shape() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path.clone()); + + apply( + &service, + 42, + serde_json::json!({ "delivery-tool-installation": true }), + ); + + let persisted: serde_json::Value = + serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + persisted, + serde_json::json!({ + "machine_id": MACHINE, + "sequence": 42, + "flags": { "delivery-tool-installation": true }, + "awaiting_restart": { + "type": "CLIENT_FEATURE_FLAGS", + "targetId": "tenant", + "dispatchId": "dispatch-42", + "sequence": 42 + } + }) + ); +} + +#[test] +fn a_new_registration_accepts_any_sequence_and_keeps_flags_until_then() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path.clone()); + apply( + &service, + 50, + serde_json::json!({ "delivery-tool-installation": true }), + ); + + let restarted = FeatureFlagsService::load(path); + + assert_eq!(restarted.applied_sequence(MACHINE), 50); + assert_eq!(restarted.applied_sequence("machine-2"), 0); + assert!(tool_installation_delivery(&restarted)); +} + +#[test] +fn a_file_without_machine_id_accepts_the_first_set() { + let (path, _dir) = flags_path(); + std::fs::write( + &path, + br#"{ "sequence": 9, "flags": { "delivery-tool-installation": true } }"#, + ) + .unwrap(); + + let service = FeatureFlagsService::load(path); + + assert!(tool_installation_delivery(&service)); + assert_eq!(service.applied_sequence(MACHINE), 0); +} + +#[test] +fn the_backend_flag_names_map_to_their_features() { + let (path, _dir) = flags_path(); + let service = FeatureFlagsService::load(path.clone()); + + apply( + &service, + 1, + serde_json::json!({ + "delivery-tool-installation": true, + "delivery-client-uninstall": true, + "delivery-client-update": false + }), + ); + + let restarted = FeatureFlagsService::load(path); + assert!(restarted + .flags() + .is_enabled(Feature::DeliveryToolInstallation)); + assert!(restarted + .flags() + .is_enabled(Feature::DeliveryClientUninstall)); + assert!(!restarted.flags().is_enabled(Feature::DeliveryClientUpdate)); +} diff --git a/clients/openframe-client/src/services/mod.rs b/clients/openframe-client/src/services/mod.rs index d0a33fa855..277d25ad07 100644 --- a/clients/openframe-client/src/services/mod.rs +++ b/clients/openframe-client/src/services/mod.rs @@ -6,6 +6,7 @@ pub mod deregistration_service; pub mod device_data_fetcher; pub mod encryption_service; pub mod execution_service; +pub mod feature_flags_service; pub mod github_download_service; pub mod hostname_report_publisher; pub mod initial_authentication_processor; diff --git a/clients/openframe-client/src/services/nats_connection_manager.rs b/clients/openframe-client/src/services/nats_connection_manager.rs index 549ceb1dac..53dd87db52 100644 --- a/clients/openframe-client/src/services/nats_connection_manager.rs +++ b/clients/openframe-client/src/services/nats_connection_manager.rs @@ -330,6 +330,11 @@ impl NatsConnectionManager { Ok(format!("{}/ws/nats?authorization={}", host, token)) } + #[cfg(test)] + pub(crate) fn use_client_for_tests(&self, client: Client) { + self.client.send_replace(Some(Arc::new(client))); + } + pub async fn get_client(&self) -> Result> { self.client .borrow() diff --git a/clients/openframe-client/src/services/result_outbox_run_manager.rs b/clients/openframe-client/src/services/result_outbox_run_manager.rs index 3fcde33a62..7155495149 100644 --- a/clients/openframe-client/src/services/result_outbox_run_manager.rs +++ b/clients/openframe-client/src/services/result_outbox_run_manager.rs @@ -5,7 +5,7 @@ use tokio::time::{interval, timeout, Duration}; use tracing::{error, info, warn}; use crate::config::update_config::{FLUSH_PUBLISH_TIMEOUT_SECS, OUTBOX_FLUSH_INTERVAL_SECS}; -use crate::models::is_ack_subject; +use crate::models::requires_stream_ack; use crate::services::result_store::{payload_limit, ResultPublisher, ResultStore}; pub struct ResultOutboxRunManager { @@ -76,7 +76,7 @@ impl ResultOutboxRunManager

{ continue; } - let outcome = if is_ack_subject(&subject) { + let outcome = if requires_stream_ack(&subject) { timeout(publish_timeout, publisher.publish_acked(&subject, &bytes)).await } else { timeout(publish_timeout, publisher.publish_raw(&subject, &bytes)).await diff --git a/clients/openframe-client/src/test_support.rs b/clients/openframe-client/src/test_support.rs new file mode 100644 index 0000000000..2fb6882f9e --- /dev/null +++ b/clients/openframe-client/src/test_support.rs @@ -0,0 +1,68 @@ +use std::net::{TcpListener, TcpStream}; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +use tempfile::TempDir; + +const READY_BOUND: Duration = Duration::from_secs(10); + +pub(crate) struct NatsServer { + child: Child, + port: u16, + _store: TempDir, +} + +impl NatsServer { + pub(crate) fn start() -> Option { + let binary = std::env::var("NATS_SERVER_BIN").unwrap_or_else(|_| "nats-server".to_string()); + let port = TcpListener::bind("127.0.0.1:0") + .and_then(|listener| listener.local_addr()) + .expect("no free local port") + .port(); + let store = tempfile::tempdir().unwrap(); + + let spawned = Command::new(&binary) + .args(["-a", "127.0.0.1", "-p", &port.to_string(), "-js", "-sd"]) + .arg(store.path()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .spawn(); + let child = match spawned { + Ok(child) => child, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + if std::env::var("REQUIRE_NATS_SERVER").as_deref() == Ok("1") { + panic!("REQUIRE_NATS_SERVER=1 but `{binary}` is not installed"); + } + eprintln!("skipping: `{binary}` is not installed (brew install nats-server)"); + return None; + } + Err(e) => panic!("failed to start `{binary}`: {e}"), + }; + + let server = Self { + child, + port, + _store: store, + }; + let started = Instant::now(); + while TcpStream::connect(("127.0.0.1", port)).is_err() { + assert!( + started.elapsed() < READY_BOUND, + "nats-server did not start listening on {port}" + ); + std::thread::sleep(Duration::from_millis(20)); + } + Some(server) + } + + pub(crate) fn url(&self) -> String { + format!("nats://127.0.0.1:{}", self.port) + } +} + +impl Drop for NatsServer { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} From 4e267c6f27820517ce9bd3e2b868138dc377392e Mon Sep 17 00:00:00 2001 From: denys-gif Date: Tue, 13 Oct 2026 20:22:28 +0100 Subject: [PATCH 2/5] feat: min ver for feature flags --- .../docs/NATS_INTEGRATION_TESTS.md | 2 +- .../src/config/update_config.rs | 3 +- clients/openframe-client/src/lib.rs | 38 +- .../src/listener/command/handlers_tests.rs | 80 ++-- .../src/listener/command/integration_tests.rs | 150 +++++++- .../src/models/feature_flags.rs | 43 ++- .../src/models/feature_flags_tests.rs | 209 +++++++++++ clients/openframe-client/src/models/mod.rs | 4 +- .../src/platform/detached_process.rs | 11 +- .../src/platform/uninstall.rs | 4 +- .../src/services/feature_flags_service.rs | 32 +- .../services/feature_flags_service_tests.rs | 346 +++++++++++++----- 12 files changed, 739 insertions(+), 183 deletions(-) create mode 100644 clients/openframe-client/src/models/feature_flags_tests.rs diff --git a/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md b/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md index 702e78e778..c5f2430f84 100644 --- a/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md +++ b/clients/openframe-client/docs/NATS_INTEGRATION_TESTS.md @@ -209,4 +209,4 @@ and the subject in `send()`, and keep the helpers. environment is not caught here — check the NATS config when you add a subject. - Process-level effects (the service manager restarting the client, `exit`) are out of reach; tests stop at the signal the code raises. -- Each test starts its own server; the 7 tests of the example suite take about a second in total. Fine for dozens of tests, not thousands. +- Each test starts its own server; the example suite (11 tests) takes about a second in total. Fine for dozens of tests, not thousands. diff --git a/clients/openframe-client/src/config/update_config.rs b/clients/openframe-client/src/config/update_config.rs index 2b92afaaa1..ee50f7363e 100644 --- a/clients/openframe-client/src/config/update_config.rs +++ b/clients/openframe-client/src/config/update_config.rs @@ -20,7 +20,8 @@ pub const RECONNECTION_DELAY_MS: u64 = 5000; // 5 seconds pub const SELF_RESTART_EXIT_CODE: i32 = 75; pub const SELF_RESTART_IDLE_WAIT_SECS: u64 = 600; pub const SELF_RESTART_IDLE_POLL_SECS: u64 = 5; -pub const SELF_RESTART_HANDOFF_SECS: u64 = 120; +pub const SELF_RESTART_HANDOFF_SECS: u64 = 900; +pub const DEV_CLIENT_VERSION: &str = "0.0.0"; // Durable result outbox pub const OUTBOX_MAX_ENTRIES: usize = 500; diff --git a/clients/openframe-client/src/lib.rs b/clients/openframe-client/src/lib.rs index 2ea4fc54f5..29e2dcecf1 100644 --- a/clients/openframe-client/src/lib.rs +++ b/clients/openframe-client/src/lib.rs @@ -378,6 +378,7 @@ impl Client { let feature_flags_service = Arc::new(FeatureFlagsService::load( directory_manager.secured_dir().join("feature_flags.json"), + env!("OPENFRAME_VERSION").to_string(), )); // Initialize tool run manager @@ -751,13 +752,38 @@ impl Client { &crate::service::Service::get_install_location(), cli::RESTART_SERVICE_SUBCOMMAND, ) { - Ok(pid) => { - info!(pid, "Service restart handed to a detached process"); - tokio::time::sleep(Duration::from_secs(SELF_RESTART_HANDOFF_SECS)).await; - warn!( - "Service was not restarted within {}s, exiting instead", - SELF_RESTART_HANDOFF_SECS + Ok(mut child) => { + info!( + pid = child.id(), + "Service restart handed to a detached process" ); + let finished = tokio::time::timeout( + Duration::from_secs(SELF_RESTART_HANDOFF_SECS), + tokio::task::spawn_blocking(move || child.wait()), + ) + .await; + match finished { + Ok(Ok(Ok(status))) => warn!( + %status, + "Detached restart finished but did not stop this process, exiting instead" + ), + Ok(Ok(Err(e))) => { + warn!( + "Cannot wait for the detached restart, exiting instead: {:#}", + e + ) + } + Ok(Err(e)) => { + warn!( + "Cannot wait for the detached restart, exiting instead: {:#}", + e + ) + } + Err(_) => warn!( + "Detached restart still running after {}s, exiting instead", + SELF_RESTART_HANDOFF_SECS + ), + } } Err(e) => warn!( "Cannot hand the service restart off, exiting instead: {:#}", diff --git a/clients/openframe-client/src/listener/command/handlers_tests.rs b/clients/openframe-client/src/listener/command/handlers_tests.rs index d0b2d80b13..c1f0703c45 100644 --- a/clients/openframe-client/src/listener/command/handlers_tests.rs +++ b/clients/openframe-client/src/listener/command/handlers_tests.rs @@ -1,11 +1,14 @@ use super::*; +use serde_json::{json, Value}; use tempfile::TempDir; const MACHINE: &str = "machine-1"; +const VERSION: &str = "1.5.1"; fn handler() -> (FeatureFlagsHandler, TempDir) { let dir = tempfile::tempdir().unwrap(); - let service = FeatureFlagsService::load(dir.path().join("feature_flags.json")); + let service = + FeatureFlagsService::load(dir.path().join("feature_flags.json"), VERSION.to_string()); (FeatureFlagsHandler::new(Arc::new(service)), dir) } @@ -18,12 +21,16 @@ fn delivery(sequence: Option) -> DeliveryRef { } } -fn message(flags: serde_json::Value) -> FeatureFlagsMessage { - serde_json::from_value(serde_json::json!({ "flags": flags })).unwrap() +fn tool_installation(value: bool, min_version: &str) -> FeatureFlagsMessage { + message(json!({ "delivery-tool-installation": { "value": value, "minVersion": min_version } })) +} + +fn message(flags: Value) -> FeatureFlagsMessage { + serde_json::from_value(json!({ "flags": flags })).unwrap() } #[test] -fn subject_is_machine_scoped() { +fn the_subject_is_machine_scoped() { assert_eq!( FeatureFlagsHandler::subject("abc123"), "machine.abc123.feature-flags" @@ -35,11 +42,7 @@ async fn flags_without_a_sequence_fail() { let (handler, _dir) = handler(); let outcome = handler - .run( - MACHINE, - &delivery(None), - message(serde_json::json!({ "delivery-tool-installation": true })), - ) + .run(MACHINE, &delivery(None), tool_installation(true, "1.0.0")) .await; assert!(matches!(outcome, Outcome::Failed(_))); @@ -47,46 +50,47 @@ async fn flags_without_a_sequence_fail() { } #[tokio::test] -async fn a_changed_set_requires_a_restart_and_an_unchanged_one_does_not() { - let (handler, _dir) = handler(); - - let unchanged = handler - .run( - MACHINE, - &delivery(Some(1)), - message(serde_json::json!({ "delivery-tool-installation": false })), - ) - .await; - let changed = handler - .run( - MACHINE, - &delivery(Some(2)), - message(serde_json::json!({ "delivery-tool-installation": true })), - ) - .await; - - assert_eq!(unchanged, Outcome::Done); - assert_eq!(changed, Outcome::RestartRequired); +async fn an_applicable_change_restarts_and_anything_else_is_done() { + let cases = [ + ( + tool_installation(false, "1.0.0"), + Outcome::Done, + "switched off, as running", + ), + ( + tool_installation(true, "2.0.0"), + Outcome::Done, + "for a newer client", + ), + ( + tool_installation(true, "1.0.0"), + Outcome::RestartRequired, + "switched on for this client", + ), + ]; + for (message, expected, why) in cases { + let (handler, _dir) = handler(); + + let outcome = handler.run(MACHINE, &delivery(Some(1)), message).await; + + assert_eq!(outcome, expected, "{why}"); + } } #[tokio::test] -async fn stale_or_repeated_sequence_is_done_without_applying() { +async fn a_stale_or_repeated_sequence_is_done_without_applying() { let (handler, _dir) = handler(); handler .run( MACHINE, &delivery(Some(5)), - message(serde_json::json!({ "delivery-tool-installation": true })), + tool_installation(true, "1.0.0"), ) .await; for stale in [5, 4] { let outcome = handler - .run( - MACHINE, - &delivery(Some(stale)), - message(serde_json::json!({})), - ) + .run(MACHINE, &delivery(Some(stale)), message(json!({}))) .await; assert_eq!(outcome, Outcome::Done, "sequence {stale}"); @@ -102,7 +106,7 @@ async fn a_reregistered_machine_applies_a_lower_sequence() { .run( MACHINE, &delivery(Some(50)), - message(serde_json::json!({ "delivery-tool-installation": true })), + tool_installation(true, "1.0.0"), ) .await; @@ -110,7 +114,7 @@ async fn a_reregistered_machine_applies_a_lower_sequence() { .run( "machine-2", &delivery(Some(3)), - message(serde_json::json!({ "delivery-tool-installation": true })), + tool_installation(true, "1.0.0"), ) .await; diff --git a/clients/openframe-client/src/listener/command/integration_tests.rs b/clients/openframe-client/src/listener/command/integration_tests.rs index b821d2fe70..659d5b4523 100644 --- a/clients/openframe-client/src/listener/command/integration_tests.rs +++ b/clients/openframe-client/src/listener/command/integration_tests.rs @@ -9,7 +9,8 @@ use tokio::time::{timeout, Duration}; use super::{CommandContext, CommandListener, FeatureFlagsHandler}; use crate::clients::AuthClient; -use crate::models::{DeliveryResult, DeliveryResultMessage, FeatureFlagsState}; +use crate::config::update_config::DEV_CLIENT_VERSION; +use crate::models::{DeliveryResult, DeliveryResultMessage, Feature, FeatureFlagsState}; use crate::platform::DirectoryManager; use crate::services::encryption_service::EncryptionService; use crate::services::feature_flags_service::FeatureFlagsService; @@ -30,6 +31,7 @@ const MACHINE: &str = "it-machine"; const BOUND: Duration = Duration::from_secs(10); const QUIET: Duration = Duration::from_millis(500); const PROBE: &str = "probe"; +const VERSION: &str = "1.5.1"; struct Agent { _server: NatsServer, @@ -44,10 +46,15 @@ struct Agent { store: Arc, flush: Arc, flags_path: PathBuf, + version: String, } impl Agent { async fn start() -> Option<(Self, CommandContext)> { + Self::start_as(VERSION).await + } + + async fn start_as(version: &str) -> Option<(Self, CommandContext)> { let server = NatsServer::start()?; let dir = tempfile::tempdir().unwrap(); let directories = DirectoryManager::with_custom_dirs( @@ -114,6 +121,7 @@ impl Agent { let mut agent = Self { _server: server, flags_path: directories.secured_dir().join("feature_flags.json"), + version: version.to_string(), _dir: dir, tester, results, @@ -125,11 +133,11 @@ impl Agent { store, flush, }; - let context = agent.launch().await; + let (context, _service) = agent.launch().await; Some((agent, context)) } - async fn launch(&mut self) -> CommandContext { + async fn launch(&mut self) -> (CommandContext, Arc) { let context = CommandContext::new( self.manager.clone(), self.config.clone(), @@ -139,13 +147,21 @@ impl Agent { self.store.clone(), self.flush.clone(), ); - let service = Arc::new(FeatureFlagsService::load(self.flags_path.clone())); - CommandListener::new(FeatureFlagsHandler::new(service), context.clone()) + let service = Arc::new(FeatureFlagsService::load( + self.flags_path.clone(), + self.version.clone(), + )); + CommandListener::new(FeatureFlagsHandler::new(service.clone()), context.clone()) .start() .await .unwrap(); self.until_listening().await; - context + (context, service) + } + + async fn launch_as(&mut self, version: &str) -> (CommandContext, Arc) { + self.version = version.to_string(); + self.launch().await } async fn until_listening(&mut self) { @@ -222,6 +238,15 @@ impl Agent { } } +fn entry(value: bool, min_version: &str) -> serde_json::Value { + json!({ "value": value, "minVersion": min_version }) +} + +async fn expect_no_restart(context: &CommandContext) { + let requested = timeout(QUIET, context.restart_requested()).await; + assert!(requested.is_err(), "a restart was requested"); +} + fn command( dispatch_id: &str, sequence: Option, @@ -248,7 +273,7 @@ async fn unchanged_flags_are_acked_done_and_persisted() { .send(command( "d1", Some(1), - json!({ "delivery-tool-installation": false }), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), )) .await; @@ -269,7 +294,7 @@ async fn changed_flags_restart_and_the_next_process_reports_done() { .send(command( "d2", Some(2), - json!({ "delivery-tool-installation": true }), + json!({ "delivery-tool-installation": entry(true, "1.0.0") }), )) .await; @@ -307,7 +332,7 @@ async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { .send(command( "d3", Some(3), - json!({ "delivery-tool-installation": true }), + json!({ "delivery-tool-installation": entry(true, "1.0.0") }), )) .await; timeout(BOUND, context.restart_requested()) @@ -318,7 +343,7 @@ async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { .send(command( "d4", Some(4), - json!({ "delivery-tool-installation": false }), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), )) .await; @@ -339,7 +364,7 @@ async fn stale_sequence_is_done_without_applying() { .send(command( "d6", Some(4), - json!({ "delivery-tool-installation": false }), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), )) .await; @@ -380,13 +405,13 @@ async fn a_message_without_delivery_gets_no_answer() { }; agent - .send(json!({ "flags": { "delivery-tool-installation": true } })) + .send(json!({ "flags": { "delivery-tool-installation": entry(true, "1.0.0") } })) .await; agent .send(command( "d8", Some(8), - json!({ "delivery-tool-installation": false }), + json!({ "delivery-tool-installation": entry(false, "1.0.0") }), )) .await; @@ -415,7 +440,7 @@ async fn flags_without_a_sequence_fail() { .send(command( "d9", None, - json!({ "delivery-tool-installation": true }), + json!({ "delivery-tool-installation": entry(true, "1.0.0") }), )) .await; @@ -423,3 +448,100 @@ async fn flags_without_a_sequence_fail() { let failed = agent.expect("d9", DeliveryResult::Failed).await; assert!(failed.error.unwrap().contains("sequence")); } + +#[tokio::test] +async fn a_flag_for_a_newer_client_is_stored_without_a_restart_and_turns_on_after_the_update() { + let Some((mut agent, context)) = Agent::start_as("1.5.9").await else { + return; + }; + + agent + .send(command( + "d20", + Some(20), + json!({ "delivery-tool-installation": entry(true, "1.6.0") }), + )) + .await; + + agent.expect("d20", DeliveryResult::Acked).await; + agent.expect("d20", DeliveryResult::Done).await; + expect_no_restart(&context).await; + assert!(agent + .persisted() + .flags + .contains_key("delivery-tool-installation")); + + let (_updated, service) = agent.launch_as("1.6.0").await; + + assert!(service + .flags() + .is_enabled(Feature::DeliveryToolInstallation)); + agent.expect_silence("d20").await; +} + +#[tokio::test] +async fn a_rollback_below_the_min_version_runs_with_the_flag_off() { + let Some((mut agent, context)) = Agent::start().await else { + return; + }; + agent + .send(command( + "d21", + Some(21), + json!({ "delivery-tool-installation": entry(true, "1.5.0") }), + )) + .await; + timeout(BOUND, context.restart_requested()) + .await + .expect("restart was never requested"); + + let (_rolled_back, service) = agent.launch_as("1.4.0").await; + + assert!(!service + .flags() + .is_enabled(Feature::DeliveryToolInstallation)); +} + +#[tokio::test] +async fn a_dev_build_applies_a_true_flag_whatever_its_min_version() { + let Some((mut agent, context)) = Agent::start_as(DEV_CLIENT_VERSION).await else { + return; + }; + + agent + .send(command( + "d22", + Some(22), + json!({ "delivery-tool-installation": entry(true, "99.0.0") }), + )) + .await; + + agent.expect("d22", DeliveryResult::Acked).await; + timeout(BOUND, context.restart_requested()) + .await + .expect("restart was never requested"); +} + +#[tokio::test] +async fn invalid_flag_entries_are_done_without_a_restart() { + let Some((mut agent, context)) = Agent::start().await else { + return; + }; + + agent + .send(command( + "d23", + Some(23), + json!({ + "delivery-tool-installation": true, + "delivery-client-uninstall": { "value": true }, + "delivery-client-update": { "value": true, "minVersion": "latest" } + }), + )) + .await; + + agent.expect("d23", DeliveryResult::Acked).await; + agent.expect("d23", DeliveryResult::Done).await; + expect_no_restart(&context).await; + assert!(agent.persisted().flags.is_empty()); +} diff --git a/clients/openframe-client/src/models/feature_flags.rs b/clients/openframe-client/src/models/feature_flags.rs index 66b44aa76d..1912ea6a24 100644 --- a/clients/openframe-client/src/models/feature_flags.rs +++ b/clients/openframe-client/src/models/feature_flags.rs @@ -1,9 +1,12 @@ +use std::cmp::Ordering; use std::collections::{BTreeMap, HashSet}; use serde::{Deserialize, Serialize}; use serde_json::Value; use super::DeliveryRef; +use crate::config::update_config::DEV_CLIENT_VERSION; +use crate::utils::version_comparator::{parse_version, precedence}; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum Feature { @@ -46,26 +49,48 @@ pub struct FeatureFlagsState { #[serde(default)] pub machine_id: String, pub sequence: i64, - pub flags: BTreeMap, + pub flags: BTreeMap, #[serde(default, skip_serializing_if = "Option::is_none")] pub awaiting_restart: Option, } +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct FeatureFlag { + pub value: bool, + pub min_version: String, +} + +impl FeatureFlag { + fn effective(&self, running_version: &str, default: bool) -> bool { + if running_version == DEV_CLIENT_VERSION { + return self.value; + } + match ( + parse_version(running_version), + parse_version(&self.min_version), + ) { + (Some(running), Some(min)) if precedence(&running, &min) != Ordering::Less => { + self.value + } + _ => default, + } + } +} + #[derive(Debug, Clone, Default, PartialEq)] pub struct FeatureFlags { enabled: HashSet, } impl FeatureFlags { - pub fn from_flags(flags: &BTreeMap) -> Self { + pub fn resolve(flags: &BTreeMap, running_version: &str) -> Self { Self { enabled: Feature::ALL .into_iter() - .filter(|feature| { - flags - .get(feature.name()) - .copied() - .unwrap_or(feature.default_enabled()) + .filter(|feature| match flags.get(feature.name()) { + Some(flag) => flag.effective(running_version, feature.default_enabled()), + None => feature.default_enabled(), }) .collect(), } @@ -75,3 +100,7 @@ impl FeatureFlags { self.enabled.contains(&feature) } } + +#[cfg(test)] +#[path = "feature_flags_tests.rs"] +mod tests; diff --git a/clients/openframe-client/src/models/feature_flags_tests.rs b/clients/openframe-client/src/models/feature_flags_tests.rs new file mode 100644 index 0000000000..736e0fe7cd --- /dev/null +++ b/clients/openframe-client/src/models/feature_flags_tests.rs @@ -0,0 +1,209 @@ +use super::*; + +const TOOL_INSTALLATION: &str = "delivery-tool-installation"; + +fn flag(value: bool, min_version: &str) -> FeatureFlag { + FeatureFlag { + value, + min_version: min_version.to_string(), + } +} + +fn resolve(entries: &[(&str, FeatureFlag)], running_version: &str) -> FeatureFlags { + let flags = entries + .iter() + .map(|(name, flag)| (name.to_string(), flag.clone())) + .collect(); + FeatureFlags::resolve(&flags, running_version) +} + +#[test] +fn a_flag_applies_only_from_its_min_version() { + let cases = [ + (true, "1.5.1", "1.5.1", true, "same version"), + (true, "1.5.1", "1.5.2", true, "newer patch"), + (true, "1.5.1", "1.6.0", true, "newer minor"), + (true, "1.5.1", "2.0.0", true, "newer major"), + (true, "1.5.1", "1.5.0", false, "older patch"), + (true, "1.5.1", "1.4.9", false, "older minor"), + (true, "1.5.1", "0.9.9", false, "older major"), + ( + false, + "1.5.1", + "2.0.0", + false, + "switched off on a supporting version", + ), + ( + false, + "1.5.1", + "1.0.0", + false, + "switched off below the min version", + ), + ( + true, + "1.5.1", + "1.5.1-rc.1", + false, + "a pre-release is below its release", + ), + ( + true, + "1.5.1-rc.1", + "1.5.1-rc.2", + true, + "a later pre-release", + ), + ( + true, + "1.5.1-rc.2", + "1.5.1-rc.1", + false, + "an earlier pre-release", + ), + ( + true, + "1.5.1-rc.1", + "1.5.1", + true, + "a release is above its pre-release", + ), + ( + true, + "1.5.1", + "1.5.1+build.7", + true, + "build metadata carries no order", + ), + ( + true, + "v1.5.1", + "1.5.1", + true, + "a v prefix on the min version", + ), + ( + true, + " 1.5.1 ", + "1.5.1", + true, + "whitespace around the min version", + ), + ( + true, + "latest", + "2.0.0", + false, + "a min version that is not semver", + ), + (true, "", "2.0.0", false, "an empty min version"), + ( + true, + "1.5.1", + "not-a-version", + false, + "a running version that is not semver", + ), + ]; + for (value, min_version, running, expected, why) in cases { + let flags = resolve(&[(TOOL_INSTALLATION, flag(value, min_version))], running); + + assert_eq!( + flags.is_enabled(Feature::DeliveryToolInstallation), + expected, + "{why}: value={value} minVersion={min_version:?} running={running:?}" + ); + } +} + +#[test] +fn the_dev_build_takes_the_value_whatever_the_min_version() { + let cases = [ + (true, "99.0.0", true), + (true, "latest", true), + (true, "0.0.1", true), + (false, "0.0.1", false), + (false, "99.0.0", false), + ]; + for (value, min_version, expected) in cases { + let flags = resolve( + &[(TOOL_INSTALLATION, flag(value, min_version))], + DEV_CLIENT_VERSION, + ); + + assert_eq!( + flags.is_enabled(Feature::DeliveryToolInstallation), + expected, + "value={value} minVersion={min_version:?}" + ); + } +} + +#[test] +fn a_flag_that_does_not_apply_falls_back_to_the_feature_default() { + let below = flag(false, "2.0.0"); + let above = flag(false, "1.0.0"); + + assert!(below.effective("1.5.1", true)); + assert!(!below.effective("1.5.1", false)); + assert!(!above.effective("1.5.1", true)); +} + +#[test] +fn a_missing_flag_resolves_to_the_default() { + let flags = resolve(&[], "1.5.1"); + + for feature in Feature::ALL { + assert_eq!( + flags.is_enabled(feature), + feature.default_enabled(), + "{}", + feature.name() + ); + } +} + +#[test] +fn each_backend_name_switches_only_its_own_feature() { + for feature in Feature::ALL { + let flags = resolve(&[(feature.name(), flag(true, "1.0.0"))], "1.5.1"); + + for other in Feature::ALL { + assert_eq!( + flags.is_enabled(other), + other == feature, + "{} set, {} checked", + feature.name(), + other.name() + ); + } + } +} + +#[test] +fn unknown_flags_do_not_touch_known_features() { + let with_unknown = resolve( + &[ + ("future-feature", flag(true, "1.0.0")), + (TOOL_INSTALLATION, flag(true, "1.0.0")), + ], + "1.5.1", + ); + let without = resolve(&[(TOOL_INSTALLATION, flag(true, "1.0.0"))], "1.5.1"); + + assert_eq!(with_unknown, without); +} + +#[test] +fn flag_entries_use_the_backend_field_names() { + let parsed: FeatureFlag = + serde_json::from_value(serde_json::json!({ "value": true, "minVersion": "1.5.1" })) + .unwrap(); + + assert_eq!(parsed, flag(true, "1.5.1")); + assert_eq!( + serde_json::to_value(&parsed).unwrap(), + serde_json::json!({ "value": true, "minVersion": "1.5.1" }) + ); +} diff --git a/clients/openframe-client/src/models/mod.rs b/clients/openframe-client/src/models/mod.rs index fc68f4c9c7..6e22ac7240 100644 --- a/clients/openframe-client/src/models/mod.rs +++ b/clients/openframe-client/src/models/mod.rs @@ -43,7 +43,9 @@ pub use delivery::{ }; pub use device_tag::DeviceTag; pub use download_configuration::{DownloadConfiguration, InstallationType}; -pub use feature_flags::{Feature, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState}; +pub use feature_flags::{ + Feature, FeatureFlag, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState, +}; pub use hostname_report_message::HostnameReportMessage; pub use initial_configuration::InitialConfiguration; pub use installed_agent_message::InstalledAgentMessage; diff --git a/clients/openframe-client/src/platform/detached_process.rs b/clients/openframe-client/src/platform/detached_process.rs index 31af1932b0..e08593a188 100644 --- a/clients/openframe-client/src/platform/detached_process.rs +++ b/clients/openframe-client/src/platform/detached_process.rs @@ -1,9 +1,10 @@ use anyhow::Result; use std::path::Path; +use std::process::Child; /// Spawn ` ` detached, so it survives the service manager stopping /// our own `com.openframe.client` service. macOS + Windows only. -pub fn spawn_detached(install_path: &Path, subcommand: &str) -> Result { +pub fn spawn_detached(install_path: &Path, subcommand: &str) -> Result { #[cfg(target_os = "macos")] { spawn_detached_macos(install_path, subcommand) @@ -22,7 +23,7 @@ pub fn spawn_detached(install_path: &Path, subcommand: &str) -> Result { /// Detached child in its own session (setsid) so it outlives launchd stopping our own /// `com.openframe.client` service — no launchd job/plist needed. #[cfg(target_os = "macos")] -fn spawn_detached_macos(install_path: &Path, subcommand: &str) -> Result { +fn spawn_detached_macos(install_path: &Path, subcommand: &str) -> Result { use anyhow::Context; use std::os::unix::process::CommandExt; use std::process::{Command, Stdio}; @@ -45,12 +46,12 @@ fn spawn_detached_macos(install_path: &Path, subcommand: &str) -> Result { } .with_context(|| format!("Failed to spawn detached `{subcommand}` process"))?; - Ok(child.id()) + Ok(child) } /// Detached child (inherits LocalSystem) so it survives the SCM stopping our service. #[cfg(target_os = "windows")] -fn spawn_detached_windows(install_path: &Path, subcommand: &str) -> Result { +fn spawn_detached_windows(install_path: &Path, subcommand: &str) -> Result { use anyhow::Context; use std::os::windows::process::CommandExt; use std::process::Command; @@ -64,5 +65,5 @@ fn spawn_detached_windows(install_path: &Path, subcommand: &str) -> Result .spawn() .with_context(|| format!("Failed to spawn detached `{subcommand}` process"))?; - Ok(child.id()) + Ok(child) } diff --git a/clients/openframe-client/src/platform/uninstall.rs b/clients/openframe-client/src/platform/uninstall.rs index 9285e1c92d..990cc16c03 100644 --- a/clients/openframe-client/src/platform/uninstall.rs +++ b/clients/openframe-client/src/platform/uninstall.rs @@ -93,9 +93,9 @@ fn remove_update_temp_files(install_path: &Path) { /// Self-uninstall stops the `com.openframe.client` service (our own process), so it must run /// out-of-process. macOS + Windows only — Linux client self-uninstall is unsupported. pub fn spawn_detached_uninstall(install_path: &Path) -> Result<()> { - let pid = + let child = crate::platform::detached_process::spawn_detached(install_path, UNINSTALL_SUBCOMMAND)?; - info!("Self-uninstall process launched (PID: {})", pid); + info!("Self-uninstall process launched (PID: {})", child.id()); Ok(()) } diff --git a/clients/openframe-client/src/services/feature_flags_service.rs b/clients/openframe-client/src/services/feature_flags_service.rs index 3b6f21ed48..ac92c2d8e2 100644 --- a/clients/openframe-client/src/services/feature_flags_service.rs +++ b/clients/openframe-client/src/services/feature_flags_service.rs @@ -3,14 +3,17 @@ use std::path::PathBuf; use std::sync::{Mutex, PoisonError}; use anyhow::{Context, Result}; -use serde_json::Value; use tracing::{info, warn}; -use crate::models::{DeliveryRef, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState}; +use crate::models::{ + DeliveryRef, FeatureFlag, FeatureFlags, FeatureFlagsMessage, FeatureFlagsState, +}; use crate::utils::fs::atomic_write; +use crate::utils::version_comparator::parse_version; pub struct FeatureFlagsService { path: PathBuf, + version: String, running: FeatureFlags, state: Mutex, } @@ -22,7 +25,7 @@ pub enum FlagsApplied { } impl FeatureFlagsService { - pub fn load(path: PathBuf) -> Self { + pub fn load(path: PathBuf, version: String) -> Self { let state = match std::fs::read(&path) { Ok(bytes) => match serde_json::from_slice::(&bytes) { Ok(state) => state, @@ -37,10 +40,17 @@ impl FeatureFlagsService { FeatureFlagsState::default() } }; - let running = FeatureFlags::from_flags(&state.flags); - info!(sequence = state.sequence, flags = ?state.flags, "Feature flags loaded"); + let running = FeatureFlags::resolve(&state.flags, &version); + info!( + version, + sequence = state.sequence, + flags = ?state.flags, + effective = ?running, + "Feature flags loaded" + ); Self { path, + version, running, state: Mutex::new(state), } @@ -72,16 +82,16 @@ impl FeatureFlagsService { ) -> Result { let mut flags = BTreeMap::new(); for (name, value) in message.flags { - match value { - Value::Bool(enabled) => { - flags.insert(name, enabled); + match serde_json::from_value::(value.clone()) { + Ok(flag) if parse_version(&flag.min_version).is_some() => { + flags.insert(name, flag); } - other => { - warn!(flag = %name, value = %other, "Ignoring feature flag with a non-boolean value") + _ => { + warn!(flag = %name, %value, "Ignoring feature flag without a boolean value and a semver minVersion") } } } - let applied = if FeatureFlags::from_flags(&flags) == self.running { + let applied = if FeatureFlags::resolve(&flags, &self.version) == self.running { FlagsApplied::Unchanged } else { FlagsApplied::RestartRequired diff --git a/clients/openframe-client/src/services/feature_flags_service_tests.rs b/clients/openframe-client/src/services/feature_flags_service_tests.rs index a3f749b1d5..9d1baa7538 100644 --- a/clients/openframe-client/src/services/feature_flags_service_tests.rs +++ b/clients/openframe-client/src/services/feature_flags_service_tests.rs @@ -1,14 +1,22 @@ use super::*; +use crate::config::update_config::DEV_CLIENT_VERSION; use crate::models::Feature; +use serde_json::{json, Value}; use tempfile::TempDir; const MACHINE: &str = "machine-1"; +const VERSION: &str = "1.5.1"; +const TOOL_INSTALLATION: &str = "delivery-tool-installation"; fn flags_path() -> (PathBuf, TempDir) { let dir = tempfile::tempdir().unwrap(); (dir.path().join("feature_flags.json"), dir) } +fn load(path: &std::path::Path, version: &str) -> FeatureFlagsService { + FeatureFlagsService::load(path.to_path_buf(), version.to_string()) +} + fn delivery(sequence: i64) -> DeliveryRef { DeliveryRef { delivery_type: "CLIENT_FEATURE_FLAGS".to_string(), @@ -18,130 +26,281 @@ fn delivery(sequence: i64) -> DeliveryRef { } } -fn message(flags: serde_json::Value) -> FeatureFlagsMessage { - serde_json::from_value(serde_json::json!({ "flags": flags })).unwrap() +fn entry(value: bool, min_version: &str) -> Value { + json!({ "value": value, "minVersion": min_version }) } -fn apply(service: &FeatureFlagsService, sequence: i64, flags: serde_json::Value) -> FlagsApplied { +fn apply(service: &FeatureFlagsService, sequence: i64, flags: Value) -> FlagsApplied { + let message: FeatureFlagsMessage = serde_json::from_value(json!({ "flags": flags })).unwrap(); service - .apply(MACHINE, &delivery(sequence), sequence, message(flags)) + .apply(MACHINE, &delivery(sequence), sequence, message) .unwrap() } -fn tool_installation_delivery(service: &FeatureFlagsService) -> bool { +fn tool_installation(service: &FeatureFlagsService) -> bool { service .flags() .is_enabled(Feature::DeliveryToolInstallation) } +fn persisted(path: &std::path::Path) -> FeatureFlagsState { + serde_json::from_slice(&std::fs::read(path).unwrap()).unwrap() +} + +fn running_with_tool_installation_on(path: &std::path::Path) -> FeatureFlagsService { + let first = load(path, VERSION); + apply( + &first, + 10, + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ); + let running = load(path, VERSION); + running.restart_reported().unwrap(); + assert!(tool_installation(&running)); + running +} + #[test] -fn missing_file_means_every_feature_off() { +fn a_missing_file_means_every_feature_off() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path); + let service = load(&path, VERSION); - assert!(!tool_installation_delivery(&service)); + assert!(!tool_installation(&service)); assert_eq!(service.applied_sequence(MACHINE), 0); assert_eq!(service.awaiting_restart(), None); } #[test] -fn corrupt_file_means_every_feature_off() { - let (path, _dir) = flags_path(); - std::fs::write(&path, b"{ not json").unwrap(); +fn an_unreadable_file_means_every_feature_off() { + for contents in [ + &b"{ not json"[..], + br#"{ "sequence": 9, "flags": { "delivery-tool-installation": true } }"#, + br#"{ "sequence": 9 }"#, + ] { + let (path, _dir) = flags_path(); + std::fs::write(&path, contents).unwrap(); - let service = FeatureFlagsService::load(path); + let service = load(&path, VERSION); - assert!(!tool_installation_delivery(&service)); - assert_eq!(service.applied_sequence(MACHINE), 0); + assert!( + !tool_installation(&service), + "{}", + String::from_utf8_lossy(contents) + ); + assert_eq!(service.applied_sequence(MACHINE), 0); + } } #[test] -fn enabling_a_feature_needs_a_restart_and_takes_effect_on_next_start() { +fn turning_on_an_applicable_flag_restarts_and_takes_effect_in_the_next_process() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path.clone()); + let service = load(&path, VERSION); let applied = apply( &service, 7, - serde_json::json!({ "delivery-tool-installation": true }), + json!({ TOOL_INSTALLATION: entry(true, "1.5.1") }), ); assert_eq!(applied, FlagsApplied::RestartRequired); - assert!(!tool_installation_delivery(&service)); + assert!(!tool_installation(&service)); assert_eq!(service.awaiting_restart(), Some(delivery(7))); - let restarted = FeatureFlagsService::load(path); - assert!(tool_installation_delivery(&restarted)); + let restarted = load(&path, VERSION); + assert!(tool_installation(&restarted)); assert_eq!(restarted.applied_sequence(MACHINE), 7); assert_eq!(restarted.awaiting_restart(), Some(delivery(7))); } #[test] -fn reported_restart_is_not_reported_again() { +fn turning_off_a_running_flag_restarts() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path.clone()); - apply( + let service = running_with_tool_installation_on(&path); + + let applied = apply( &service, - 7, - serde_json::json!({ "delivery-tool-installation": true }), + 11, + json!({ TOOL_INSTALLATION: entry(false, "1.0.0") }), ); - let restarted = FeatureFlagsService::load(path.clone()); - restarted.restart_reported().unwrap(); - - assert_eq!(restarted.awaiting_restart(), None); - let next_start = FeatureFlagsService::load(path); - assert_eq!(next_start.awaiting_restart(), None); - assert!(tool_installation_delivery(&next_start)); + assert_eq!(applied, FlagsApplied::RestartRequired); + assert!(!tool_installation(&load(&path, VERSION))); } #[test] -fn absent_false_and_non_boolean_all_resolve_to_the_default() { - for flags in [ - serde_json::json!({}), - serde_json::json!({ "delivery-tool-installation": false }), - serde_json::json!({ "delivery-tool-installation": "true" }), - serde_json::json!({ "delivery-tool-installation": 1 }), - serde_json::json!({ "delivery-tool-installation": null }), - ] { +fn changes_that_leave_the_effective_set_alone_do_not_restart() { + let cases = [ + ( + "a flag for a newer client", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "delivery-client-uninstall": entry(true, "2.0.0") }), + ), + ( + "the min version moved but still applies", + json!({ TOOL_INSTALLATION: entry(true, "1.5.1") }), + ), + ( + "an unknown flag added", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "future-feature": entry(true, "1.0.0") }), + ), + ( + "an invalid entry for another feature", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "delivery-client-update": true }), + ), + ( + "a switched-off flag for a newer client", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0"), "delivery-client-update": entry(false, "9.0.0") }), + ), + ( + "the same set with a newer sequence", + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), + ), + ]; + for (why, flags) in cases { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path); + let service = running_with_tool_installation_on(&path); - let applied = apply(&service, 1, flags.clone()); + let applied = apply(&service, 11, flags); - assert_eq!(applied, FlagsApplied::Unchanged, "flags: {flags}"); - assert_eq!(service.awaiting_restart(), None, "flags: {flags}"); + assert_eq!(applied, FlagsApplied::Unchanged, "{why}"); + assert_eq!(service.awaiting_restart(), None, "{why}"); + assert_eq!(service.applied_sequence(MACHINE), 11, "{why}"); } } #[test] -fn unknown_flags_are_kept_on_disk_but_change_nothing() { +fn a_flag_for_a_newer_client_is_stored_without_a_restart() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path.clone()); + let service = load(&path, VERSION); - let applied = apply(&service, 3, serde_json::json!({ "FUTURE_FEATURE": true })); + let applied = apply( + &service, + 3, + json!({ TOOL_INSTALLATION: entry(true, "1.6.0") }), + ); assert_eq!(applied, FlagsApplied::Unchanged); - let persisted: FeatureFlagsState = - serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!(persisted.flags.get("FUTURE_FEATURE"), Some(&true)); + assert_eq!(service.awaiting_restart(), None); + assert_eq!( + persisted(&path).flags.get(TOOL_INSTALLATION), + Some(&FeatureFlag { + value: true, + min_version: "1.6.0".to_string() + }) + ); +} + +#[test] +fn the_stored_flag_turns_on_when_the_client_is_updated() { + let (path, _dir) = flags_path(); + let before_update = load(&path, "1.5.9"); + apply( + &before_update, + 3, + json!({ TOOL_INSTALLATION: entry(true, "1.6.0") }), + ); + + let after_update = load(&path, "1.6.0"); + + assert!(tool_installation(&after_update)); + assert_eq!(after_update.awaiting_restart(), None); +} + +#[test] +fn a_rollback_below_the_min_version_turns_the_flag_off() { + let (path, _dir) = flags_path(); + running_with_tool_installation_on(&path); + + let rolled_back = load(&path, "0.9.0"); + + assert!(!tool_installation(&rolled_back)); +} + +#[test] +fn the_dev_build_turns_on_a_flag_whatever_its_min_version() { + let (path, _dir) = flags_path(); + let service = load(&path, DEV_CLIENT_VERSION); + + let applied = apply( + &service, + 1, + json!({ TOOL_INSTALLATION: entry(true, "99.0.0") }), + ); + + assert_eq!(applied, FlagsApplied::RestartRequired); + assert!(tool_installation(&load(&path, DEV_CLIENT_VERSION))); + assert!(!tool_installation(&load(&path, VERSION))); +} + +#[test] +fn invalid_entries_are_dropped_and_resolve_to_the_default() { + let cases = [ + ("a plain boolean", json!(true)), + ("no minVersion", json!({ "value": true })), + ("no value", json!({ "minVersion": "1.0.0" })), + ( + "a string value", + json!({ "value": "true", "minVersion": "1.0.0" }), + ), + ( + "a numeric minVersion", + json!({ "value": true, "minVersion": 1 }), + ), + ( + "a non-semver minVersion", + json!({ "value": true, "minVersion": "latest" }), + ), + ("null", json!(null)), + ("a string", json!("on")), + ]; + for (why, invalid) in cases { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + let applied = apply( + &service, + 1, + json!({ TOOL_INSTALLATION: invalid, "delivery-client-uninstall": entry(false, "1.0.0") }), + ); + + assert_eq!(applied, FlagsApplied::Unchanged, "{why}"); + let stored = persisted(&path); + assert_eq!(stored.flags.get(TOOL_INSTALLATION), None, "{why}"); + assert!( + stored.flags.contains_key("delivery-client-uninstall"), + "{why}" + ); + } +} + +#[test] +fn unknown_flags_are_kept_on_disk() { + let (path, _dir) = flags_path(); + let service = load(&path, VERSION); + + apply( + &service, + 3, + json!({ "future-feature": entry(true, "1.0.0") }), + ); + + assert!(persisted(&path).flags.contains_key("future-feature")); } #[test] fn flipping_back_before_the_restart_cancels_it() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path); + let service = load(&path, VERSION); apply( &service, 1, - serde_json::json!({ "delivery-tool-installation": true }), + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), ); let applied = apply( &service, 2, - serde_json::json!({ "delivery-tool-installation": false }), + json!({ TOOL_INSTALLATION: entry(false, "1.0.0") }), ); assert_eq!(applied, FlagsApplied::Unchanged); @@ -150,49 +309,39 @@ fn flipping_back_before_the_restart_cancels_it() { } #[test] -fn persisted_state_matches_the_documented_shape() { +fn a_reported_restart_is_not_reported_again() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path.clone()); - + let service = load(&path, VERSION); apply( &service, - 42, - serde_json::json!({ "delivery-tool-installation": true }), + 7, + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), ); - let persisted: serde_json::Value = - serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); - assert_eq!( - persisted, - serde_json::json!({ - "machine_id": MACHINE, - "sequence": 42, - "flags": { "delivery-tool-installation": true }, - "awaiting_restart": { - "type": "CLIENT_FEATURE_FLAGS", - "targetId": "tenant", - "dispatchId": "dispatch-42", - "sequence": 42 - } - }) - ); + let restarted = load(&path, VERSION); + restarted.restart_reported().unwrap(); + + assert_eq!(restarted.awaiting_restart(), None); + let next_start = load(&path, VERSION); + assert_eq!(next_start.awaiting_restart(), None); + assert!(tool_installation(&next_start)); } #[test] fn a_new_registration_accepts_any_sequence_and_keeps_flags_until_then() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path.clone()); + let service = load(&path, VERSION); apply( &service, 50, - serde_json::json!({ "delivery-tool-installation": true }), + json!({ TOOL_INSTALLATION: entry(true, "1.0.0") }), ); - let restarted = FeatureFlagsService::load(path); + let restarted = load(&path, VERSION); assert_eq!(restarted.applied_sequence(MACHINE), 50); assert_eq!(restarted.applied_sequence("machine-2"), 0); - assert!(tool_installation_delivery(&restarted)); + assert!(tool_installation(&restarted)); } #[test] @@ -200,37 +349,40 @@ fn a_file_without_machine_id_accepts_the_first_set() { let (path, _dir) = flags_path(); std::fs::write( &path, - br#"{ "sequence": 9, "flags": { "delivery-tool-installation": true } }"#, + br#"{ "sequence": 9, "flags": { "delivery-tool-installation": { "value": true, "minVersion": "1.0.0" } } }"#, ) .unwrap(); - let service = FeatureFlagsService::load(path); + let service = load(&path, VERSION); - assert!(tool_installation_delivery(&service)); + assert!(tool_installation(&service)); assert_eq!(service.applied_sequence(MACHINE), 0); } #[test] -fn the_backend_flag_names_map_to_their_features() { +fn the_persisted_state_matches_the_documented_shape() { let (path, _dir) = flags_path(); - let service = FeatureFlagsService::load(path.clone()); + let service = load(&path, VERSION); apply( &service, - 1, - serde_json::json!({ - "delivery-tool-installation": true, - "delivery-client-uninstall": true, - "delivery-client-update": false - }), + 42, + json!({ TOOL_INSTALLATION: entry(true, "1.5.1") }), ); - let restarted = FeatureFlagsService::load(path); - assert!(restarted - .flags() - .is_enabled(Feature::DeliveryToolInstallation)); - assert!(restarted - .flags() - .is_enabled(Feature::DeliveryClientUninstall)); - assert!(!restarted.flags().is_enabled(Feature::DeliveryClientUpdate)); + let stored: Value = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap(); + assert_eq!( + stored, + json!({ + "machine_id": MACHINE, + "sequence": 42, + "flags": { TOOL_INSTALLATION: { "value": true, "minVersion": "1.5.1" } }, + "awaiting_restart": { + "type": "CLIENT_FEATURE_FLAGS", + "targetId": "tenant", + "dispatchId": "dispatch-42", + "sequence": 42 + } + }) + ); } From 784459dc5ab2d48d2e275d504c7cb5b8b52c1710 Mon Sep 17 00:00:00 2001 From: denys-gif Date: Wed, 14 Oct 2026 13:03:10 +0100 Subject: [PATCH 3/5] chore: dev ver --- clients/openframe-client/src/config/update_config.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/clients/openframe-client/src/config/update_config.rs b/clients/openframe-client/src/config/update_config.rs index ee50f7363e..78ff5a91dd 100644 --- a/clients/openframe-client/src/config/update_config.rs +++ b/clients/openframe-client/src/config/update_config.rs @@ -21,7 +21,7 @@ pub const SELF_RESTART_EXIT_CODE: i32 = 75; pub const SELF_RESTART_IDLE_WAIT_SECS: u64 = 600; pub const SELF_RESTART_IDLE_POLL_SECS: u64 = 5; pub const SELF_RESTART_HANDOFF_SECS: u64 = 900; -pub const DEV_CLIENT_VERSION: &str = "0.0.0"; +pub const DEV_CLIENT_VERSION: &str = "0.0.0-dev"; // Durable result outbox pub const OUTBOX_MAX_ENTRIES: usize = 500; From c0487d5c73017e86228a80add792a8e715582b3c Mon Sep 17 00:00:00 2001 From: denys-gif Date: Wed, 14 Oct 2026 16:19:48 +0100 Subject: [PATCH 4/5] fix: flag field name, restart race, startup order --- clients/openframe-client/src/lib.rs | 12 ++++++------ .../src/listener/command/handlers_tests.rs | 4 +++- .../src/listener/command/integration_tests.rs | 6 +++--- .../src/listener/command/restart.rs | 3 ++- .../openframe-client/src/models/feature_flags.rs | 6 +++--- .../src/models/feature_flags_tests.rs | 6 +++--- .../src/services/feature_flags_service.rs | 2 +- .../src/services/feature_flags_service_tests.rs | 16 ++++++++-------- 8 files changed, 29 insertions(+), 26 deletions(-) diff --git a/clients/openframe-client/src/lib.rs b/clients/openframe-client/src/lib.rs index 29e2dcecf1..0370b0b592 100644 --- a/clients/openframe-client/src/lib.rs +++ b/clients/openframe-client/src/lib.rs @@ -744,6 +744,12 @@ impl Client { self.nats_connection_manager.start_connection_watchdog(); + // Handle any pending update from previous run (after NATS is connected) + if let Err(e) = self.update_handler_service.handle_pending_update().await { + error!("Failed to handle pending update: {:#}", e); + // Continue startup even if update handling fails - don't block the client + } + self.feature_flags_listener.start().await?; let command_context = self.command_context.clone(); tokio::spawn(async move { @@ -797,12 +803,6 @@ impl Client { std::process::exit(SELF_RESTART_EXIT_CODE); }); - // Handle any pending update from previous run (after NATS is connected) - if let Err(e) = self.update_handler_service.handle_pending_update().await { - error!("Failed to handle pending update: {:#}", e); - // Continue startup even if update handling fails - don't block the client - } - // Start machine heartbeat run manager self.machine_heartbeat_run_manager.start(); self.machine_timezone_run_manager.start(); diff --git a/clients/openframe-client/src/listener/command/handlers_tests.rs b/clients/openframe-client/src/listener/command/handlers_tests.rs index c1f0703c45..50a5fb9bb3 100644 --- a/clients/openframe-client/src/listener/command/handlers_tests.rs +++ b/clients/openframe-client/src/listener/command/handlers_tests.rs @@ -22,7 +22,9 @@ fn delivery(sequence: Option) -> DeliveryRef { } fn tool_installation(value: bool, min_version: &str) -> FeatureFlagsMessage { - message(json!({ "delivery-tool-installation": { "value": value, "minVersion": min_version } })) + message( + json!({ "delivery-tool-installation": { "enabled": value, "minVersion": min_version } }), + ) } fn message(flags: Value) -> FeatureFlagsMessage { diff --git a/clients/openframe-client/src/listener/command/integration_tests.rs b/clients/openframe-client/src/listener/command/integration_tests.rs index 659d5b4523..f5960c34b0 100644 --- a/clients/openframe-client/src/listener/command/integration_tests.rs +++ b/clients/openframe-client/src/listener/command/integration_tests.rs @@ -239,7 +239,7 @@ impl Agent { } fn entry(value: bool, min_version: &str) -> serde_json::Value { - json!({ "value": value, "minVersion": min_version }) + json!({ "enabled": value, "minVersion": min_version }) } async fn expect_no_restart(context: &CommandContext) { @@ -534,8 +534,8 @@ async fn invalid_flag_entries_are_done_without_a_restart() { Some(23), json!({ "delivery-tool-installation": true, - "delivery-client-uninstall": { "value": true }, - "delivery-client-update": { "value": true, "minVersion": "latest" } + "delivery-client-uninstall": { "enabled": true }, + "delivery-client-update": { "enabled": true, "minVersion": "latest" } }), )) .await; diff --git a/clients/openframe-client/src/listener/command/restart.rs b/clients/openframe-client/src/listener/command/restart.rs index 2072df2054..cb8924134f 100644 --- a/clients/openframe-client/src/listener/command/restart.rs +++ b/clients/openframe-client/src/listener/command/restart.rs @@ -30,10 +30,11 @@ impl RestartGate { } pub(super) fn admit(&self) -> Option> { + let admitted = self.work.clone().try_read_owned().ok()?; if self.closing.load(Ordering::Acquire) { return None; } - self.work.clone().try_read_owned().ok() + Some(admitted) } pub(super) async fn drain_and_request(&self) { diff --git a/clients/openframe-client/src/models/feature_flags.rs b/clients/openframe-client/src/models/feature_flags.rs index 1912ea6a24..1d230351fa 100644 --- a/clients/openframe-client/src/models/feature_flags.rs +++ b/clients/openframe-client/src/models/feature_flags.rs @@ -57,21 +57,21 @@ pub struct FeatureFlagsState { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] pub struct FeatureFlag { - pub value: bool, + pub enabled: bool, pub min_version: String, } impl FeatureFlag { fn effective(&self, running_version: &str, default: bool) -> bool { if running_version == DEV_CLIENT_VERSION { - return self.value; + return self.enabled; } match ( parse_version(running_version), parse_version(&self.min_version), ) { (Some(running), Some(min)) if precedence(&running, &min) != Ordering::Less => { - self.value + self.enabled } _ => default, } diff --git a/clients/openframe-client/src/models/feature_flags_tests.rs b/clients/openframe-client/src/models/feature_flags_tests.rs index 736e0fe7cd..36a6820548 100644 --- a/clients/openframe-client/src/models/feature_flags_tests.rs +++ b/clients/openframe-client/src/models/feature_flags_tests.rs @@ -4,7 +4,7 @@ const TOOL_INSTALLATION: &str = "delivery-tool-installation"; fn flag(value: bool, min_version: &str) -> FeatureFlag { FeatureFlag { - value, + enabled: value, min_version: min_version.to_string(), } } @@ -198,12 +198,12 @@ fn unknown_flags_do_not_touch_known_features() { #[test] fn flag_entries_use_the_backend_field_names() { let parsed: FeatureFlag = - serde_json::from_value(serde_json::json!({ "value": true, "minVersion": "1.5.1" })) + serde_json::from_value(serde_json::json!({ "enabled": true, "minVersion": "1.5.1" })) .unwrap(); assert_eq!(parsed, flag(true, "1.5.1")); assert_eq!( serde_json::to_value(&parsed).unwrap(), - serde_json::json!({ "value": true, "minVersion": "1.5.1" }) + serde_json::json!({ "enabled": true, "minVersion": "1.5.1" }) ); } diff --git a/clients/openframe-client/src/services/feature_flags_service.rs b/clients/openframe-client/src/services/feature_flags_service.rs index ac92c2d8e2..3406d85df8 100644 --- a/clients/openframe-client/src/services/feature_flags_service.rs +++ b/clients/openframe-client/src/services/feature_flags_service.rs @@ -87,7 +87,7 @@ impl FeatureFlagsService { flags.insert(name, flag); } _ => { - warn!(flag = %name, %value, "Ignoring feature flag without a boolean value and a semver minVersion") + warn!(flag = %name, %value, "Ignoring feature flag without a boolean enabled and a semver minVersion") } } } diff --git a/clients/openframe-client/src/services/feature_flags_service_tests.rs b/clients/openframe-client/src/services/feature_flags_service_tests.rs index 9d1baa7538..1f09834362 100644 --- a/clients/openframe-client/src/services/feature_flags_service_tests.rs +++ b/clients/openframe-client/src/services/feature_flags_service_tests.rs @@ -27,7 +27,7 @@ fn delivery(sequence: i64) -> DeliveryRef { } fn entry(value: bool, min_version: &str) -> Value { - json!({ "value": value, "minVersion": min_version }) + json!({ "enabled": value, "minVersion": min_version }) } fn apply(service: &FeatureFlagsService, sequence: i64, flags: Value) -> FlagsApplied { @@ -184,7 +184,7 @@ fn a_flag_for_a_newer_client_is_stored_without_a_restart() { assert_eq!( persisted(&path).flags.get(TOOL_INSTALLATION), Some(&FeatureFlag { - value: true, + enabled: true, min_version: "1.6.0".to_string() }) ); @@ -236,19 +236,19 @@ fn the_dev_build_turns_on_a_flag_whatever_its_min_version() { fn invalid_entries_are_dropped_and_resolve_to_the_default() { let cases = [ ("a plain boolean", json!(true)), - ("no minVersion", json!({ "value": true })), + ("no minVersion", json!({ "enabled": true })), ("no value", json!({ "minVersion": "1.0.0" })), ( "a string value", - json!({ "value": "true", "minVersion": "1.0.0" }), + json!({ "enabled": "true", "minVersion": "1.0.0" }), ), ( "a numeric minVersion", - json!({ "value": true, "minVersion": 1 }), + json!({ "enabled": true, "minVersion": 1 }), ), ( "a non-semver minVersion", - json!({ "value": true, "minVersion": "latest" }), + json!({ "enabled": true, "minVersion": "latest" }), ), ("null", json!(null)), ("a string", json!("on")), @@ -349,7 +349,7 @@ fn a_file_without_machine_id_accepts_the_first_set() { let (path, _dir) = flags_path(); std::fs::write( &path, - br#"{ "sequence": 9, "flags": { "delivery-tool-installation": { "value": true, "minVersion": "1.0.0" } } }"#, + br#"{ "sequence": 9, "flags": { "delivery-tool-installation": { "enabled": true, "minVersion": "1.0.0" } } }"#, ) .unwrap(); @@ -376,7 +376,7 @@ fn the_persisted_state_matches_the_documented_shape() { json!({ "machine_id": MACHINE, "sequence": 42, - "flags": { TOOL_INSTALLATION: { "value": true, "minVersion": "1.5.1" } }, + "flags": { TOOL_INSTALLATION: { "enabled": true, "minVersion": "1.5.1" } }, "awaiting_restart": { "type": "CLIENT_FEATURE_FLAGS", "targetId": "tenant", From 610add0f1ae0c7a53b8aa29b4263fdecdcb4548c Mon Sep 17 00:00:00 2001 From: denys-gif Date: Wed, 14 Oct 2026 16:47:47 +0100 Subject: [PATCH 5/5] refactor: restart gate --- clients/openframe-client/src/lib.rs | 20 ++- .../src/listener/client_update_gate.rs | 4 +- .../src/listener/command/integration_tests.rs | 65 ++++---- .../src/listener/command/mod.rs | 15 +- .../src/listener/command/restart.rs | 41 +++-- .../src/listener/command/restart_tests.rs | 142 +++++++++++------- .../src/services/execution_service.rs | 40 ++++- .../src/services/execution_service_tests.rs | 23 +++ .../openframe_client_update_service.rs | 10 +- .../openframe-client/src/services/tool_ops.rs | 25 ++- .../src/services/tool_ops_tests.rs | 12 ++ .../src/services/tool_run_manager.rs | 55 ++++--- .../src/services/tool_run_manager_tests.rs | 63 +++++--- 13 files changed, 341 insertions(+), 174 deletions(-) diff --git a/clients/openframe-client/src/lib.rs b/clients/openframe-client/src/lib.rs index 0370b0b592..ee0458d34f 100644 --- a/clients/openframe-client/src/lib.rs +++ b/clients/openframe-client/src/lib.rs @@ -46,7 +46,7 @@ use crate::config::update_config::{ SELF_RESTART_HANDOFF_SECS, }; use crate::listener::client_uninstall_message_listener::ClientUninstallMessageListener; -use crate::listener::command::{CommandContext, CommandListener, FeatureFlagsHandler}; +use crate::listener::command::{CommandContext, CommandListener, FeatureFlagsHandler, RestartGate}; use crate::listener::execution_listener::ExecutionListener; use crate::listener::openframe_client_update_listener::OpenFrameClientUpdateListener; use crate::listener::remote_access_message_listener::RemoteAccessMessageListener; @@ -181,7 +181,7 @@ pub struct Client { client_uninstall_message_listener: ClientUninstallMessageListener, remote_access_message_listener: RemoteAccessMessageListener, feature_flags_listener: CommandListener, - command_context: CommandContext, + restart_gate: RestartGate, command_execution_listener: ExecutionListener, script_execution_listener: ExecutionListener, script_bootstrap_execution_listener: ExecutionListener, @@ -568,18 +568,22 @@ impl Client { ); let result_store_for_recovery = result_store.clone(); + let restart_gate = RestartGate::new( + tool_ops.busy_tools(), + execution_service.running_scripts(), + tool_run_manager.client_update_pending(), + ); let command_context = CommandContext::new( nats_connection_manager.clone(), config_service.clone(), - tool_ops.clone(), - tool_run_manager.clone(), + restart_gate.clone(), nats_message_publisher.clone(), result_store.clone(), flush_notify.clone(), ); let feature_flags_listener = CommandListener::new( FeatureFlagsHandler::new(feature_flags_service), - command_context.clone(), + command_context, ); let command_execution_listener = ExecutionListener::::new( @@ -665,7 +669,7 @@ impl Client { client_uninstall_message_listener, remote_access_message_listener, feature_flags_listener, - command_context, + restart_gate, command_execution_listener, script_execution_listener, script_bootstrap_execution_listener, @@ -751,9 +755,9 @@ impl Client { } self.feature_flags_listener.start().await?; - let command_context = self.command_context.clone(); + let restart_gate = self.restart_gate.clone(); tokio::spawn(async move { - command_context.restart_requested().await; + restart_gate.requested().await; match crate::platform::detached_process::spawn_detached( &crate::service::Service::get_install_location(), cli::RESTART_SERVICE_SUBCOMMAND, diff --git a/clients/openframe-client/src/listener/client_update_gate.rs b/clients/openframe-client/src/listener/client_update_gate.rs index d5943ec41f..d0cc76860c 100644 --- a/clients/openframe-client/src/listener/client_update_gate.rs +++ b/clients/openframe-client/src/listener/client_update_gate.rs @@ -14,7 +14,7 @@ pub async fn park_or_dispatch( F: FnOnce(Message) -> Fut + Send + 'static, Fut: Future + Send + 'static, { - if !manager.is_client_update_pending().await { + if !manager.is_client_update_pending() { dispatch(message).await; return; } @@ -25,7 +25,7 @@ pub async fn park_or_dispatch( ); tokio::spawn(async move { - while manager.is_client_update_pending().await { + while manager.is_client_update_pending() { if let Err(e) = message.ack_with(AckKind::Progress).await { warn!( "Failed to send Progress ack for parked {}: {} — abandoning park, redelivery takes over", diff --git a/clients/openframe-client/src/listener/command/integration_tests.rs b/clients/openframe-client/src/listener/command/integration_tests.rs index f5960c34b0..b7e8036728 100644 --- a/clients/openframe-client/src/listener/command/integration_tests.rs +++ b/clients/openframe-client/src/listener/command/integration_tests.rs @@ -7,12 +7,13 @@ use tempfile::TempDir; use tokio::sync::Notify; use tokio::time::{timeout, Duration}; -use super::{CommandContext, CommandListener, FeatureFlagsHandler}; +use super::{CommandContext, CommandListener, FeatureFlagsHandler, RestartGate}; use crate::clients::AuthClient; use crate::config::update_config::DEV_CLIENT_VERSION; use crate::models::{DeliveryResult, DeliveryResultMessage, Feature, FeatureFlagsState}; use crate::platform::DirectoryManager; use crate::services::encryption_service::EncryptionService; +use crate::services::execution_service::ExecutionService; use crate::services::feature_flags_service::FeatureFlagsService; use crate::services::local_tls_config_provider::LocalTlsConfigProvider; use crate::services::nats_connection_manager::NatsConnectionManager; @@ -50,11 +51,11 @@ struct Agent { } impl Agent { - async fn start() -> Option<(Self, CommandContext)> { + async fn start() -> Option<(Self, RestartGate)> { Self::start_as(VERSION).await } - async fn start_as(version: &str) -> Option<(Self, CommandContext)> { + async fn start_as(version: &str) -> Option<(Self, RestartGate)> { let server = NatsServer::start()?; let dir = tempfile::tempdir().unwrap(); let directories = DirectoryManager::with_custom_dirs( @@ -133,16 +134,20 @@ impl Agent { store, flush, }; - let (context, _service) = agent.launch().await; - Some((agent, context)) + let (gate, _service) = agent.launch().await; + Some((agent, gate)) } - async fn launch(&mut self) -> (CommandContext, Arc) { + async fn launch(&mut self) -> (RestartGate, Arc) { + let gate = RestartGate::new( + ToolOps::default().busy_tools(), + ExecutionService::new().running_scripts(), + self.tool_run_manager.client_update_pending(), + ); let context = CommandContext::new( self.manager.clone(), self.config.clone(), - ToolOps::default(), - self.tool_run_manager.clone(), + gate.clone(), self.publisher.clone(), self.store.clone(), self.flush.clone(), @@ -151,15 +156,15 @@ impl Agent { self.flags_path.clone(), self.version.clone(), )); - CommandListener::new(FeatureFlagsHandler::new(service.clone()), context.clone()) + CommandListener::new(FeatureFlagsHandler::new(service.clone()), context) .start() .await .unwrap(); self.until_listening().await; - (context, service) + (gate, service) } - async fn launch_as(&mut self, version: &str) -> (CommandContext, Arc) { + async fn launch_as(&mut self, version: &str) -> (RestartGate, Arc) { self.version = version.to_string(); self.launch().await } @@ -242,8 +247,8 @@ fn entry(value: bool, min_version: &str) -> serde_json::Value { json!({ "enabled": value, "minVersion": min_version }) } -async fn expect_no_restart(context: &CommandContext) { - let requested = timeout(QUIET, context.restart_requested()).await; +async fn expect_no_restart(gate: &RestartGate) { + let requested = timeout(QUIET, gate.requested()).await; assert!(requested.is_err(), "a restart was requested"); } @@ -265,7 +270,7 @@ fn command( #[tokio::test] async fn unchanged_flags_are_acked_done_and_persisted() { - let Some((mut agent, _context)) = Agent::start().await else { + let Some((mut agent, _gate)) = Agent::start().await else { return; }; @@ -286,7 +291,7 @@ async fn unchanged_flags_are_acked_done_and_persisted() { #[tokio::test] async fn changed_flags_restart_and_the_next_process_reports_done() { - let Some((mut agent, context)) = Agent::start().await else { + let Some((mut agent, gate)) = Agent::start().await else { return; }; @@ -299,7 +304,7 @@ async fn changed_flags_restart_and_the_next_process_reports_done() { .await; agent.expect("d2", DeliveryResult::Acked).await; - timeout(BOUND, context.restart_requested()) + timeout(BOUND, gate.requested()) .await .expect("restart was never requested"); agent.expect_silence("d2").await; @@ -325,7 +330,7 @@ async fn changed_flags_restart_and_the_next_process_reports_done() { #[tokio::test] async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { - let Some((mut agent, context)) = Agent::start().await else { + let Some((mut agent, gate)) = Agent::start().await else { return; }; agent @@ -335,7 +340,7 @@ async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { json!({ "delivery-tool-installation": entry(true, "1.0.0") }), )) .await; - timeout(BOUND, context.restart_requested()) + timeout(BOUND, gate.requested()) .await .expect("restart was never requested"); @@ -353,7 +358,7 @@ async fn commands_after_a_requested_restart_are_left_to_the_backend_retry() { #[tokio::test] async fn stale_sequence_is_done_without_applying() { - let Some((mut agent, _context)) = Agent::start().await else { + let Some((mut agent, _gate)) = Agent::start().await else { return; }; agent.send(command("d5", Some(5), json!({}))).await; @@ -375,7 +380,7 @@ async fn stale_sequence_is_done_without_applying() { #[tokio::test] async fn malformed_flags_fail_with_a_reason() { - let Some((mut agent, _context)) = Agent::start().await else { + let Some((mut agent, _gate)) = Agent::start().await else { return; }; @@ -400,7 +405,7 @@ async fn malformed_flags_fail_with_a_reason() { #[tokio::test] async fn a_message_without_delivery_gets_no_answer() { - let Some((mut agent, _context)) = Agent::start().await else { + let Some((mut agent, _gate)) = Agent::start().await else { return; }; @@ -432,7 +437,7 @@ async fn a_message_without_delivery_gets_no_answer() { #[tokio::test] async fn flags_without_a_sequence_fail() { - let Some((mut agent, _context)) = Agent::start().await else { + let Some((mut agent, _gate)) = Agent::start().await else { return; }; @@ -451,7 +456,7 @@ async fn flags_without_a_sequence_fail() { #[tokio::test] async fn a_flag_for_a_newer_client_is_stored_without_a_restart_and_turns_on_after_the_update() { - let Some((mut agent, context)) = Agent::start_as("1.5.9").await else { + let Some((mut agent, gate)) = Agent::start_as("1.5.9").await else { return; }; @@ -465,7 +470,7 @@ async fn a_flag_for_a_newer_client_is_stored_without_a_restart_and_turns_on_afte agent.expect("d20", DeliveryResult::Acked).await; agent.expect("d20", DeliveryResult::Done).await; - expect_no_restart(&context).await; + expect_no_restart(&gate).await; assert!(agent .persisted() .flags @@ -481,7 +486,7 @@ async fn a_flag_for_a_newer_client_is_stored_without_a_restart_and_turns_on_afte #[tokio::test] async fn a_rollback_below_the_min_version_runs_with_the_flag_off() { - let Some((mut agent, context)) = Agent::start().await else { + let Some((mut agent, gate)) = Agent::start().await else { return; }; agent @@ -491,7 +496,7 @@ async fn a_rollback_below_the_min_version_runs_with_the_flag_off() { json!({ "delivery-tool-installation": entry(true, "1.5.0") }), )) .await; - timeout(BOUND, context.restart_requested()) + timeout(BOUND, gate.requested()) .await .expect("restart was never requested"); @@ -504,7 +509,7 @@ async fn a_rollback_below_the_min_version_runs_with_the_flag_off() { #[tokio::test] async fn a_dev_build_applies_a_true_flag_whatever_its_min_version() { - let Some((mut agent, context)) = Agent::start_as(DEV_CLIENT_VERSION).await else { + let Some((mut agent, gate)) = Agent::start_as(DEV_CLIENT_VERSION).await else { return; }; @@ -517,14 +522,14 @@ async fn a_dev_build_applies_a_true_flag_whatever_its_min_version() { .await; agent.expect("d22", DeliveryResult::Acked).await; - timeout(BOUND, context.restart_requested()) + timeout(BOUND, gate.requested()) .await .expect("restart was never requested"); } #[tokio::test] async fn invalid_flag_entries_are_done_without_a_restart() { - let Some((mut agent, context)) = Agent::start().await else { + let Some((mut agent, gate)) = Agent::start().await else { return; }; @@ -542,6 +547,6 @@ async fn invalid_flag_entries_are_done_without_a_restart() { agent.expect("d23", DeliveryResult::Acked).await; agent.expect("d23", DeliveryResult::Done).await; - expect_no_restart(&context).await; + expect_no_restart(&gate).await; assert!(agent.persisted().flags.is_empty()); } diff --git a/clients/openframe-client/src/listener/command/mod.rs b/clients/openframe-client/src/listener/command/mod.rs index 1d69ea2e19..6b8c13b773 100644 --- a/clients/openframe-client/src/listener/command/mod.rs +++ b/clients/openframe-client/src/listener/command/mod.rs @@ -17,14 +17,12 @@ use crate::models::DeliveryRef; use crate::services::nats_connection_manager::NatsConnectionManager; use crate::services::nats_message_publisher::NatsMessagePublisher; use crate::services::result_store::ResultStore; -use crate::services::tool_ops::ToolOps; -use crate::services::tool_run_manager::ToolRunManager; use crate::services::AgentConfigurationService; use delivery::Delivery; -use restart::RestartGate; pub use handlers::FeatureFlagsHandler; +pub use restart::RestartGate; pub trait CommandHandler: Clone + Send + Sync + 'static { type Command: DeserializeOwned + Send + 'static; @@ -68,8 +66,7 @@ impl CommandContext { pub fn new( nats_connection_manager: NatsConnectionManager, config_service: AgentConfigurationService, - tool_ops: ToolOps, - tool_run_manager: ToolRunManager, + restart: RestartGate, publisher: NatsMessagePublisher, result_store: Arc, flush_notify: Arc, @@ -77,18 +74,12 @@ impl CommandContext { Self { nats_connection_manager, config_service, - restart: RestartGate::new(tool_ops, tool_run_manager), + restart, delivery: Delivery::new(publisher, result_store, flush_notify), } } } -impl CommandContext { - pub async fn restart_requested(&self) { - self.restart.requested().await; - } -} - #[derive(Clone)] pub struct CommandListener { handler: H, diff --git a/clients/openframe-client/src/listener/command/restart.rs b/clients/openframe-client/src/listener/command/restart.rs index cb8924134f..d16a655078 100644 --- a/clients/openframe-client/src/listener/command/restart.rs +++ b/clients/openframe-client/src/listener/command/restart.rs @@ -6,26 +6,33 @@ use tokio::time::{sleep, timeout_at, Duration, Instant}; use tracing::{info, warn}; use crate::config::update_config::{SELF_RESTART_IDLE_POLL_SECS, SELF_RESTART_IDLE_WAIT_SECS}; -use crate::services::tool_ops::ToolOps; -use crate::services::tool_run_manager::ToolRunManager; +use crate::services::execution_service::RunningScripts; +use crate::services::tool_ops::BusyTools; +use crate::services::tool_run_manager::ClientUpdatePending; #[derive(Clone)] -pub(super) struct RestartGate { +pub struct RestartGate { closing: Arc, work: Arc>, requested: Arc, - tool_ops: ToolOps, - tool_run_manager: ToolRunManager, + busy_tools: BusyTools, + running_scripts: RunningScripts, + client_update: ClientUpdatePending, } impl RestartGate { - pub(super) fn new(tool_ops: ToolOps, tool_run_manager: ToolRunManager) -> Self { + pub fn new( + busy_tools: BusyTools, + running_scripts: RunningScripts, + client_update: ClientUpdatePending, + ) -> Self { Self { closing: Arc::new(AtomicBool::new(false)), work: Arc::new(RwLock::new(())), requested: Arc::new(Notify::new()), - tool_ops, - tool_run_manager, + busy_tools, + running_scripts, + client_update, } } @@ -48,23 +55,29 @@ impl RestartGate { SELF_RESTART_IDLE_WAIT_SECS ); } - while self.busy().await && Instant::now() < deadline { + while self.busy() && Instant::now() < deadline { sleep(Duration::from_secs(SELF_RESTART_IDLE_POLL_SECS)).await; } - if self.busy().await { + if self.busy() { warn!( - "Tool operation or client update still running after {}s, restarting anyway", + "Tool operation, script or client update still running after {}s, restarting anyway", SELF_RESTART_IDLE_WAIT_SECS ); } + // A running client update restarts the client itself; stopping the service under it could + // interrupt the binary swap. With a tool still busy the update is only deferred and keeps + // re-marking the flag on redelivery, so it is not worth waiting for. + while self.client_update.is_pending() && !self.busy_tools.any() { + sleep(Duration::from_secs(SELF_RESTART_IDLE_POLL_SECS)).await; + } self.requested.notify_one(); } - async fn busy(&self) -> bool { - self.tool_ops.any_busy() || self.tool_run_manager.is_client_update_pending().await + fn busy(&self) -> bool { + self.busy_tools.any() || self.running_scripts.any() || self.client_update.is_pending() } - pub(super) async fn requested(&self) { + pub async fn requested(&self) { self.requested.notified().await; } } diff --git a/clients/openframe-client/src/listener/command/restart_tests.rs b/clients/openframe-client/src/listener/command/restart_tests.rs index 12a60677d7..0f09de6a02 100644 --- a/clients/openframe-client/src/listener/command/restart_tests.rs +++ b/clients/openframe-client/src/listener/command/restart_tests.rs @@ -1,14 +1,29 @@ use super::*; -use crate::platform::DirectoryManager; -use crate::services::{ - AgentConfigurationService, InitialConfigurationService, InstalledToolsService, - ToolCommandParamsResolver, ToolKillService, -}; -use tempfile::TempDir; +use crate::config::update_config::SELF_RESTART_IDLE_WAIT_SECS; +use crate::services::execution_service::ExecutionService; +use crate::services::tool_ops::ToolOps; +use crate::services::tool_run_manager::ClientUpdatePendingFlag; use tokio::time::timeout; const BOUND: Duration = Duration::from_secs(10); +#[derive(Default)] +struct Work { + tool_ops: ToolOps, + scripts: ExecutionService, + client_update: ClientUpdatePendingFlag, +} + +impl Work { + fn gate(&self) -> RestartGate { + RestartGate::new( + self.tool_ops.busy_tools(), + self.scripts.running_scripts(), + self.client_update.view(), + ) + } +} + async fn until_closed(gate: &RestartGate) { timeout(BOUND, async { while gate.admit().is_some() { @@ -24,39 +39,24 @@ fn start_drain(gate: &RestartGate) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { gate.drain_and_request().await }) } -fn gate_with(tool_ops: ToolOps) -> (RestartGate, ToolRunManager, TempDir) { - let dir = TempDir::new().unwrap(); - let directories = DirectoryManager::with_custom_dirs( - dir.path().join("logs"), - dir.path().join("app_support"), - dir.path().join("secured"), - ); - let tool_run_manager = ToolRunManager::new( - InstalledToolsService::new(directories.clone()).unwrap(), - ToolCommandParamsResolver::new( - directories.clone(), - InitialConfigurationService::new(directories.clone()).unwrap(), - AgentConfigurationService::new(directories).unwrap(), - ), - ToolKillService::new(), - tool_ops.clone(), - ); - ( - RestartGate::new(tool_ops, tool_run_manager.clone()), - tool_run_manager, - dir, - ) -} - async fn until_requested(gate: &RestartGate) { timeout(BOUND, gate.requested()) .await .expect("restart was never requested"); } +async fn drain_waits_until_released(gate: &RestartGate, release: impl FnOnce()) { + let draining = start_drain(gate); + until_closed(gate).await; + assert!(!draining.is_finished()); + + release(); + until_requested(gate).await; +} + #[tokio::test] async fn admits_concurrent_commands_while_no_restart_is_pending() { - let (gate, _manager, _dir) = gate_with(ToolOps::default()); + let gate = Work::default().gate(); let first = gate.admit(); let second = gate.admit(); @@ -67,36 +67,48 @@ async fn admits_concurrent_commands_while_no_restart_is_pending() { #[tokio::test] async fn drain_refuses_new_commands_and_waits_for_the_admitted_one() { - let (gate, _manager, _dir) = gate_with(ToolOps::default()); + let gate = Work::default().gate(); let running = gate.admit().expect("admitted before the restart"); - let draining = start_drain(&gate); - until_closed(&gate).await; - assert!(!draining.is_finished()); - - drop(running); - until_requested(&gate).await; + drain_waits_until_released(&gate, || drop(running)).await; assert!(gate.admit().is_none()); } #[tokio::test] async fn drain_waits_for_a_busy_tool() { - let tool_ops = ToolOps::default(); - let (gate, _manager, _dir) = gate_with(tool_ops.clone()); - let op = tool_ops.try_lock("fleetmdm-agent").unwrap().mark_busy(); + let work = Work::default(); + let gate = work.gate(); + let op = work + .tool_ops + .try_lock("fleetmdm-agent") + .unwrap() + .mark_busy(); + + drain_waits_until_released(&gate, || drop(op)).await; +} - let draining = start_drain(&gate); - until_closed(&gate).await; - assert!(!draining.is_finished()); +#[tokio::test] +async fn drain_waits_for_a_running_script() { + let work = Work::default(); + let gate = work.gate(); + let script = work.scripts.running_for_tests(); - drop(op); - until_requested(&gate).await; + drain_waits_until_released(&gate, || drop(script)).await; +} + +#[tokio::test] +async fn drain_waits_for_a_pending_client_update() { + let work = Work::default(); + let gate = work.gate(); + work.client_update.mark(); + + drain_waits_until_released(&gate, || work.client_update.clear()).await; } #[tokio::test] async fn the_gate_stays_closed_after_the_restart_is_requested() { - let (gate, _manager, _dir) = gate_with(ToolOps::default()); + let gate = Work::default().gate(); start_drain(&gate); until_requested(&gate).await; @@ -104,15 +116,39 @@ async fn the_gate_stays_closed_after_the_restart_is_requested() { assert!(gate.admit().is_none()); } -#[tokio::test] -async fn drain_waits_for_a_pending_client_update() { - let (gate, tool_run_manager, _dir) = gate_with(ToolOps::default()); - tool_run_manager.mark_client_update_pending().await; +#[tokio::test(start_paused = true)] +async fn a_running_client_update_holds_the_restart_past_the_deadline() { + let work = Work::default(); + let gate = work.gate(); + work.client_update.mark(); let draining = start_drain(&gate); - until_closed(&gate).await; + sleep(Duration::from_secs(SELF_RESTART_IDLE_WAIT_SECS * 2)).await; assert!(!draining.is_finished()); - tool_run_manager.clear_client_update_pending().await; + work.client_update.clear(); until_requested(&gate).await; } + +#[tokio::test(start_paused = true)] +async fn a_client_update_deferred_by_a_busy_tool_does_not_hold_the_restart() { + let work = Work::default(); + let gate = work.gate(); + work.client_update.mark(); + let _op = work + .tool_ops + .try_lock("fleetmdm-agent") + .unwrap() + .mark_busy(); + + let draining = start_drain(&gate); + timeout( + Duration::from_secs(SELF_RESTART_IDLE_WAIT_SECS * 2), + gate.requested(), + ) + .await + .expect("restart was never requested"); + + assert!(work.client_update.view().is_pending()); + draining.await.unwrap(); +} diff --git a/clients/openframe-client/src/services/execution_service.rs b/clients/openframe-client/src/services/execution_service.rs index 669a950294..6876861aab 100644 --- a/clients/openframe-client/src/services/execution_service.rs +++ b/clients/openframe-client/src/services/execution_service.rs @@ -1,3 +1,5 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; use std::time::Instant; use crate::models::{ExecutionRequest, RmmResult}; @@ -8,14 +10,48 @@ use crate::models::PrivilegePolicy; use crate::{executor::Privilege, models::PrivilegeLevel}; #[derive(Clone, Default)] -pub struct ExecutionService; +pub struct ExecutionService { + running: Arc, +} + +struct Running(Arc); + +#[derive(Clone)] +pub struct RunningScripts(Arc); + +impl RunningScripts { + pub fn any(&self) -> bool { + self.0.load(Ordering::Acquire) > 0 + } +} + +impl Drop for Running { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::AcqRel); + } +} impl ExecutionService { pub fn new() -> Self { - Self + Self::default() + } + + pub fn running_scripts(&self) -> RunningScripts { + RunningScripts(self.running.clone()) + } + + fn track(&self) -> Running { + self.running.fetch_add(1, Ordering::AcqRel); + Running(self.running.clone()) + } + + #[cfg(test)] + pub(crate) fn running_for_tests(&self) -> impl Drop { + self.track() } pub async fn execute(&self, req: &ExecutionRequest<'_>, machine_id: &str) -> RmmResult { + let _running = self.track(); let start = Instant::now(); #[cfg(any(unix, windows))] diff --git a/clients/openframe-client/src/services/execution_service_tests.rs b/clients/openframe-client/src/services/execution_service_tests.rs index f2a5ee2de1..9dd4c4a15e 100644 --- a/clients/openframe-client/src/services/execution_service_tests.rs +++ b/clients/openframe-client/src/services/execution_service_tests.rs @@ -47,3 +47,26 @@ async fn maps_spawn_failure_to_error() { assert!(r.error.is_some()); assert!(!r.timed_out); } + +#[tokio::test] +async fn a_script_counts_as_running_until_it_finishes() { + let service = ExecutionService::new(); + let running = service.running_scripts(); + assert!(!running.any()); + + let script = tokio::spawn({ + let service = service.clone(); + async move { service.execute(&req("#!/bin/sh\nsleep 1\n"), "m").await } + }); + tokio::time::timeout(std::time::Duration::from_secs(10), async { + while !running.any() { + tokio::task::yield_now().await; + } + }) + .await + .expect("the script never started"); + + let result = script.await.unwrap(); + assert_eq!(result.exit_code, 0); + assert!(!running.any()); +} diff --git a/clients/openframe-client/src/services/openframe_client_update_service.rs b/clients/openframe-client/src/services/openframe_client_update_service.rs index 46921d74ec..784f0b6de7 100644 --- a/clients/openframe-client/src/services/openframe_client_update_service.rs +++ b/clients/openframe-client/src/services/openframe_client_update_service.rs @@ -78,7 +78,7 @@ impl OpenFrameClientUpdateService { let requested_version = message.version.trim(); info!("Received update request for version: {}", requested_version); - self.tool_run_manager.mark_client_update_pending().await; + self.tool_run_manager.mark_client_update_pending(); if self.tool_ops.any_busy() { warn!("Tool operation in progress, deferring client update to version {} (will redeliver)", requested_version); @@ -123,7 +123,7 @@ impl OpenFrameClientUpdateService { "Already running version {}, ignoring update request", canonical_version ); - self.tool_run_manager.clear_client_update_pending().await; + self.tool_run_manager.clear_client_update_pending(); return Ok(()); } @@ -148,7 +148,7 @@ impl OpenFrameClientUpdateService { "refusing downgrade to {} — anchored at {}", requested_version, anchor ); - self.tool_run_manager.clear_client_update_pending().await; + self.tool_run_manager.clear_client_update_pending(); return Ok(()); } Ok(_) => {} @@ -262,7 +262,7 @@ impl OpenFrameClientUpdateService { binary_bytes.len() ); - self.tool_run_manager.mark_client_update_pending().await; + self.tool_run_manager.mark_client_update_pending(); // 3. Extract binary update_state.set_phase(UpdatePhase::Extracting); @@ -433,7 +433,7 @@ impl OpenFrameClientUpdateService { ); } - self.tool_run_manager.clear_client_update_pending().await; + self.tool_run_manager.clear_client_update_pending(); if let Err(e) = self.tool_run_manager.resume_after_update_failure().await { warn!( "Failed to resume tool supervision after the aborted update to {}: {:#}", diff --git a/clients/openframe-client/src/services/tool_ops.rs b/clients/openframe-client/src/services/tool_ops.rs index c35d48836a..7b81b6bd2e 100644 --- a/clients/openframe-client/src/services/tool_ops.rs +++ b/clients/openframe-client/src/services/tool_ops.rs @@ -18,6 +18,20 @@ struct OpState { #[derive(Debug)] pub struct Busy; +#[derive(Clone)] +pub struct BusyTools(Arc>); + +impl BusyTools { + pub fn any(&self) -> bool { + !self + .0 + .lock() + .unwrap_or_else(PoisonError::into_inner) + .busy + .is_empty() + } +} + pub struct ToolLock { ops: ToolOps, tool_agent_id: String, @@ -54,12 +68,11 @@ impl ToolOps { } pub fn any_busy(&self) -> bool { - !self - .state - .lock() - .unwrap_or_else(PoisonError::into_inner) - .busy - .is_empty() + self.busy_tools().any() + } + + pub fn busy_tools(&self) -> BusyTools { + BusyTools(self.state.clone()) } pub(crate) fn is_released(&self, tool_agent_id: &str) -> bool { diff --git a/clients/openframe-client/src/services/tool_ops_tests.rs b/clients/openframe-client/src/services/tool_ops_tests.rs index 71b8f92108..85449f25f3 100644 --- a/clients/openframe-client/src/services/tool_ops_tests.rs +++ b/clients/openframe-client/src/services/tool_ops_tests.rs @@ -59,6 +59,18 @@ async fn any_busy_follows_every_tool() { assert!(!ops.any_busy()); } +#[tokio::test] +async fn a_busy_tools_view_follows_ops_started_after_it() { + let ops = ToolOps::default(); + let view = ops.busy_tools(); + + let op = ops.lock(TOOL).await.mark_busy(); + assert!(view.any()); + + drop(op); + assert!(!view.any()); +} + #[tokio::test] async fn waiting_lock_sees_the_flag_cleared() { let ops = ToolOps::default(); diff --git a/clients/openframe-client/src/services/tool_run_manager.rs b/clients/openframe-client/src/services/tool_run_manager.rs index a80ae295f5..97aee25895 100644 --- a/clients/openframe-client/src/services/tool_run_manager.rs +++ b/clients/openframe-client/src/services/tool_run_manager.rs @@ -69,20 +69,41 @@ async fn shutdown_break( #[derive(Clone, Default)] pub(crate) struct ClientUpdatePendingFlag { - since: Arc>>, + since: Arc>>, } impl ClientUpdatePendingFlag { - pub(crate) async fn mark(&self) { - *self.since.write().await = Some(std::time::Instant::now()); + pub(crate) fn mark(&self) { + *self.since() = Some(std::time::Instant::now()); } - pub(crate) async fn is_pending(&self, ttl: Duration) -> bool { - matches!(*self.since.read().await, Some(since) if since.elapsed() < ttl) + pub(crate) fn is_pending(&self, ttl: Duration) -> bool { + matches!(*self.since(), Some(since) if since.elapsed() < ttl) } - pub(crate) async fn clear(&self) { - *self.since.write().await = None; + pub(crate) fn clear(&self) { + *self.since() = None; + } + + pub(crate) fn view(&self) -> ClientUpdatePending { + ClientUpdatePending(self.clone()) + } + + fn since(&self) -> std::sync::MutexGuard<'_, Option> { + self.since + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } +} + +#[derive(Clone)] +pub struct ClientUpdatePending(ClientUpdatePendingFlag); + +impl ClientUpdatePending { + pub fn is_pending(&self) -> bool { + self.0.is_pending(Duration::from_secs( + crate::config::update_config::CLIENT_UPDATE_PENDING_TTL_SECS, + )) } } @@ -212,21 +233,21 @@ impl ToolRunManager { Ok(()) } - pub async fn mark_client_update_pending(&self) { - self.client_update_pending.mark().await; + pub fn mark_client_update_pending(&self) { + self.client_update_pending.mark(); info!("Client update pending: new tool operations will be parked"); } - pub async fn is_client_update_pending(&self) -> bool { - self.client_update_pending - .is_pending(Duration::from_secs( - crate::config::update_config::CLIENT_UPDATE_PENDING_TTL_SECS, - )) - .await + pub fn is_client_update_pending(&self) -> bool { + self.client_update_pending.view().is_pending() + } + + pub fn client_update_pending(&self) -> ClientUpdatePending { + self.client_update_pending.view() } - pub async fn clear_client_update_pending(&self) { - self.client_update_pending.clear().await; + pub fn clear_client_update_pending(&self) { + self.client_update_pending.clear(); info!("Client update no longer pending: parked tool operations released"); } diff --git a/clients/openframe-client/src/services/tool_run_manager_tests.rs b/clients/openframe-client/src/services/tool_run_manager_tests.rs index 60fd97a0bd..ac088c677e 100644 --- a/clients/openframe-client/src/services/tool_run_manager_tests.rs +++ b/clients/openframe-client/src/services/tool_run_manager_tests.rs @@ -7,42 +7,42 @@ use tokio::sync::RwLock; const LONG_TTL: Duration = Duration::from_secs(3600); -#[tokio::test] -async fn not_pending_before_first_mark() { +#[test] +fn not_pending_before_first_mark() { let flag = ClientUpdatePendingFlag::default(); - assert!(!flag.is_pending(LONG_TTL).await); + assert!(!flag.is_pending(LONG_TTL)); } -#[tokio::test] -async fn pending_after_mark_within_ttl() { +#[test] +fn pending_after_mark_within_ttl() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; - assert!(flag.is_pending(LONG_TTL).await); + flag.mark(); + assert!(flag.is_pending(LONG_TTL)); } -#[tokio::test] -async fn expired_when_ttl_elapsed() { +#[test] +fn expired_when_ttl_elapsed() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; - assert!(!flag.is_pending(Duration::ZERO).await); + flag.mark(); + assert!(!flag.is_pending(Duration::ZERO)); } #[tokio::test] async fn remark_refreshes_the_ttl() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; + flag.mark(); tokio::time::sleep(Duration::from_millis(30)).await; - assert!(!flag.is_pending(Duration::from_millis(10)).await); - flag.mark().await; - assert!(flag.is_pending(Duration::from_millis(10)).await); + assert!(!flag.is_pending(Duration::from_millis(10))); + flag.mark(); + assert!(flag.is_pending(Duration::from_millis(10))); } -#[tokio::test] -async fn clones_share_state() { +#[test] +fn clones_share_state() { let flag = ClientUpdatePendingFlag::default(); let clone = flag.clone(); - clone.mark().await; - assert!(flag.is_pending(LONG_TTL).await); + clone.mark(); + assert!(flag.is_pending(LONG_TTL)); } #[tokio::test] @@ -57,13 +57,26 @@ async fn shutdown_break_drops_supervision_only_once_signalled() { assert!(!running_tools.read().await.contains("mesh")); } -#[tokio::test] -async fn clear_releases_the_flag() { +#[test] +fn clear_releases_the_flag() { + let flag = ClientUpdatePendingFlag::default(); + flag.mark(); + assert!(flag.is_pending(LONG_TTL)); + flag.clear(); + assert!(!flag.is_pending(LONG_TTL)); +} + +#[test] +fn the_view_follows_the_flag() { let flag = ClientUpdatePendingFlag::default(); - flag.mark().await; - assert!(flag.is_pending(LONG_TTL).await); - flag.clear().await; - assert!(!flag.is_pending(LONG_TTL).await); + let view = flag.view(); + assert!(!view.is_pending()); + + flag.mark(); + assert!(view.is_pending()); + + flag.clear(); + assert!(!view.is_pending()); } fn gui_app_record(state: ToolRecordState) -> InstalledTool {