Merge pull request #4915 from khaneliman/custom-gio

include/modules/custom.hpp

@@ -3,10 +3,12 @@
 #include <fmt/format.h>
 
 #include <csignal>
+#include <memory>
 #include <string>
 
 #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<std::string> class_;
   int percentage_;
-  FILE* fp_;
-  int pid_;
   util::command::res output_;
   util::JsonParser parser_;
+  std::unique_ptr<util::command::LineStream> continuous_stream_;
+  sigc::connection restart_connection_;
 
   util::SleeperThread thread_;
 };

include/util/command_line_stream.hpp

@@ -0,0 +1,40 @@
+#pragma once
+
+#include <glibmm/main.h>
+#include <glibmm/spawn.h>
+
+#include <functional>
+#include <string>
+
+namespace waybar::util::command {
+
+class LineStream {
+ public:
+  using OutputCallback = std::function<void(const std::string&)>;
+  using ExitCallback = std::function<void(int)>;
+
+  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

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(

src/modules/custom.cpp

@@ -2,10 +2,10 @@
 
 #include <spdlog/spdlog.h>
 
+#include <cerrno>
+#include <stdexcept>
 #include <utility>
 
-#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,12 +51,20 @@ 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<pid_t>(*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()) {
@@ -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);
+  continuous_stream_ = std::make_unique<util::command::LineStream>(
+      output_name_,
+      [this](const std::string& output) {
+        output_ = {.exit_code = 0, .out = output};
+        dp.emit();
+      },
+      [this](int exit_code) { handleContinuousProcessExit(exit_code); });
+  startContinuousProcess(true);
 }
-  thread_ = [this, cmd] {
-    char* buff = nullptr;
-    waybar::util::ScopeGuard buff_deleter([&buff]() {
-      if (buff) {
-        free(buff);
+
+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());
     }
-    });
-    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;
+    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, ""};
+    output_ = {.exit_code = exit_code, .out = ""};
     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<long>(config_["restart-interval"].asDouble() * 1000))));
-        fp_ = util::command::open(cmd, pid_, output_name_);
-        if (!fp_) {
-          throw std::runtime_error("Unable to open " + cmd);
+
+  scheduleContinuousRestart();
 }
-      } else {
+
+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".
-        thread_.stop();
     return;
   }
-    } else {
-      std::string output = buff;
 
-      // Remove last newline
-      if (!output.empty() && output[output.length() - 1] == '\n') {
-        output.erase(output.length() - 1);
-      }
-      output_ = {0, output};
-      dp.emit();
-    }
-  };
+  restart_connection_ = Glib::signal_timeout().connect(
+      [this] {
+        startContinuousProcess(false);
+        return false;
+      },
+      std::max(1U, static_cast<unsigned>(config_["restart-interval"].asDouble() * 1000)));
 }
 
 void waybar::modules::Custom::waitingWorker() {

src/util/command_line_stream.cpp

@@ -0,0 +1,219 @@
+#include "util/command_line_stream.hpp"
+
+#include <fcntl.h>
+#include <signal.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+#include <array>
+#include <cerrno>
+#include <cstring>
+#include <stdexcept>
+#include <utility>
+#include <vector>
+
+#ifdef __linux__
+#include <sys/prctl.h>
+#endif
+#ifdef __FreeBSD__
+#include <sys/procctl.h>
+#endif
+
+#include "util/command.hpp"
+
+namespace {
+
+auto buildChildEnvironment(const std::string& output_name) -> std::vector<std::string> {
+  auto names = Glib::listenv();
+  std::vector<std::string> 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<void*>(&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<std::string> 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<bool>(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<char, 4096> chunk = {};
+  while (true) {
+    const auto bytes_read = ::read(stdout_fd_, chunk.data(), chunk.size());
+    if (bytes_read > 0) {
+      buffer_.append(chunk.data(), static_cast<size_t>(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;
+}

test/utils/command_line_stream.cpp

@@ -0,0 +1,68 @@
+#if __has_include(<catch2/catch_test_macros.hpp>)
+#include <catch2/catch_test_macros.hpp>
+#else
+#include <catch2/catch.hpp>
+#endif
+
+#include <glibmm/main.h>
+
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "util/command_line_stream.hpp"
+
+namespace {
+
+struct StreamResult {
+  std::vector<std::string> lines;
+  std::optional<int> 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<std::string>{"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<std::string>{"first", "second"});
+}

test/utils/fixtures/GlibTestsFixture.hpp

@@ -6,9 +6,15 @@
 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_.disconnect();
+    timeout_ = Glib::signal_timeout().connect(
+        []() {
+          throw std::runtime_error("Test timed out");
+          return false;
+        },
         timeout);
   }
 
@@ -21,4 +27,5 @@ class GlibTestsFixture : public sigc::trackable {
 
  protected:
   Glib::RefPtr<Glib::MainLoop> main_loop_;
+  sigc::connection timeout_;
 };

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()