z-bar-qt/Plugins/ZShell/Services/gpu.cpp

#include "gpu.hpp"

#include "sensorslib.hpp"

#include "config/rootnodes.hpp"
#include "config/services.hpp"

#include <qdir.h>
#include <qfile.h>
#include <qdiriterator.h>
#include <qregularexpression.h>

namespace ZShell::services {

using Qt::StringLiterals::operator""_s;

namespace {

QStringList cardFiles(const QString& relative) {
	static const QRegularExpression k_cardRe(u"^card\\d+$"_s);

	QStringList files;
	QDirIterator it(u"/sys/class/drm"_s, QDir::Dirs | QDir::NoDotAndDotDot);
	while (it.hasNext()) {
		const QString path = it.next();
		if (!k_cardRe.match(it.fileName()).hasMatch()) continue;
		const QString file = path + relative;
		if (QFile::exists(file)) files << file;
	}
	return files;
}

std::optional<qreal> readNumber(const QString& path) {
	QFile f(path);
	if (!f.open(QIODevice::ReadOnly | QIODevice::Text)) return std::nullopt;
	bool ok = false;
	const qreal v = f.readAll().trimmed().toDouble(&ok);
	if (!ok) return std::nullopt;
	return v;
}

QString cleanName(QString s) {
	static const QRegularExpression k_noise(
		u"\\(R\\)|\\(TM\\)|Graphics"_s,
		QRegularExpression::CaseInsensitiveOption);
	static const QRegularExpression k_spaces(u"\\s+"_s);
	s.replace(k_noise, {});
	s.replace(k_spaces, u" "_s);
	return s.trimmed();
}

QString parseNvidiaName(const QByteArray& out) {
	const QString first =
		QString::fromUtf8(out).split(u'\n').value(0).trimmed();
	return first.isEmpty() ? QString() : cleanName(first);
}

QString parseGlxinfoName(const QByteArray& out) {
	const QStringList lines = QString::fromUtf8(out).split(u'\n');
	for (const QString& line : lines) {
		const qsizetype idx = line.indexOf(u"Device:"_s);
		if (idx < 0) {
			continue;
		}

		QString rest = line.mid(idx + 7);
		const qsizetype paren = rest.indexOf(u'(');
		if (paren >= 0) {
			rest = rest.left(paren);
		}

		const QString cleaned = cleanName(rest);
		if (!cleaned.isEmpty()) {
			return cleaned;
		}
	}

	return {};
}

QString parseLspciName(const QByteArray& out) {
	static const QRegularExpression k_lineRe(
		u"vga|3d controller|display"_s,
		QRegularExpression::CaseInsensitiveOption);

	const QStringList lines = QString::fromUtf8(out).split(u'\n');
	QString match;
	for (const QString& line : lines) {
		if (k_lineRe.match(line).hasMatch()) {
			match = line;
			break;
		}
	}

	if (match.isEmpty()) {
		return {};
	}

	static const QRegularExpression k_bracketRe(u"\\[([^\\]]+)\\][^\\[]*$"_s);
	const auto bracket = k_bracketRe.match(match);
	if (bracket.hasMatch()) {
		return cleanName(bracket.captured(1));
	}

	static const QRegularExpression k_colonRe(u":\\s+(.+)"_s);
	const auto colon = k_colonRe.match(match);
	if (colon.hasMatch()) {
		return cleanName(colon.captured(1));
	}

	return {};
}

struct NameSource {
	QString program;
	QStringList args;
	QString (*parse)(const QByteArray&);
};

const std::array<NameSource, 3>& nameSources() {
	static const std::array<NameSource, 3> k_sources = {{
		{
			.program = u"nvidia-smi"_s,
			.args = {u"--query-gpu=name"_s, u"--format=csv,noheader"_s},
			.parse = &parseNvidiaName,
		},
		{
			.program = u"glxinfo"_s,
			.args = {u"-B"_s},
			.parse = &parseGlxinfoName,
		},
		{
			.program = u"lspci"_s,
			.args = {},
			.parse = &parseLspciName,
		},
	}};
	return k_sources;
}

constexpr int k_nvidiaSource = 0;
constexpr int k_firstGenericSource = 1;

} // namespace

Gpu::Gpu(QObject* parent) : TickingService(parent) {
	auto* svc = config::ConfigSingleton::instance()->services();
	m_userType = svc->gpuType();
	QObject::connect(svc, &config::Services::gpuTypeChanged, this, [this, svc] {
		const GpuType value = svc->gpuType();
		if (value == m_userType) return;
		m_userType = value;
		resolveGpu();
	});

	resolveGpu();

	m_restartTimer = new QTimer(this);
	m_restartTimer->setSingleShot(true);
	m_restartTimer->setInterval(500);
	QObject::connect(m_restartTimer, &QTimer::timeout, this, [this] {
		if (!m_nvidiaProc || m_type != GpuType::Nvidia) return;

		stopNvidiaUsage();
		startNvidiaUsage();
	});

	QObject::connect(this, &TickingService::updateIntervalChanged, this, [this] {
		if (!m_nvidiaProc) return;
		m_restartTimer->start();
	});
}

Gpu::~Gpu() {
	if (!m_nvidiaProc) return;
	m_nvidiaProc->disconnect(this);
	m_nvidiaProc->terminate();
	if (!m_nvidiaProc->waitForFinished(500)) {
		m_nvidiaProc->kill();
		m_nvidiaProc->waitForFinished(200);
	}
}

GpuType Gpu::type() const {
	return m_type;
}

QString Gpu::name() const {
	return m_name;
}

qreal Gpu::percentage() const {
	return m_percentage;
}

qreal Gpu::temperature() const {
	return m_temperature;
}

qreal Gpu::memoryUsed() const {
	return m_memoryUsed;
}

qreal Gpu::memoryTotal() const {
	return m_memoryTotal;
}

void Gpu::setType(GpuType value) {
	if (value == m_type) return;

	m_type = value;
	if (m_type != GpuType::Nvidia) stopNvidiaUsage();
	if (m_type == GpuType::None) resetUsage();
	emit typeChanged();
}

void Gpu::setName(QString value) {
	if (value == m_name) {
		return;
	}
	m_name = std::move(value);
	Q_EMIT nameChanged();
}

void Gpu::setDetecting(bool value) {
	if (value == m_detecting) return;
	m_detecting = value;
	emit detectingChanged();
}

void Gpu::setMemoryUsed(qreal value) {
	if (std::abs(m_memoryUsed - value) < 0.001) return;
	m_memoryUsed = value;
	Q_EMIT memoryUsedChanged();
}

void Gpu::setMemoryTotal(qreal value) {
	if (std::abs(m_memoryTotal - value) < 0.001) return;
	m_memoryTotal = value;
	Q_EMIT memoryTotalChanged();
}

void Gpu::tick() {
	if (m_type == GpuType::Generic) {
		readGenericUsage();
		readGenericMemory();
		readGpuTemperature();
	} else if (m_type == GpuType::Nvidia) {
		startNvidiaUsage();
	} else {
		resetUsage();
	}
}

void Gpu::refreshSysfsCache() {
	m_busyFiles = cardFiles(u"/device/gpu_busy_percent"_s);

	m_vramFiles.clear();
	const QStringList totals = cardFiles(u"/device/mem_info_vram_total"_s);
	for (const QString& total : totals) {
		QString used = total;
		used.replace(u"mem_info_vram_total"_s, u"mem_info_vram_used"_s);
		if (QFile::exists(used)) m_vramFiles.append({used, total});
	}
}

void Gpu::resolveGpu() {
	refreshSysfsCache();
	const int generation = ++m_generation;

	if (m_userType != GpuType::Auto) setType(m_userType);

	if (m_userType == GpuType::None) {
		setName({});
		setDetecting(false);
		return;
	}

	setName({});
	setDetecting(true);
	tryNameSource(
		m_userType == GpuType::Generic ? k_firstGenericSource : k_nvidiaSource,
		generation);
}

int Gpu::probeEnd() const {
	return m_userType == GpuType::Nvidia
			   ? k_firstGenericSource
			   : static_cast<int>(nameSources().size());
}

void Gpu::tryNameSource(int index, int generation) {
	const NameSource& src = nameSources().at(static_cast<std::size_t>(index));
	runProcess(
		src.program,
		src.args,
		[this, index, generation, parse = src.parse](const QByteArray& out) {
			finishNameSource(index, generation, parse(out));
		});
}

void Gpu::finishNameSource(int index, int generation, QString name) {
	if (generation != m_generation) {
		return;
	}

	if (m_userType == GpuType::Auto && index == k_nvidiaSource) {
		if (!name.isEmpty())
			setType(GpuType::Nvidia);
		else
			setType(m_busyFiles.isEmpty() ? GpuType::None : GpuType::Generic);

		if (m_type == GpuType::None) {
			setName({});
			setDetecting(false);
			return;
		}
	}

	if (!name.isEmpty()) {
		setName(std::move(name));
		setDetecting(false);
		return;
	}

	const int next = index + 1;
	if (next < probeEnd()) {
		tryNameSource(next, generation);
	} else {
		setName({});
		setDetecting(false);
	}
}

void Gpu::runProcess(
	const QString& program,
	const QStringList& args,
	std::function<void(const QByteArray&)> callback) {
	auto* proc = new QProcess(this);
	proc->setStandardErrorFile(QProcess::nullDevice());

	const auto finish =
		[proc, callback = std::move(callback)](const QByteArray& out) {
			callback(out);
			proc->deleteLater();
		};

	QObject::connect(
		proc,
		&QProcess::finished,
		this,
		[finish, proc](int code, QProcess::ExitStatus status) {
			const bool ok = status == QProcess::NormalExit && code == 0;
			finish(ok ? proc->readAllStandardOutput() : QByteArray());
		});
	QObject::connect(
		proc,
		&QProcess::errorOccurred,
		this,
		[finish](QProcess::ProcessError err) {
			if (err == QProcess::FailedToStart) {
				finish(QByteArray());
			}
		});

	proc->start(program, args);
}

void Gpu::onStopped() {
	stopNvidiaUsage();
}

void Gpu::readGenericMemory() {
	qreal total = 0.0;
	qreal used = 0.0;
	for (const VramFiles& v : std::as_const(m_vramFiles)) {
		const auto t = readNumber(v.total);
		const auto u = readNumber(v.used);
		if (!t || !u) continue;
		total += *t;
		used += *u;
	}
	setMemoryTotal(total);
	setMemoryUsed(used);
}

void Gpu::readGenericUsage() {
	qreal sum = 0.0;
	int count = 0;
	for (const QString& path : std::as_const(m_busyFiles)) {
		if (const auto v = readNumber(path)) {
			sum += *v;
			++count;
		}
	}
	const qreal newPerc = count > 0 ? sum / count / 100.0 : 0.0;
	if (std::abs(newPerc - m_percentage) > 0.0001) {
		m_percentage = newPerc;
		emit percentageChanged();
	}
}

void Gpu::stopNvidiaUsage() {
	m_restartTimer->stop();

	if (!m_nvidiaProc) return;

	QProcess* proc = m_nvidiaProc;
	m_nvidiaProc = nullptr;

	proc->disconnect(this);
	if (proc->state() == QProcess::NotRunning) {
		proc->deleteLater();
		return;
	}
	QObject::connect(proc, &QProcess::finished, proc, &QObject::deleteLater);
	proc->terminate();
}

void Gpu::startNvidiaUsage() {
	if (m_nvidiaProc) return;

	auto* proc = new QProcess(this);
	m_nvidiaProc = proc;

	QObject::connect(
		proc, &QProcess::readyReadStandardOutput, this, [this, proc]() {
			while (proc->canReadLine()) {
				const QByteArray out = proc->readLine();
				const QString output = QString::fromUtf8(out).trimmed();
				if (output.isEmpty()) continue;

				const QList<QString> parts = output.split(',');
				if (parts.size() < 4) return;

				bool ok1 = false;
				bool ok2 = false;
				bool ok3 = false;
				bool ok4 = false;
				const qreal usage =
					parts.at(0).trimmed().toDouble(&ok1) / 100.0;
				const qreal temp = parts.at(1).trimmed().toDouble(&ok2);
				const qreal memUsed =
					parts.at(2).trimmed().toDouble(&ok3) * 1024;
				const qreal memTotal =
					parts.at(3).trimmed().toDouble(&ok4) * 1024;

				if (ok1 && std::abs(usage - m_percentage) > 0.0001) {
					m_percentage = usage;
					Q_EMIT percentageChanged();
				}
				if (ok2 && std::abs(temp - m_temperature) > 0.05) {
					m_temperature = temp;
					Q_EMIT temperatureChanged();
				}
				if (ok3) {
					setMemoryUsed(memUsed);
				}
				if (ok4) {
					setMemoryTotal(memTotal);
				}
			}
		});

	auto cleanup = [this, proc]() {
		if (m_nvidiaProc == proc) m_nvidiaProc = nullptr;

		proc->deleteLater();
	};

	QObject::connect(proc, &QProcess::finished, this, cleanup);

	QObject::connect(
		proc,
		&QProcess::errorOccurred,
		this,
		[cleanup](QProcess::ProcessError) { cleanup(); });

	proc->start(
		QStringLiteral("nvidia-smi"),
		{QStringLiteral(
			 "--query-gpu=utilization.gpu,temperature.gpu,"
			 "memory.used,memory.total"),
		 QStringLiteral("--format=csv,noheader,nounits"),
		 QStringLiteral("-lms"),
		 QString::number(updateInterval())});
}

void Gpu::readGpuTemperature() {
	const auto t = sensorslib::gpuPciAverageTemp();
	const qreal newTemp = t.value_or(0.0);
	if (std::abs(newTemp - m_temperature) > 0.05) {
		m_temperature = newTemp;
		emit temperatureChanged();
	}
}

void Gpu::resetUsage() {
	if (std::abs(m_percentage) > 0.0001) {
		m_percentage = 0.0;
		emit percentageChanged();
	}
	if (std::abs(m_temperature) > 0.05) {
		m_temperature = 0.0;
		emit temperatureChanged();
	}
	setMemoryUsed(0.0);
	setMemoryTotal(0.0);
}

} // namespace ZShell::services