LLVM: Prevent crash on disk space shortage

Utilities/JIT.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "util/types.hpp"
+#include "util/atomic.hpp"
 
 // Include asmjit with warnings ignored
 #define ASMJIT_EMBED
@@ -508,6 +509,9 @@ class jit_compiler final
 	// Arch
 	std::string m_cpu{};
 
+	// Disk Space left
+	atomic_t<usz> m_disk_space = umax;
+
 public:
 	jit_compiler(const std::unordered_map<std::string, u64>& _link, const std::string& _cpu, u32 flags = 0);
 	~jit_compiler();
@@ -530,7 +534,7 @@ public:
 	void add(std::unique_ptr<llvm::Module> _module);
 
 	// Add object (path to obj file)
-	void add(const std::string& path);
+	bool add(const std::string& path);
 
 	// Update global mapping for a single value
 	void update_global_mapping(const std::string& name, u64 addr);
@@ -552,6 +556,8 @@ public:
 
 	// Get system triple (SPU)
 	static std::string triple2();
+
+	bool add_sub_disk_space(ssz space);
 };
 
 #endif // LLVM_AVAILABLE

Utilities/JITLLVM.cpp

@@ -320,10 +320,12 @@ struct MemoryManager2 : llvm::RTDyldMemoryManager
 class ObjectCache final : public llvm::ObjectCache
 {
 	const std::string& m_path;
+	const std::add_pointer_t<jit_compiler> m_compiler = nullptr;
 
 public:
-	ObjectCache(const std::string& path)
+	ObjectCache(const std::string& path, jit_compiler* compiler = nullptr)
 		: m_path(path)
+		, m_compiler(compiler)
 	{
 	}
 
@@ -332,10 +334,19 @@ public:
 	void notifyObjectCompiled(const llvm::Module* _module, llvm::MemoryBufferRef obj) override
 	{
 		std::string name = m_path;
+
 		name.append(_module->getName().data());
 		//fs::file(name, fs::rewrite).write(obj.getBufferStart(), obj.getBufferSize());
 		name.append(".gz");
 
+		if (!obj.getBufferSize())
+		{
+			jit_log.error("LLVM: Nothing to write: %s", name);
+			return;
+		}
+
+		ensure(m_compiler);
+
 		fs::file module_file(name, fs::rewrite);
 
 		if (!module_file)
@@ -344,6 +355,15 @@ public:
 			return;
 		}
 
+		// Bold assumption about upper limit of space consumption
+		const usz max_size = obj.getBufferSize() * 4;
+
+		if (!m_compiler->add_sub_disk_space(0 - max_size))
+		{
+			jit_log.error("LLVM: Failed to create module file: %s (not enough disk space left)", name);
+			return;
+		}
+
 		if (!zip(obj.getBufferStart(), obj.getBufferSize(), module_file))
 		{
 			jit_log.error("LLVM: Failed to compress module: %s", _module->getName().data());
@@ -353,6 +373,9 @@ public:
 		}
 
 		jit_log.notice("LLVM: Created module: %s", _module->getName().data());
+
+		// Restore space that was overestimated
+		ensure(m_compiler->add_sub_disk_space(max_size - module_file.size()));
 	}
 
 	static std::unique_ptr<llvm::MemoryBuffer> load(const std::string& path)
@@ -388,7 +411,6 @@ public:
 
 			auto buf = llvm::WritableMemoryBuffer::getNewUninitMemBuffer(cached.size());
 			cached.read(buf->getBufferStart(), buf->getBufferSize());
-
 			return buf;
 		}
 
@@ -510,6 +532,26 @@ std::string jit_compiler::triple2()
 #endif
 }
 
+bool jit_compiler::add_sub_disk_space(ssz space)
+{
+	if (space >= 0)
+	{
+		ensure(m_disk_space.fetch_add(space) < ~static_cast<usz>(space));
+		return true;
+	}
+
+	return m_disk_space.fetch_op([sub_size = static_cast<usz>(0 - space)](usz& val)
+	{
+		if (val >= sub_size)
+		{
+			val -= sub_size;
+			return true;
+		}
+
+		return false;
+	}).second;
+}
+
 jit_compiler::jit_compiler(const std::unordered_map<std::string, u64>& _link, const std::string& _cpu, u32 flags)
 	: m_context(new llvm::LLVMContext)
 	, m_cpu(cpu(_cpu))
@@ -576,6 +618,13 @@ jit_compiler::jit_compiler(const std::unordered_map<std::string, u64>& _link, co
 	{
 		fmt::throw_exception("LLVM: Failed to create ExecutionEngine: %s", result);
 	}
+
+	fs::device_stat stats{};
+
+	if (fs::statfs(fs::get_cache_dir(), stats))
+	{
+		m_disk_space = stats.avail_free / 4;
+	}
 }
 
 jit_compiler::~jit_compiler()
@@ -584,7 +633,7 @@ jit_compiler::~jit_compiler()
 
 void jit_compiler::add(std::unique_ptr<llvm::Module> _module, const std::string& path)
 {
-	ObjectCache cache{path};
+	ObjectCache cache{path, this};
 	m_engine->setObjectCache(&cache);
 
 	const auto ptr = _module.get();
@@ -612,23 +661,26 @@ void jit_compiler::add(std::unique_ptr<llvm::Module> _module)
 	}
 }
 
-void jit_compiler::add(const std::string& path)
+bool jit_compiler::add(const std::string& path)
 {
 	auto cache = ObjectCache::load(path);
 
 	if (!cache)
 	{
 		jit_log.error("ObjectCache: Failed to read file. (path='%s', error=%s)", path, fs::g_tls_error);
-		return;		
+		return false;		
 	}
 
 	if (auto object_file = llvm::object::ObjectFile::createObjectFile(*cache))
 	{
 		m_engine->addObjectFile(llvm::object::OwningBinary<llvm::object::ObjectFile>(std::move(*object_file), std::move(cache)));
+		jit_log.trace("ObjectCache: Successfully added %s", path);
+		return true;
 	}
 	else
 	{
 		jit_log.error("ObjectCache: Adding failed: %s", path);
+		return false;
 	}
 }
 

rpcs3/Emu/Cell/PPUThread.cpp

@@ -5011,6 +5011,7 @@ bool ppu_initialize(const ppu_module& info, bool check_only, u64 file_size)
 		g_watchdog_hold_ctr--;
 	}
 
+	bool failed_to_load = false;
 	{
 		if (!is_being_used_in_emulation || (cpu ? cpu->state.all_of(cpu_flag::exit) : Emu.IsStopped()))
 		{
@@ -5030,7 +5031,21 @@ bool ppu_initialize(const ppu_module& info, bool check_only, u64 file_size)
 				break;
 			}
 
-			jit->add(cache_path + obj_name);
+			if (!failed_to_load && !jit->add(cache_path + obj_name))
+			{
+				ppu_log.error("LLVM: Failed to load module %s", obj_name);
+				failed_to_load = true;
+			}
+
+			if (failed_to_load)
+			{
+				if (!is_compiled)
+				{
+					g_progr_pdone++;
+				}
+
+				continue;
+			}
 
 			if (!is_compiled)
 			{
@@ -5040,7 +5055,7 @@ bool ppu_initialize(const ppu_module& info, bool check_only, u64 file_size)
 		}
 	}
 
-	if (!is_being_used_in_emulation || (cpu ? cpu->state.all_of(cpu_flag::exit) : Emu.IsStopped()))
+	if (failed_to_load || !is_being_used_in_emulation || (cpu ? cpu->state.all_of(cpu_flag::exit) : Emu.IsStopped()))
 	{
 		return compiled_new;
 	}
@@ -5073,6 +5088,7 @@ bool ppu_initialize(const ppu_module& info, bool check_only, u64 file_size)
 	if (is_first)
 	{
 		jit_mod.symbol_resolver = reinterpret_cast<void(*)(u8*, u64)>(jit->get("__resolve_symbols"));
+		ensure(jit_mod.symbol_resolver);
 	}
 	else
 	{

rpcs3/util/types.hpp

@@ -141,6 +141,7 @@ using s8  = std::int8_t;
 using s16 = std::int16_t;
 using s32 = std::int32_t;
 using s64 = std::int64_t;
+using ssz = std::make_signed_t<std::size_t>;
 
 // Get integral type from type size
 template <usz N>