diff --git a/include/modules/custom.hpp b/include/modules/custom.hpp index 442b1c37..dbdcb698 100644 --- a/include/modules/custom.hpp +++ b/include/modules/custom.hpp @@ -3,10 +3,12 @@ #include #include +#include #include #include "AIconLabel.hpp" #include "util/command.hpp" +#include "util/command_line_stream.hpp" #include "util/json.hpp" #include "util/sleeper_thread.hpp" @@ -22,6 +24,9 @@ class Custom : public AIconLabel { private: void delayWorker(); void continuousWorker(); + void startContinuousProcess(bool throw_on_failure); + void handleContinuousProcessExit(int exit_code); + void scheduleContinuousRestart(); void waitingWorker(); void parseOutputRaw(); void parseOutputJson(); @@ -42,10 +47,10 @@ class Custom : public AIconLabel { const bool tooltip_format_enabled_; std::vector class_; int percentage_; - FILE* fp_; - int pid_; util::command::res output_; util::JsonParser parser_; + std::unique_ptr continuous_stream_; + sigc::connection restart_connection_; util::SleeperThread thread_; }; diff --git a/include/util/command_line_stream.hpp b/include/util/command_line_stream.hpp new file mode 100644 index 00000000..05c82fc4 --- /dev/null +++ b/include/util/command_line_stream.hpp @@ -0,0 +1,40 @@ +#pragma once + +#include +#include + +#include +#include + +namespace waybar::util::command { + +class LineStream { + public: + using OutputCallback = std::function; + using ExitCallback = std::function; + + LineStream(std::string output_name, OutputCallback on_output, ExitCallback on_exit); + ~LineStream(); + + void start(const std::string& cmd); + void stop(); + bool running() const; + + private: + bool handleStdout(Glib::IOCondition condition); + void handleExit(Glib::Pid pid, int status); + void closeStdout(); + void drainStdout(bool flush_trailing_line); + static int statusToExitCode(int status); + + std::string output_name_; + OutputCallback on_output_; + ExitCallback on_exit_; + std::string buffer_; + Glib::Pid pid_; + int stdout_fd_; + sigc::connection stdout_connection_; + sigc::connection child_connection_; +}; + +} // namespace waybar::util::command diff --git a/meson.build b/meson.build index f96139a0..e3839cf5 100644 --- a/meson.build +++ b/meson.build @@ -196,7 +196,8 @@ src_files = files( 'src/util/regex_collection.cpp', 'src/util/css_reload_helper.cpp', 'src/util/transform_8bit_to_rgba.cpp', - 'src/util/utf8_string.cpp' + 'src/util/utf8_string.cpp', + 'src/util/command_line_stream.cpp' ) man_files = files( diff --git a/src/modules/custom.cpp b/src/modules/custom.cpp index 83f36797..64577509 100644 --- a/src/modules/custom.cpp +++ b/src/modules/custom.cpp @@ -2,10 +2,10 @@ #include +#include +#include #include -#include "util/scope_guard.hpp" - waybar::modules::Custom::Custom(const std::string& name, const std::string& id, const Json::Value& config, const std::string& output_name) : AIconLabel(config, "custom-" + name, id, "{}"), @@ -13,9 +13,7 @@ waybar::modules::Custom::Custom(const std::string& name, const std::string& id, output_name_(output_name), id_(id), tooltip_format_enabled_{config_["tooltip-format"].isString()}, - percentage_(0), - fp_(nullptr), - pid_(-1) { + percentage_(0) { if (config.isNull()) { spdlog::warn("There is no configuration for 'custom/{}', element will be hidden", name); } @@ -40,10 +38,9 @@ waybar::modules::Custom::Custom(const std::string& name, const std::string& id, } waybar::modules::Custom::~Custom() { - if (pid_ != -1) { - killpg(pid_, SIGTERM); - waitpid(pid_, NULL, 0); - pid_ = -1; + restart_connection_.disconnect(); + if (continuous_stream_) { + continuous_stream_->stop(); } } @@ -54,13 +51,21 @@ void waybar::modules::Custom::delayWorker() { } thread_ = [this] { - for (int i : this->pid_children_) { - int status; - waitpid(i, &status, 0); + for (auto it = this->pid_children_.begin(); it != this->pid_children_.end();) { + int status = 0; + const auto pid = static_cast(*it); + const auto waited = waitpid(pid, &status, WNOHANG); + if (waited == 0) { + ++it; + continue; + } + if (waited == -1 && errno != ECHILD) { + ++it; + continue; + } + it = this->pid_children_.erase(it); } - this->pid_children_.clear(); - bool can_update = true; if (config_["exec-if"].isString()) { output_ = util::command::execNoRead(config_["exec-if"].asString()); @@ -80,57 +85,64 @@ void waybar::modules::Custom::delayWorker() { } void waybar::modules::Custom::continuousWorker() { - auto cmd = config_["exec"].asString(); - pid_ = -1; - fp_ = util::command::open(cmd, pid_, output_name_); - if (!fp_) { - throw std::runtime_error("Unable to open " + cmd); - } - thread_ = [this, cmd] { - char* buff = nullptr; - waybar::util::ScopeGuard buff_deleter([&buff]() { - if (buff) { - free(buff); - } - }); - size_t len = 0; - if (getline(&buff, &len, fp_) == -1) { - int exit_code = 1; - if (fp_) { - exit_code = WEXITSTATUS(util::command::close(fp_, pid_)); - fp_ = nullptr; - } - if (exit_code != 0) { - output_ = {exit_code, ""}; + continuous_stream_ = std::make_unique( + output_name_, + [this](const std::string& output) { + output_ = {.exit_code = 0, .out = output}; dp.emit(); - spdlog::error("{} stopped unexpectedly, is it endless?", name_); - } - if (config_["restart-interval"].isNumeric() && config_["restart-interval"].asDouble() > 0) { - pid_ = -1; - thread_.sleep_for(std::chrono::milliseconds( - std::max(1L, // Minimum 1ms due to millisecond precision - static_cast(config_["restart-interval"].asDouble() * 1000)))); - fp_ = util::command::open(cmd, pid_, output_name_); - if (!fp_) { - throw std::runtime_error("Unable to open " + cmd); - } - } else { - // A non-positive restart-interval must not busy-respawn the script - // (that starves the GTK main loop); treat it as "do not restart". - thread_.stop(); - return; - } - } else { - std::string output = buff; + }, + [this](int exit_code) { handleContinuousProcessExit(exit_code); }); + startContinuousProcess(true); +} - // Remove last newline - if (!output.empty() && output[output.length() - 1] == '\n') { - output.erase(output.length() - 1); - } - output_ = {0, output}; - dp.emit(); +void waybar::modules::Custom::startContinuousProcess(bool throw_on_failure) { + const auto cmd = config_["exec"].asString(); + + try { + continuous_stream_->start(cmd); + } catch (const Glib::SpawnError& e) { + if (throw_on_failure) { + throw std::runtime_error("Unable to open " + cmd + ": " + e.what().raw()); } - }; + output_ = {.exit_code = 1, .out = ""}; + dp.emit(); + spdlog::error("Unable to restart {}: {}", name_, e.what().raw()); + scheduleContinuousRestart(); + } catch (const std::exception& e) { + if (throw_on_failure) { + throw; + } + output_ = {.exit_code = 1, .out = ""}; + dp.emit(); + spdlog::error("Unable to restart {}: {}", name_, e.what()); + scheduleContinuousRestart(); + } +} + +void waybar::modules::Custom::handleContinuousProcessExit(int exit_code) { + if (exit_code != 0) { + output_ = {.exit_code = exit_code, .out = ""}; + dp.emit(); + spdlog::error("{} stopped unexpectedly, is it endless?", name_); + } + + scheduleContinuousRestart(); +} + +void waybar::modules::Custom::scheduleContinuousRestart() { + restart_connection_.disconnect(); + if (!config_["restart-interval"].isNumeric() || config_["restart-interval"].asDouble() <= 0) { + // A non-positive restart-interval must not busy-respawn the script + // (that starves the GTK main loop); treat it as "do not restart". + return; + } + + restart_connection_ = Glib::signal_timeout().connect( + [this] { + startContinuousProcess(false); + return false; + }, + std::max(1U, static_cast(config_["restart-interval"].asDouble() * 1000))); } void waybar::modules::Custom::waitingWorker() { diff --git a/src/util/command_line_stream.cpp b/src/util/command_line_stream.cpp new file mode 100644 index 00000000..7abb79b8 --- /dev/null +++ b/src/util/command_line_stream.cpp @@ -0,0 +1,219 @@ +#include "util/command_line_stream.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#ifdef __linux__ +#include +#endif +#ifdef __FreeBSD__ +#include +#endif + +#include "util/command.hpp" + +namespace { + +auto buildChildEnvironment(const std::string& output_name) -> std::vector { + auto names = Glib::listenv(); + std::vector envp; + envp.reserve(names.size() + 1); + + for (const auto& name : names) { + bool found = false; + auto value = Glib::getenv(name, found); + if (!found || name == "WAYBAR_OUTPUT_NAME") { + continue; + } + envp.push_back(name + "=" + value); + } + + if (!output_name.empty()) { + envp.push_back("WAYBAR_OUTPUT_NAME=" + output_name); + } + + return envp; +} + +void prepareChild() { + sigset_t mask; + sigfillset(&mask); + + (void)pthread_sigmask(SIG_UNBLOCK, &mask, nullptr); + + int deathsig = SIGTERM; +#ifdef __linux__ + (void)prctl(PR_SET_PDEATHSIG, deathsig); +#endif +#ifdef __FreeBSD__ + (void)procctl(P_PID, 0, PROC_PDEATHSIG_CTL, reinterpret_cast(&deathsig)); +#endif + + (void)setpgid(0, 0); +} + +void emitBufferedLines(std::string& buffer, + const waybar::util::command::LineStream::OutputCallback& on_output, + bool flush_trailing_line) { + for (auto newline_pos = buffer.find('\n'); newline_pos != std::string::npos; + newline_pos = buffer.find('\n')) { + auto line = buffer.substr(0, newline_pos); + buffer.erase(0, newline_pos + 1); + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + on_output(line); + } + + if (flush_trailing_line && !buffer.empty()) { + auto line = std::move(buffer); + buffer.clear(); + if (!line.empty() && line.back() == '\r') { + line.pop_back(); + } + on_output(line); + } +} + +} // namespace + +waybar::util::command::LineStream::LineStream(std::string output_name, OutputCallback on_output, + ExitCallback on_exit) + : output_name_(std::move(output_name)), + on_output_(std::move(on_output)), + on_exit_(std::move(on_exit)), + pid_(0), + stdout_fd_(-1) {} + +waybar::util::command::LineStream::~LineStream() { stop(); } + +void waybar::util::command::LineStream::start(const std::string& cmd) { + stop(); + + std::vector argv{"/bin/sh", "-c", cmd}; + auto envp = buildChildEnvironment(output_name_); + Glib::spawn_async_with_pipes("", argv, envp, + Glib::SPAWN_DO_NOT_REAP_CHILD | Glib::SPAWN_CLOEXEC_PIPES, + sigc::ptr_fun(&prepareChild), &pid_, nullptr, &stdout_fd_, nullptr); + + const auto flags = fcntl(stdout_fd_, F_GETFL, 0); + if (flags == -1 || fcntl(stdout_fd_, F_SETFL, flags | O_NONBLOCK) == -1) { + const auto saved_errno = errno; + stop(); + throw std::runtime_error("Unable to configure child stdout: " + + std::string(std::strerror(saved_errno))); + } + + stdout_connection_ = + Glib::signal_io().connect(sigc::mem_fun(*this, &LineStream::handleStdout), stdout_fd_, + Glib::IO_IN | Glib::IO_HUP | Glib::IO_ERR | Glib::IO_NVAL); + child_connection_ = + Glib::signal_child_watch().connect(sigc::mem_fun(*this, &LineStream::handleExit), pid_); +} + +void waybar::util::command::LineStream::stop() { + stdout_connection_.disconnect(); + child_connection_.disconnect(); + + if (pid_ != 0) { + killpg(pid_, SIGTERM); + waitpid(pid_, nullptr, 0); + Glib::spawn_close_pid(pid_); + pid_ = 0; + } + + closeStdout(); + buffer_.clear(); +} + +bool waybar::util::command::LineStream::running() const { return pid_ != 0; } + +bool waybar::util::command::LineStream::handleStdout(Glib::IOCondition condition) { + const auto should_flush = + static_cast(condition & (Glib::IO_HUP | Glib::IO_ERR | Glib::IO_NVAL)); + drainStdout(should_flush); + + if (!running() || should_flush) { + closeStdout(); + return false; + } + + return true; +} + +void waybar::util::command::LineStream::handleExit(Glib::Pid pid, int status) { + child_connection_.disconnect(); + + if (stdout_fd_ != -1) { + drainStdout(true); + stdout_connection_.disconnect(); + closeStdout(); + } + + if (pid_ == pid) { + Glib::spawn_close_pid(pid_); + pid_ = 0; + } else { + Glib::spawn_close_pid(pid); + } + + on_exit_(statusToExitCode(status)); +} + +void waybar::util::command::LineStream::closeStdout() { + if (stdout_fd_ != -1) { + ::close(stdout_fd_); + stdout_fd_ = -1; + } +} + +void waybar::util::command::LineStream::drainStdout(bool flush_trailing_line) { + if (stdout_fd_ == -1) { + return; + } + + std::array chunk = {}; + while (true) { + const auto bytes_read = ::read(stdout_fd_, chunk.data(), chunk.size()); + if (bytes_read > 0) { + buffer_.append(chunk.data(), static_cast(bytes_read)); + emitBufferedLines(buffer_, on_output_, false); + continue; + } + + if (bytes_read == 0) { + emitBufferedLines(buffer_, on_output_, flush_trailing_line); + return; + } + + if (errno == EINTR) { + continue; + } + if (errno == EAGAIN || errno == EWOULDBLOCK) { + return; + } + + spdlog::error("Reading command stdout failed: {}", std::strerror(errno)); + emitBufferedLines(buffer_, on_output_, flush_trailing_line); + return; + } +} + +int waybar::util::command::LineStream::statusToExitCode(int status) { + if (WIFEXITED(status)) { + return WEXITSTATUS(status); + } + if (WIFSIGNALED(status)) { + return 128 + WTERMSIG(status); + } + return 1; +} diff --git a/test/utils/command_line_stream.cpp b/test/utils/command_line_stream.cpp new file mode 100644 index 00000000..b852a6e8 --- /dev/null +++ b/test/utils/command_line_stream.cpp @@ -0,0 +1,68 @@ +#if __has_include() +#include +#else +#include +#endif + +#include + +#include +#include +#include + +#include "util/command_line_stream.hpp" + +namespace { + +struct StreamResult { + std::vector lines; + std::optional exit_code; + bool timed_out = false; +}; + +auto run_stream_command(const std::string& cmd) -> StreamResult { + StreamResult result; + auto loop = Glib::MainLoop::create(); + + auto timeout = Glib::signal_timeout().connect( + [&]() { + result.timed_out = true; + loop->quit(); + return false; + }, + 3000); + + waybar::util::command::LineStream stream( + "", [&](const std::string& line) { result.lines.push_back(line); }, + [&](int exit_code) { + result.exit_code = exit_code; + loop->quit(); + }); + + stream.start(cmd); + loop->run(); + timeout.disconnect(); + return result; +} + +} // namespace + +TEST_CASE("command::LineStream emits complete lines and flushes trailing output", + "[util][command_line_stream]") { + const auto result = run_stream_command("printf 'first\\nsecond'"); + + REQUIRE_FALSE(result.timed_out); + REQUIRE(result.exit_code.has_value()); + REQUIRE(*result.exit_code == 0); + REQUIRE(result.lines == std::vector{"first", "second"}); +} + +TEST_CASE("command::LineStream does not emit an extra line after newline-terminated output", + "[util][command_line_stream]") { + const auto result = run_stream_command("printf 'first\\nsecond\\n'"); + + REQUIRE_FALSE(result.timed_out); + REQUIRE(result.exit_code.has_value()); + REQUIRE(*result.exit_code == 0); + REQUIRE(result.lines == std::vector{"first", "second"}); +} diff --git a/test/utils/fixtures/GlibTestsFixture.hpp b/test/utils/fixtures/GlibTestsFixture.hpp index a21c8e07..95ce7f95 100644 --- a/test/utils/fixtures/GlibTestsFixture.hpp +++ b/test/utils/fixtures/GlibTestsFixture.hpp @@ -6,10 +6,16 @@ class GlibTestsFixture : public sigc::trackable { public: GlibTestsFixture() : main_loop_{Glib::MainLoop::create()} {} + ~GlibTestsFixture() { timeout_.disconnect(); } void setTimeout(int timeout) { - Glib::signal_timeout().connect_once([]() { throw std::runtime_error("Test timed out"); }, - timeout); + timeout_.disconnect(); + timeout_ = Glib::signal_timeout().connect( + []() { + throw std::runtime_error("Test timed out"); + return false; + }, + timeout); } void run(std::function fn) { @@ -21,4 +27,5 @@ class GlibTestsFixture : public sigc::trackable { protected: Glib::RefPtr main_loop_; + sigc::connection timeout_; }; diff --git a/test/utils/meson.build b/test/utils/meson.build index e8dd37fa..0d9130de 100644 --- a/test/utils/meson.build +++ b/test/utils/meson.build @@ -15,8 +15,10 @@ test_src = files( 'SafeSignal.cpp', 'sleeper_thread.cpp', 'command.cpp', + 'command_line_stream.cpp', 'css_reload_helper.cpp', '../../src/util/css_reload_helper.cpp', + '../../src/util/command_line_stream.cpp', ) if tz_dep.found()