From 62d80b97e65237539a103ded87f4e650ddafe4b8 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Thu, 6 Feb 2025 03:14:37 +0000
Subject: [PATCH] add systemd unit logging mode

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/core/log/console.rs      | 77 +++++++++++++++++++++++++++++++++---
 src/core/log/mod.rs          |  4 +-
 src/main/logging.rs          |  4 +-
 src/service/admin/console.rs |  5 ++-
 4 files changed, 78 insertions(+), 12 deletions(-)

diff --git a/src/core/log/console.rs b/src/core/log/console.rs
index 0bc44fa7..1f04ba26 100644
--- a/src/core/log/console.rs
+++ b/src/core/log/console.rs
@@ -1,3 +1,5 @@
+use std::{env, io, sync::LazyLock};
+
 use tracing::{
 	field::{Field, Visit},
 	Event, Level, Subscriber,
@@ -7,12 +9,59 @@ use tracing_subscriber::{
 	fmt,
 	fmt::{
 		format::{Compact, DefaultVisitor, Format, Full, Pretty, Writer},
-		FmtContext, FormatEvent, FormatFields,
+		FmtContext, FormatEvent, FormatFields, MakeWriter,
 	},
 	registry::LookupSpan,
 };
 
-use crate::{Config, Result};
+use crate::{apply, Config, Result};
+
+static SYSTEMD_MODE: LazyLock<bool> =
+	LazyLock::new(|| env::var("SYSTEMD_EXEC_PID").is_ok() && env::var("JOURNAL_STREAM").is_ok());
+
+pub struct ConsoleWriter {
+	stdout: io::Stdout,
+	stderr: io::Stderr,
+	_journal_stream: [u64; 2],
+	use_stderr: bool,
+}
+
+impl ConsoleWriter {
+	#[must_use]
+	pub fn new(_config: &Config) -> Self {
+		let journal_stream = get_journal_stream();
+		Self {
+			stdout: io::stdout(),
+			stderr: io::stderr(),
+			_journal_stream: journal_stream.into(),
+			use_stderr: journal_stream.0 != 0,
+		}
+	}
+}
+
+impl<'a> MakeWriter<'a> for ConsoleWriter {
+	type Writer = &'a Self;
+
+	fn make_writer(&'a self) -> Self::Writer { self }
+}
+
+impl io::Write for &'_ ConsoleWriter {
+	fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+		if self.use_stderr {
+			self.stderr.lock().write(buf)
+		} else {
+			self.stdout.lock().write(buf)
+		}
+	}
+
+	fn flush(&mut self) -> io::Result<()> {
+		if self.use_stderr {
+			self.stderr.lock().flush()
+		} else {
+			self.stdout.lock().flush()
+		}
+	}
+}
 
 pub struct ConsoleFormat {
 	_compact: Format<Compact>,
@@ -20,10 +69,6 @@ pub struct ConsoleFormat {
 	pretty: Format<Pretty>,
 }
 
-struct ConsoleVisitor<'a> {
-	visitor: DefaultVisitor<'a>,
-}
-
 impl ConsoleFormat {
 	#[must_use]
 	pub fn new(config: &Config) -> Self {
@@ -68,6 +113,10 @@ where
 	}
 }
 
+struct ConsoleVisitor<'a> {
+	visitor: DefaultVisitor<'a>,
+}
+
 impl<'writer> FormatFields<'writer> for ConsoleFormat {
 	fn format_fields<R>(&self, writer: Writer<'writer>, fields: R) -> Result<(), std::fmt::Error>
 	where
@@ -92,3 +141,19 @@ impl Visit for ConsoleVisitor<'_> {
 		self.visitor.record_debug(field, value);
 	}
 }
+
+#[must_use]
+fn get_journal_stream() -> (u64, u64) {
+	is_systemd_mode()
+		.then(|| env::var("JOURNAL_STREAM").ok())
+		.flatten()
+		.as_deref()
+		.and_then(|s| s.split_once(':'))
+		.map(apply!(2, str::parse))
+		.map(apply!(2, Result::unwrap_or_default))
+		.unwrap_or((0, 0))
+}
+
+#[inline]
+#[must_use]
+pub fn is_systemd_mode() -> bool { *SYSTEMD_MODE }
diff --git a/src/core/log/mod.rs b/src/core/log/mod.rs
index 0c51a383..0c1840d0 100644
--- a/src/core/log/mod.rs
+++ b/src/core/log/mod.rs
@@ -2,14 +2,14 @@
 
 pub mod capture;
 pub mod color;
-mod console;
+pub mod console;
 pub mod fmt;
 pub mod fmt_span;
 mod reload;
 mod suppress;
 
 pub use capture::Capture;
-pub use console::ConsoleFormat;
+pub use console::{is_systemd_mode, ConsoleFormat, ConsoleWriter};
 pub use reload::{LogLevelReloadHandles, ReloadHandle};
 pub use suppress::Suppress;
 pub use tracing::Level;
diff --git a/src/main/logging.rs b/src/main/logging.rs
index 85945e8a..35e482de 100644
--- a/src/main/logging.rs
+++ b/src/main/logging.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use conduwuit::{
 	config::Config,
 	debug_warn, err,
-	log::{capture, fmt_span, ConsoleFormat, LogLevelReloadHandles},
+	log::{capture, fmt_span, ConsoleFormat, ConsoleWriter, LogLevelReloadHandles},
 	result::UnwrapOrErr,
 	Result,
 };
@@ -30,7 +30,7 @@ pub(crate) fn init(
 		.with_span_events(console_span_events)
 		.event_format(ConsoleFormat::new(config))
 		.fmt_fields(ConsoleFormat::new(config))
-		.map_writer(|w| w);
+		.with_writer(ConsoleWriter::new(config));
 
 	let (console_reload_filter, console_reload_handle) =
 		reload::Layer::new(console_filter.clone());
diff --git a/src/service/admin/console.rs b/src/service/admin/console.rs
index de201f4b..59b9a31b 100644
--- a/src/service/admin/console.rs
+++ b/src/service/admin/console.rs
@@ -1,10 +1,11 @@
 #![cfg(feature = "console")]
+
 use std::{
 	collections::VecDeque,
 	sync::{Arc, Mutex},
 };
 
-use conduwuit::{debug, defer, error, log, Server};
+use conduwuit::{debug, defer, error, log, log::is_systemd_mode, Server};
 use futures::future::{AbortHandle, Abortable};
 use ruma::events::room::message::RoomMessageEventContent;
 use rustyline_async::{Readline, ReadlineError, ReadlineEvent};
@@ -123,7 +124,7 @@ impl Console {
 	}
 
 	async fn readline(self: &Arc<Self>) -> Result<ReadlineEvent, ReadlineError> {
-		let _suppression = log::Suppress::new(&self.server);
+		let _suppression = (!is_systemd_mode()).then(|| log::Suppress::new(&self.server));
 
 		let (mut readline, _writer) = Readline::new(PROMPT.to_owned())?;
 		let self_ = Arc::clone(self);