From 9f4b65fca612185d912e2c10abb85c84c9bd4846 Mon Sep 17 00:00:00 2001 From: Jakub Luwierski Date: Sun, 18 Oct 2026 14:24:04 +0200 Subject: [PATCH] ngram: avoid draining nonexistent reads after initialization failures --- CMakeLists.txt | 19 ++++ src/ngram/ple_reader.cpp | 4 +- tests/platform/ple_reader_lifecycle.cpp | 114 ++++++++++++++++++++++++ 3 files changed, 136 insertions(+), 1 deletion(-) create mode 100644 tests/platform/ple_reader_lifecycle.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 7ba7ee6349..4ea4952874 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2140,6 +2140,25 @@ if(STRATA_BUILD_TESTS) target_include_directories(direct_file_async_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include) target_link_libraries(direct_file_async_test PRIVATE Threads::Threads) add_test(NAME direct_file_async_test COMMAND direct_file_async_test) + # Reader initialization failures must return without draining reads that were never submitted. + add_executable(ple_reader_lifecycle_test tests/platform/ple_reader_lifecycle.cpp + src/ngram/ple_reader.cpp src/platform/direct_file.cpp) + target_include_directories(ple_reader_lifecycle_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include) + target_link_libraries(ple_reader_lifecycle_test PRIVATE Threads::Threads strata_warnings) + if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + target_compile_definitions(ple_reader_lifecycle_test PRIVATE STRATA_TEST_ALLOC_FAILURE=1) + target_link_options(ple_reader_lifecycle_test PRIVATE -Wl,--wrap=posix_memalign) + endif() + foreach(_mode caller threaded) + foreach(_case reads reopen) + add_test(NAME ple_reader_${_case}_${_mode} COMMAND ple_reader_lifecycle_test ${_case} ${_mode}) + set_tests_properties(ple_reader_${_case}_${_mode} PROPERTIES TIMEOUT 15) + endforeach() + if(CMAKE_CXX_COMPILER_ID MATCHES "GNU|Clang") + add_test(NAME ple_reader_alloc_${_mode} COMMAND ple_reader_lifecycle_test alloc ${_mode}) + set_tests_properties(ple_reader_alloc_${_mode} PROPERTIES TIMEOUT 15) + endif() + endforeach() endif() # --aux-cpus: the CPU-set derivation (pure) and the placement of live threads; header-only, no GPU add_executable(aux_cpus_test tests/platform/aux_cpus_test.cpp) diff --git a/src/ngram/ple_reader.cpp b/src/ngram/ple_reader.cpp index 9d2a05051e..96351582b3 100644 --- a/src/ngram/ple_reader.cpp +++ b/src/ngram/ple_reader.cpp @@ -347,12 +347,13 @@ bool PleReader::open(const std::string& path, uint64_t table_offset, uint64_t n_ } impl_->table_offset = table_offset; impl_->n_rows = n_rows; - impl_->max_inflight = max_inflight; impl_->slab = (uint8_t*) DirectFile::alloc_aligned((size_t) max_inflight * 2 * PAGE); if (impl_->slab == nullptr) { err = "PleReader: cannot allocate read buffers"; close(); return false; } impl_->inflight.assign(max_inflight, Job{}); impl_->free_slots.clear(); for (uint32_t s = max_inflight; s-- > 0;) impl_->free_slots.push_back(s); + // Publish the slot count only after every slot is free. A failed initialization has no reads to drain. + impl_->max_inflight = max_inflight; impl_->cache.init(cache_rows, row_bytes); impl_->error.clear(); impl_->keep_us = 0; @@ -409,6 +410,7 @@ void PleReader::close() { m.delayed.clear(); m.inflight.clear(); m.free_slots.clear(); + m.max_inflight = 0; m.cache.init(0); m.threaded = false; m.stop = false; diff --git a/tests/platform/ple_reader_lifecycle.cpp b/tests/platform/ple_reader_lifecycle.cpp new file mode 100644 index 0000000000..e1357ec26a --- /dev/null +++ b/tests/platform/ple_reader_lifecycle.cpp @@ -0,0 +1,114 @@ +// CPU-only lifecycle regressions using real DirectFile I/O; allocation failure is injected at link time. +#include "strata/ngram/ple_reader.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifdef STRATA_TEST_ALLOC_FAILURE +static bool fail_alloc = false; +extern "C" int __real_posix_memalign(void**, size_t, size_t); +extern "C" int __wrap_posix_memalign(void** out, size_t alignment, size_t bytes) { + if (fail_alloc) { fail_alloc = false; return ENOMEM; } + return __real_posix_memalign(out, alignment, bytes); +} +#endif + +namespace { +using strata::ngram::PleReader; +constexpr uint32_t kRows = 128, kRowBytes = 90; +constexpr size_t kFileBytes = 16384; +void check(bool ok, const std::string& why) { if (!ok) throw std::runtime_error(why); } +uint8_t pattern(size_t i) { return (uint8_t) ((i * 17 + i / 257) & 255); } + +struct Table { + std::filesystem::path path; + Table() { + std::string name = (std::filesystem::temp_directory_path() / "strata-ple-lifecycle-XXXXXX").string(); + const int fd = mkstemp(name.data()); + check(fd >= 0, "mkstemp failed"); + ::close(fd); + path = name; + std::vector bytes(kFileBytes); + for (size_t i = 0; i < bytes.size(); ++i) bytes[i] = pattern(i); + std::ofstream f(path, std::ios::binary | std::ios::trunc); + f.write(reinterpret_cast(bytes.data()), (std::streamsize) bytes.size()); + check((bool) f, "cannot write fixture"); + } + ~Table() { std::error_code ignored; std::filesystem::remove(path, ignored); } +}; + +void open(PleReader& r, const Table& table, bool threaded) { + std::string err; + const bool ok = r.open(table.path.string(), 0, kRows, 4, 0, err, threaded, kRowBytes); + check(ok, "open failed: " + err); +} + +void read(PleReader& r) { + const uint32_t rows[] = {0, 45, 91, 127}; // different pages, including a row that crosses a page boundary + std::vector out(4 * kRowBytes, 0); + const auto ticket = r.issue(rows, 4, out.data()); + std::string err; + const bool ok = r.collect(ticket, err); + check(ok, "collect failed: " + err); + for (size_t i = 0; i < 4; ++i) + for (size_t b = 0; b < kRowBytes; ++b) + check(out[i * kRowBytes + b] == pattern((size_t) rows[i] * kRowBytes + b), "row bytes changed"); +} + +void refused(PleReader& r, const Table& table, bool threaded, bool allocation) { + std::string err; +#ifdef STRATA_TEST_ALLOC_FAILURE + if (allocation) fail_alloc = true; +#endif + const bool ok = r.open(table.path.string(), allocation ? 0 : kFileBytes, kRows, 4, 0, err, + threaded, kRowBytes); + check(!ok, "invalid initialization was accepted"); + check(err.find(allocation ? "cannot allocate read buffers" : "table extends past the end") != std::string::npos, + "unexpected initialization error: " + err); +#ifdef STRATA_TEST_ALLOC_FAILURE + check(!fail_alloc, "allocation failure was not exercised"); +#endif + check(!r.is_open(), "failed reader kept its file open"); + r.close(); + r.close(); +} +} // namespace + +int main(int argc, char** argv) { + try { + check(argc == 3, "usage: ple_reader_lifecycle_test reads|reopen|alloc caller|threaded"); + const std::string mode = argv[1]; + const bool threaded = std::string(argv[2]) == "threaded"; + check(mode == "reads" || mode == "reopen" || mode == "alloc", "unknown case"); + Table table; + PleReader r; + if (mode != "reads") refused(r, table, threaded, mode == "alloc"); + open(r, table, threaded); + read(r); + // open() must also clean up a previously initialized reader before handling the new failure. + if (mode != "reads") refused(r, table, threaded, mode == "alloc"); + else { + const uint32_t rows[] = {0, 45, 91, 127}; + std::vector out(4 * kRowBytes, 0); + (void) r.issue(rows, 4, out.data()); + r.close(); // outstanding writes must finish before out is released + } + r.close(); + open(r, table, threaded); + read(r); + r.close(); + r.close(); + std::printf("%s %s: passed\n", argv[1], argv[2]); + return 0; + } catch (const std::exception& e) { + std::fprintf(stderr, "%s\n", e.what()); + return 1; + } +}