[Engine] avoid hanging on PleReader initialization failures

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)

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;

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 <cerrno>
+#include <cstdio>
+#include <cstdlib>
+#include <filesystem>
+#include <fstream>
+#include <stdexcept>
+#include <string>
+#include <vector>
+#include <unistd.h>
+
+#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<uint8_t> 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<const char*>(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<uint8_t> 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<uint8_t> 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;
+    }
+}