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

#include "storage.hpp"

#include <algorithm>
#include <cmath>
#include <qdir.h>
#include <qfile.h>
#include <qfileinfo.h>
#include <qfuturewatcher.h>
#include <qhash.h>
#include <qloggingcategory.h>
#include <qstorageinfo.h>
#include <qthreadpool.h>
#include <qtconcurrentrun.h>
#include <sys/stat.h>
#include <sys/sysmacros.h>

Q_LOGGING_CATEGORY(lcStorage, "ZShell.services.storage", QtInfoMsg)

namespace ZShell::services {

namespace {

// A scan that finishes quickly is not worth showing as loading. Only a scan that
// stays in flight this long is reported, so the UI does not flicker every tick.
constexpr int kLoadingDelayMs = 300;

struct Accum {
	quint64 usedBytes = 0;
	quint64 totalBytes = 0;
	bool hasRoot = false;
};

QThreadPool* storagePool() {
	// Single thread: a scan blocked on an unreachable mount should not take a
	// worker away from the global pool.
	static QThreadPool* pool = [] {
		auto* p = new QThreadPool();
		p->setMaxThreadCount(1);
		return p;
	}();
	return pool;
}

[[nodiscard]] QString sysfsRealPath(uint major, uint minor) {
	const QString link =
		QStringLiteral("/sys/dev/block/%1:%2").arg(major).arg(minor);
	const QString resolved = QFileInfo(link).canonicalFilePath();
	return resolved;
}

[[nodiscard]] bool readDevtFromSysfs(
	const QString& sysfsBlockDir, uint& major, uint& minor) {
	QFile f(sysfsBlockDir + QStringLiteral("/dev"));
	if (!f.open(QIODevice::ReadOnly | QIODevice::Text)) {
		return false;
	}
	const QByteArray line = f.readLine().trimmed();
	f.close();

	const qsizetype colon = line.indexOf(':');
	if (colon <= 0) {
		return false;
	}
	bool okM = false;
	bool okN = false;
	major = line.left(colon).toUInt(&okM);
	minor = line.mid(colon + 1).toUInt(&okN);
	return okM && okN;
}

QStringList resolveByDevt(uint major, uint minor, int depth = 0);

// NOLINTNEXTLINE(misc-no-recursion)
QStringList resolveAtNode(const QString& node, int depth) {
	if (node.isEmpty() || depth > 8) {
		return {};
	}

	const QFileInfo nodeInfo(node);
	if (!nodeInfo.exists() || !nodeInfo.isDir()) {
		return {};
	}

	if (QFileInfo::exists(node + QStringLiteral("/partition"))) {
		const QString diskNode = nodeInfo.path();
		return {QFileInfo(diskNode).fileName()};
	}

	const QDir slavesDir(node + QStringLiteral("/slaves"));
	if (slavesDir.exists()) {
		const QStringList slaves =
			slavesDir.entryList(QDir::Dirs | QDir::NoDotAndDotDot);
		if (!slaves.isEmpty()) {
			QStringList out;
			for (const QString& slave : slaves) {
				uint sm = 0;
				uint sn = 0;
				const QString slaveDir =
					QStringLiteral("/sys/class/block/") + slave;
				if (!readDevtFromSysfs(slaveDir, sm, sn)) {
					continue;
				}
				const auto devs = resolveByDevt(sm, sn, depth + 1);
				for (const QString& d : devs) {
					if (!out.contains(d)) {
						out.append(d);
					}
				}
			}
			return out;
		}
	}

	return {nodeInfo.fileName()};
}

// NOLINTNEXTLINE(misc-no-recursion)
QStringList resolveByDevt(uint major, uint minor, int depth) {
	return resolveAtNode(sysfsRealPath(major, minor), depth);
}

} // namespace

Storage::Storage(QObject* parent) : TickingService(parent) {
	m_loadingTimer.setSingleShot(true);
	QObject::connect(&m_loadingTimer, &QTimer::timeout, this, [this] {
		if (m_watcher) {
			setLoading(true);
		}
	});
}

qreal Storage::percentage() const {
	qreal totalUsed = 0.0;
	qreal totalSize = 0.0;
	for (const DiskInfo* d : m_disks) {
		totalUsed += d->used();
		totalSize += d->total();
	}
	return totalSize > 0.0 ? totalUsed / totalSize : 0.0;
}

bool Storage::loading() const {
	return m_loading;
}

void Storage::setLoading(bool value) {
	if (m_loading == value) {
		return;
	}
	m_loading = value;
	Q_EMIT loadingChanged();
}

bool Storage::sameOrder(const QList<DiskInfo*>& a, const QList<DiskInfo*>& b) {
	if (a.size() != b.size()) {
		return false;
	}
	for (qsizetype i = 0; i < a.size(); ++i) {
		if (a.at(i) != b.at(i)) {
			return false;
		}
	}
	return true;
}

QQmlListProperty<DiskInfo> Storage::disksProp() {
	return QQmlListProperty<DiskInfo>(
		this, nullptr, &Storage::disksCount, &Storage::disksAt);
}

qsizetype Storage::disksCount(QQmlListProperty<DiskInfo>* prop) {
	return static_cast<Storage*>(prop->object)->m_disks.size();
}

DiskInfo* Storage::disksAt(QQmlListProperty<DiskInfo>* prop, qsizetype i) {
	return static_cast<Storage*>(prop->object)->m_disks.at(i);
}

DiskInfo* Storage::manualPrimaryDisk() const {
	return m_manualPrimaryDisk.data();
}

void Storage::setManualPrimaryDisk(DiskInfo* disk) {
	if (m_manualPrimaryDisk.data() == disk) {
		return;
	}
	m_manualPrimaryDisk = disk;
	Q_EMIT manualPrimaryDiskChanged();
	Q_EMIT primaryDiskChanged();
}

DiskInfo* Storage::primaryDisk() const {
	if (auto* m = m_manualPrimaryDisk.data()) {
		return m;
	}
	return m_disks.isEmpty() ? nullptr : m_disks.first();
}

bool Storage::isPseudoFs(QByteArrayView fsType) {
	static constexpr const char* kPseudo[] = {
		"tmpfs",
		"devtmpfs",
		"proc",
		"sysfs",
		"cgroup",
		"cgroup2",
		"overlay",
		"squashfs",
		"devpts",
		"mqueue",
		"ramfs",
		"rpc_pipefs",
		"autofs",
		"configfs",
		"debugfs",
		"tracefs",
		"securityfs",
		"pstore",
		"bpf",
		"binfmt_misc",
		"hugetlbfs",
		"fusectl",
		"efivarfs",
		"selinuxfs",
		// Real FUSE mounts (rclone, sshfs, ...) are reported like any other
		// volume, so only the virtual ones are excluded here.
		"fuse.gvfsd-fuse",
		"fuse.portal",
	};
	for (const char* p : kPseudo) {
		if (fsType == QByteArrayView(p)) {
			return true;
		}
	}
	return false;
}

QStringList Storage::resolveToPhysicalDisks(const QString& devicePath) {
	if (devicePath.isEmpty() || !devicePath.startsWith(QLatin1Char('/'))) {
		return {};
	}
	struct stat st{};
	if (::stat(devicePath.toLocal8Bit().constData(), &st) != 0) {
		return {};
	}
	if (!S_ISBLK(st.st_mode)) {
		return {};
	}
	return resolveByDevt(major(st.st_rdev), minor(st.st_rdev));
}

StorageScan Storage::scanVolumes() {
	struct DeviceEntry {
		quint64 totalBytes = 0;
		quint64 usedBytes = 0;
		bool hasRoot = false;
		// Mount point of the first mount seen for this device. Used as the entry
		// key when the device has no backing block device (network mount).
		QString mount;
	};

	QHash<QByteArray, DeviceEntry> byDevice;

	const auto mountedVols = QStorageInfo::mountedVolumes();
	for (const QStorageInfo& v : mountedVols) {
		if (!v.isReady() || !v.isValid() || v.bytesTotal() <= 0) {
			continue;
		}
		if (isPseudoFs(QByteArrayView(v.fileSystemType()))) {
			continue;
		}

		const auto totalBytes = static_cast<quint64>(v.bytesTotal());
		const auto availBytes = static_cast<quint64>(v.bytesAvailable());
		const quint64 usedBytes =
			totalBytes > availBytes ? totalBytes - availBytes : 0;
		const bool isRoot = v.rootPath() == QStringLiteral("/");

		DeviceEntry& e = byDevice[v.device()];
		e.totalBytes = totalBytes;
		e.usedBytes = usedBytes;
		e.hasRoot = e.hasRoot || isRoot;
		if (e.mount.isEmpty()) {
			e.mount = v.rootPath();
		}
	}

	QHash<QString, Accum> byEntry;

	// Multiple mounts can share a single backing filesystem (btrfs subvolumes,
	// bind mounts, etc.) and each one reports identical bytesTotal/bytesAvailable.
	// Dedupe by source device so the filesystem only contributes once per disk.
	for (auto it = byDevice.constBegin(); it != byDevice.constEnd(); ++it) {
		const DeviceEntry& e = it.value();
		const QStringList disks =
			resolveToPhysicalDisks(QString::fromLocal8Bit(it.key()));
		if (disks.isEmpty()) {
			// Network mounts (and any other filesystem with no backing block
			// device) have no physical disk to attribute the usage to, so the
			// mount itself becomes the entry.
			Accum& a = byEntry[e.mount];
			a.usedBytes += e.usedBytes;
			a.totalBytes += e.totalBytes;
			a.hasRoot = a.hasRoot || e.hasRoot;
			continue;
		}
		for (const QString& d : disks) {
			if (d.startsWith(QStringLiteral("zram"))) {
				continue;
			}
			Accum& a = byEntry[d];
			a.usedBytes += e.usedBytes;
			a.totalBytes += e.totalBytes;
			a.hasRoot = a.hasRoot || e.hasRoot;
		}
	}

	StorageScan scan;
	scan.reserve(byEntry.size());
	for (auto it = byEntry.constBegin(); it != byEntry.constEnd(); ++it) {
		StorageEntry entry;
		entry.mount = it.key();
		entry.usedBytes = it.value().usedBytes;
		entry.totalBytes = it.value().totalBytes;
		entry.hasRoot = it.value().hasRoot;
		scan.append(entry);
	}

	std::sort(
		scan.begin(),
		scan.end(),
		[](const StorageEntry& a, const StorageEntry& b) {
			if (a.hasRoot != b.hasRoot) {
				return a.hasRoot;
			}
			return a.mount < b.mount;
		});

	return scan;
}

void Storage::tick() {
	if (m_watcher) {
		// Still waiting on a mount that has not answered. Starting another scan
		// here would just queue blocked threads behind the first.
		return;
	}

	beginScan();
}

void Storage::beginScan() {
	m_loadingTimer.start(kLoadingDelayMs);

	m_watcher = new QFutureWatcher<StorageScan>(this);
	QObject::connect(
		m_watcher,
		&QFutureWatcher<StorageScan>::finished,
		this,
		&Storage::onScanFinished);

	m_watcher->setFuture(
		QtConcurrent::run(storagePool(), &Storage::scanVolumes));
}

void Storage::onScanFinished() {
	QFutureWatcher<StorageScan>* watcher = m_watcher;
	m_watcher = nullptr;
	if (!watcher) {
		return;
	}
	StorageScan scan = watcher->result();
	watcher->deleteLater();

	m_loadingTimer.stop();

	// scan is already ordered (root disk first), so applying it in order keeps the
	// QML model in the same order as before.
	applyScan(scan);

	setLoading(false);
}

void Storage::applyScan(const StorageScan& scan) {
	const qreal prevPercentage = percentage();

	QHash<QString, DiskInfo*> existing;
	existing.reserve(m_disks.size());
	for (DiskInfo* d : std::as_const(m_disks)) {
		existing.insert(d->mount(), d);
	}

	QList<DiskInfo*> next;
	next.reserve(scan.size());
	for (const StorageEntry& e : scan) {
		if (DiskInfo* survivor = existing.take(e.mount)) {
			survivor->update(e.usedBytes, e.totalBytes, e.hasRoot);
			next.append(survivor);
		} else {
			next.append(new DiskInfo(
				e.mount, e.usedBytes, e.totalBytes, e.hasRoot, this));
		}
	}

	bool manualCleared = false;
	if (DiskInfo* m = m_manualPrimaryDisk.data();
		m && existing.contains(m->mount())) {
		m_manualPrimaryDisk.clear();
		manualCleared = true;
	}
	for (DiskInfo* stale : std::as_const(existing)) {
		stale->deleteLater();
	}

	const bool listChanged = !sameOrder(m_disks, next);
	DiskInfo* prevPrimary = primaryDisk();
	m_disks = next;

	if (listChanged) {
		Q_EMIT disksChanged();
	}
	if (std::abs(percentage() - prevPercentage) > 0.0001) {
		Q_EMIT percentageChanged();
	}
	if (manualCleared) {
		Q_EMIT manualPrimaryDiskChanged();
	}
	if (primaryDisk() != prevPrimary) {
		Q_EMIT primaryDiskChanged();
	}
}

} // namespace ZShell::services