Fix OMP exception safety in IndexHNSW search (#5133)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5133

Exceptions thrown inside OpenMP parallel regions in IndexHNSW search methods call std::terminate because the OpenMP specification does not allow exceptions to escape parallel regions. This is triggered by corrupt serialized index data that causes allocation failures (e.g. VisitedTable construction with absurdly large ntotal) or other errors in DistanceComputer, but can occur with any exception thrown during the HNSW search loop.

Wrap the OpenMP parallel bodies in hnsw_search(), IndexHNSW::search_level_0(), and IndexHNSW2Level::search() with per-thread try/catch blocks that capture exceptions via std::exception_ptr and re-throw them on the main thread after the parallel region completes. This follows the same pattern established in IndexFlatCodes and IndexNNDescent/IndexNSG.

VisitedTable, SingleResultHandler, and DistanceComputer are moved from stack to heap allocation (std::unique_ptr) because the constructors are now inside try/catch blocks. If a stack-allocated object's constructor throws, the object never comes into existence, but OpenMP requires all threads to participate in the worksharing loop that follows — any reference to the uninitialized object would be undefined behavior. With std::unique_ptr, the variable starts as nullptr before the try block, remains nullptr if construction fails, and the interrupt flag ensures all loop iterations skip via continue without dereferencing it.

Reviewed By: mnorris11

Differential Revision: D101638521

fbshipit-source-id: 3920be5e41b1bff0490a14e2571b3dbeace64397

faiss/IndexHNSW.cpp

@@ -8,6 +8,7 @@
 #include <faiss/IndexHNSW.h>
 
 #include <omp.h>
+#include <atomic>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
@@ -26,6 +27,7 @@
 #include <faiss/IndexIVFPQ.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissAssert.h>
+#include <faiss/impl/FaissException.h>
 #include <faiss/impl/ResultHandler.h>
 #include <faiss/impl/VisitedTable.h>
 #include <faiss/impl/hnsw/MinimaxHeap.h>
@@ -264,29 +266,48 @@ void hnsw_search(
 
     for (idx_t i0 = 0; i0 < n; i0 += check_period) {
         idx_t i1 = std::min(i0 + check_period, n);
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 
 #pragma omp parallel if (i1 - i0 > 1)
         {
-            VisitedTable vt(index->ntotal, hnsw.use_visited_hashset);
-            typename BlockResultHandler::SingleResultHandler res(bres);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(index->storage));
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<typename BlockResultHandler::SingleResultHandler>
+                    res;
+            std::unique_ptr<DistanceComputer> dis;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        index->ntotal, hnsw.use_visited_hashset);
+                res = std::make_unique<
+                        typename BlockResultHandler::SingleResultHandler>(bres);
+                dis.reset(storage_distance_computer(index->storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
             for (idx_t i = i0; i < i1; i++) {
-                res.begin(i);
-                dis->set_query(x + i * index->d);
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
+                }
+                try {
+                    res->begin(i);
+                    dis->set_query(x + i * index->d);
 
-                HNSWStats stats = hnsw.search(*dis, index, res, vt, params);
-                n1 += stats.n1;
-                n2 += stats.n2;
-                ndis += stats.ndis;
-                nhops += stats.nhops;
-                res.end();
-                vt.advance();
+                    HNSWStats stats =
+                            hnsw.search(*dis, index, *res, *vt, params);
+                    n1 += stats.n1;
+                    n2 += stats.n2;
+                    ndis += stats.ndis;
+                    nhops += stats.nhops;
+                    res->end();
+                    vt->advance();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 
@@ -441,37 +462,54 @@ void IndexHNSW::search_level_0(
     using RH = HeapBlockResultHandler<HNSW::C>;
     RH bres(n, distances, labels, k);
 
+    std::exception_ptr ex;
+    std::atomic<bool> interrupt{false};
 #pragma omp parallel
     {
-        std::unique_ptr<DistanceComputer> qdis(
-                storage_distance_computer(storage));
+        std::unique_ptr<DistanceComputer> qdis;
         HNSWStats search_stats;
-        VisitedTable vt(hnsw_ntotal, hnsw.use_visited_hashset);
-        RH::SingleResultHandler res(bres);
+        std::unique_ptr<VisitedTable> vt;
+        std::unique_ptr<RH::SingleResultHandler> res;
+        try {
+            qdis.reset(storage_distance_computer(storage));
+            vt = std::make_unique<VisitedTable>(
+                    hnsw_ntotal, hnsw.use_visited_hashset);
+            res = std::make_unique<RH::SingleResultHandler>(bres);
+        } catch (...) {
+            omp_capture_exception(ex, [&] { interrupt = true; });
+        }
 
 #pragma omp for
         for (idx_t i = 0; i < n; i++) {
-            res.begin(i);
-            qdis->set_query(x + i * d);
+            if (interrupt.load(std::memory_order_relaxed)) {
+                continue;
+            }
+            try {
+                res->begin(i);
+                qdis->set_query(x + i * d);
 
-            hnsw.search_level_0(
-                    *qdis.get(),
-                    res,
-                    nprobe,
-                    nearest + i * nprobe,
-                    nearest_d + i * nprobe,
-                    search_type,
-                    search_stats,
-                    vt,
-                    params);
-            res.end();
-            vt.advance();
+                hnsw.search_level_0(
+                        *qdis.get(),
+                        *res,
+                        nprobe,
+                        nearest + i * nprobe,
+                        nearest_d + i * nprobe,
+                        search_type,
+                        search_stats,
+                        *vt,
+                        params);
+                res->end();
+                vt->advance();
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
         }
 #pragma omp critical
         {
             hnsw_stats.combine(search_stats);
         }
     }
+    omp_rethrow_if_exception(ex);
     if (is_similarity_metric(this->metric_type)) {
 // we need to revert the negated distances
 #pragma omp parallel for
@@ -883,73 +921,88 @@ void IndexHNSW2Level::search(
                 labels,
                 false);
 
+        std::exception_ptr ex;
+        std::atomic<bool> interrupt{false};
 #pragma omp parallel
         {
             // visited table (not hash set) for tri-state flags.
-            VisitedTable vt(ntotal, /*use_hashset=*/false);
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
-
+            std::unique_ptr<VisitedTable> vt;
+            std::unique_ptr<DistanceComputer> dis;
             constexpr int candidates_size = 1;
-            MinimaxHeap candidates(candidates_size);
+            std::unique_ptr<MinimaxHeap> candidates;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        ntotal, /*use_hashset=*/false);
+                dis.reset(storage_distance_computer(storage));
+                candidates = std::make_unique<MinimaxHeap>(candidates_size);
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for reduction(+ : n1, n2, ndis, nhops)
             for (idx_t i = 0; i < n; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
-
-                // mark all inverted list elements as visited
-
-                for (size_t j = 0; j < nprobe; j++) {
-                    idx_t key = coarse_assign[j + i * nprobe];
-                    if (key < 0) {
-                        break;
-                    }
-                    size_t list_length = index_ivfpq->get_list_size(key);
-                    const idx_t* ids = index_ivfpq->invlists->get_ids(key);
-
-                    for (size_t jj = 0; jj < list_length; jj++) {
-                        vt.set(ids[jj]);
-                    }
+                if (interrupt.load(std::memory_order_relaxed)) {
+                    continue;
                 }
+                try {
+                    idx_t* idxi = labels + i * k;
+                    float* simi = distances + i * k;
+                    dis->set_query(x + i * d);
 
-                candidates.clear();
+                    // mark all inverted list elements as visited
+                    for (size_t j = 0; j < nprobe; j++) {
+                        idx_t key = coarse_assign[j + i * nprobe];
+                        if (key < 0) {
+                            break;
+                        }
+                        size_t list_length = index_ivfpq->get_list_size(key);
+                        const idx_t* ids = index_ivfpq->invlists->get_ids(key);
 
-                for (int j = 0; j < k; j++) {
-                    if (idxi[j] < 0) {
-                        break;
+                        for (size_t jj = 0; jj < list_length; jj++) {
+                            vt->set(ids[jj]);
+                        }
                     }
-                    candidates.push(
-                            static_cast<storage_idx_t>(idxi[j]), simi[j]);
+
+                    candidates->clear();
+
+                    for (int j = 0; j < k; j++) {
+                        if (idxi[j] < 0) {
+                            break;
+                        }
+                        candidates->push(
+                                static_cast<storage_idx_t>(idxi[j]), simi[j]);
+                    }
+
+                    // reorder from sorted to heap
+                    maxheap_heapify(k, simi, idxi, simi, idxi, k);
+
+                    HNSWStats search_stats;
+                    search_from_candidates_2(
+                            hnsw,
+                            *dis,
+                            k,
+                            idxi,
+                            simi,
+                            *candidates,
+                            *vt,
+                            search_stats,
+                            0,
+                            k);
+                    n1 += search_stats.n1;
+                    n2 += search_stats.n2;
+                    ndis += search_stats.ndis;
+                    nhops += search_stats.nhops;
+
+                    vt->advance();
+                    vt->advance();
+
+                    maxheap_reorder(k, simi, idxi);
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
                 }
-
-                // reorder from sorted to heap
-                maxheap_heapify(k, simi, idxi, simi, idxi, k);
-
-                HNSWStats search_stats;
-                search_from_candidates_2(
-                        hnsw,
-                        *dis,
-                        static_cast<int>(k),
-                        idxi,
-                        simi,
-                        candidates,
-                        vt,
-                        search_stats,
-                        0,
-                        static_cast<int>(k));
-                n1 += search_stats.n1;
-                n2 += search_stats.n2;
-                ndis += search_stats.ndis;
-                nhops += search_stats.nhops;
-
-                vt.advance();
-                vt.advance();
-
-                maxheap_reorder(k, simi, idxi);
             }
         }
+        omp_rethrow_if_exception(ex);
 
         hnsw_stats.combine({n1, n2, ndis, nhops});
     }

tests/test_omp_exception_safety.cpp

@@ -13,9 +13,12 @@
 
 #include <gtest/gtest.h>
 
+#include <faiss/Index2Layer.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexFlatCodes.h>
+#include <faiss/IndexHNSW.h>
 #include <faiss/IndexIVFFlat.h>
+#include <faiss/IndexIVFPQ.h>
 #include <faiss/IndexNNDescent.h>
 #include <faiss/IndexNSG.h>
 #include <faiss/impl/AuxIndexStructures.h>
@@ -381,3 +384,336 @@ TEST(OMPExceptionSafety, nsg_search) {
             index.search(1, xq.data(), 1, distances.data(), labels.data()),
             FaissException);
 }
+
+// ---------------------------------------------------------------------------
+// IndexHNSW OMP exception safety tests.
+//
+// There are 3 OMP regions wrapped with exception capture in IndexHNSW.cpp:
+//   1. hnsw_search()             — primary search path
+//   2. IndexHNSW::search_level_0 — level-0 search
+//   3. IndexHNSW2Level::search   — mixed IVFPQ+HNSW path
+//
+// Each region has two throw locations:
+//   (a) Init phase  — DistanceComputer/VisitedTable construction
+//   (b) Loop body   — set_query, graph traversal, etc.
+// ---------------------------------------------------------------------------
+
+// DistanceComputer that throws from set_query(). Exercises loop-body
+// try/catch after successful init-phase construction.
+struct ThrowingDC : DistanceComputer {
+    void set_query(const float*) override {
+        throw std::runtime_error("corrupt storage");
+    }
+    float operator()(idx_t) override {
+        return 0;
+    }
+    float symmetric_dis(idx_t, idx_t) override {
+        return 0;
+    }
+};
+
+// Storage whose get_distance_computer() returns a ThrowingDC.
+// Exercises loop-body throws (DC construction succeeds, set_query throws).
+struct LoopThrowingStorage : IndexFlatL2 {
+    using IndexFlatL2::IndexFlatL2;
+    DistanceComputer* get_distance_computer() const override {
+        return new ThrowingDC();
+    }
+};
+
+// Storage whose get_distance_computer() itself throws.
+// Exercises init-phase throws (DC construction fails).
+struct InitThrowingStorage : IndexFlatL2 {
+    using IndexFlatL2::IndexFlatL2;
+    DistanceComputer* get_distance_computer() const override {
+        throw std::runtime_error("corrupt storage init");
+    }
+};
+
+// Helper: build a valid IndexHNSWFlat with data, then swap storage.
+struct HNSWFixture {
+    static constexpr int d = 4;
+    static constexpr int nb = 10;
+    static constexpr int M = 16;
+
+    IndexHNSWFlat index{d, M};
+    std::vector<float> xb;
+    std::vector<float> xq;
+
+    HNSWFixture() : xb(d * nb), xq(d) {
+        std::mt19937 rng(42);
+        std::uniform_real_distribution<float> dist;
+        for (auto& v : xb) {
+            v = dist(rng);
+        }
+        index.add(nb, xb.data());
+        for (auto& v : xq) {
+            v = dist(rng);
+        }
+    }
+
+    void swap_storage(Index* new_storage) {
+        index.own_fields = false;
+        index.storage = new_storage;
+    }
+
+    void restore_storage(Index* old_storage) {
+        index.storage = old_storage;
+        index.own_fields = true;
+    }
+};
+
+// ---------------------------------------------------------------------------
+// hnsw_search() init-phase throw: storage_distance_computer() throws
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, hnsw_search_init) {
+    HNSWFixture f;
+    auto throwing = std::make_unique<InitThrowingStorage>(f.d);
+    throwing->add(f.nb, f.xb.data());
+    auto* old = f.index.storage;
+    f.swap_storage(throwing.get());
+
+    std::vector<float> distances(1);
+    std::vector<idx_t> labels(1);
+    EXPECT_THROW(
+            f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()),
+            std::runtime_error);
+
+    f.restore_storage(old);
+}
+
+// ---------------------------------------------------------------------------
+// hnsw_search() loop-body throw: set_query() throws
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, hnsw_search_loop) {
+    HNSWFixture f;
+    auto throwing = std::make_unique<LoopThrowingStorage>(f.d);
+    throwing->add(f.nb, f.xb.data());
+    auto* old = f.index.storage;
+    f.swap_storage(throwing.get());
+
+    std::vector<float> distances(1);
+    std::vector<idx_t> labels(1);
+    EXPECT_THROW(
+            f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()),
+            std::runtime_error);
+
+    f.restore_storage(old);
+}
+
+// ---------------------------------------------------------------------------
+// search_level_0() init-phase throw: storage_distance_computer() throws
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, hnsw_search_level_0_init) {
+    HNSWFixture f;
+    auto throwing = std::make_unique<InitThrowingStorage>(f.d);
+    throwing->add(f.nb, f.xb.data());
+    auto* old = f.index.storage;
+    f.swap_storage(throwing.get());
+
+    idx_t k = 1;
+    int nprobe = 1;
+    std::vector<HNSW::storage_idx_t> nearest(nprobe, 0);
+    std::vector<float> nearest_d(nprobe, 0.0f);
+    std::vector<float> distances(k);
+    std::vector<idx_t> labels(k);
+
+    EXPECT_THROW(
+            f.index.search_level_0(
+                    1,
+                    f.xq.data(),
+                    k,
+                    nearest.data(),
+                    nearest_d.data(),
+                    distances.data(),
+                    labels.data(),
+                    nprobe),
+            std::runtime_error);
+
+    f.restore_storage(old);
+}
+
+// ---------------------------------------------------------------------------
+// search_level_0() loop-body throw: set_query() throws
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, hnsw_search_level_0_loop) {
+    HNSWFixture f;
+    auto throwing = std::make_unique<LoopThrowingStorage>(f.d);
+    throwing->add(f.nb, f.xb.data());
+    auto* old = f.index.storage;
+    f.swap_storage(throwing.get());
+
+    idx_t k = 1;
+    int nprobe = 1;
+    std::vector<HNSW::storage_idx_t> nearest(nprobe, 0);
+    std::vector<float> nearest_d(nprobe, 0.0f);
+    std::vector<float> distances(k);
+    std::vector<idx_t> labels(k);
+
+    EXPECT_THROW(
+            f.index.search_level_0(
+                    1,
+                    f.xq.data(),
+                    k,
+                    nearest.data(),
+                    nearest_d.data(),
+                    distances.data(),
+                    labels.data(),
+                    nprobe),
+            std::runtime_error);
+
+    f.restore_storage(old);
+}
+
+// IndexIVFPQ subclass with a counted throw mechanism. The first
+// `throw_after` calls to get_distance_computer() delegate to the
+// real implementation; subsequent calls throw (init) or return
+// ThrowingDC (loop). This lets search_preassigned succeed while
+// the HNSW2Level OMP region's storage_distance_computer() fails.
+struct CountedThrowIVFPQ : IndexIVFPQ {
+    mutable std::atomic<int> calls_remaining;
+    bool throw_from_init;
+
+    CountedThrowIVFPQ(
+            Index* quantizer,
+            size_t d,
+            size_t nlist,
+            size_t M,
+            size_t nbits,
+            MetricType metric,
+            int throw_after,
+            bool init_throw)
+            : IndexIVFPQ(quantizer, d, nlist, M, nbits, metric),
+              calls_remaining(throw_after),
+              throw_from_init(init_throw) {}
+
+    DistanceComputer* get_distance_computer() const override {
+        if (calls_remaining.fetch_sub(1) <= 0) {
+            if (throw_from_init) {
+                throw std::runtime_error("IVFPQ init throw");
+            }
+            return new ThrowingDC();
+        }
+        return IndexIVFPQ::get_distance_computer();
+    }
+};
+
+struct HNSW2LevelFixture {
+    static constexpr int d = 8;
+    static constexpr int nlist = 4;
+    static constexpr int m_pq = 4;
+    static constexpr int nb = 512;
+    static constexpr int M = 16;
+
+    IndexFlatL2 quantizer{d};
+    IndexHNSW2Level index{&quantizer, nlist, m_pq, M};
+    std::vector<float> xb;
+    std::vector<float> xq;
+
+    HNSW2LevelFixture() : xb(nb * d), xq(d) {
+        index.own_fields = false;
+        std::mt19937 rng(42);
+        std::uniform_real_distribution<float> dist;
+        for (auto& v : xb) {
+            v = dist(rng);
+        }
+        index.train(nb, xb.data());
+        index.add(nb, xb.data());
+        // Convert Index2Layer storage to IndexIVFPQ for the mixed
+        // search path. After this, storage is an IndexIVFPQ owned
+        // by the HNSW2Level index.
+        index.flip_to_ivf();
+        index.own_fields = true;
+        for (auto& v : xq) {
+            v = dist(rng);
+        }
+    }
+};
+
+// ---------------------------------------------------------------------------
+// IndexHNSW2Level::search() mixed path init-phase throw:
+// storage_distance_computer() throws during OMP init, after
+// search_preassigned has already succeeded.
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, hnsw_2level_mixed_init) {
+    HNSW2LevelFixture f;
+
+    // The replacement IVFPQ shares the real index's trained state
+    // (codebooks, inverted lists). We copy the real IVFPQ's invlists
+    // and PQ so search_preassigned succeeds, then the counted throw
+    // fires during the HNSW2Level OMP init phase.
+    auto* real_ivfpq = dynamic_cast<IndexIVFPQ*>(f.index.storage);
+    ASSERT_NE(real_ivfpq, nullptr);
+
+    CountedThrowIVFPQ throwing(
+            &f.quantizer,
+            f.d,
+            f.nlist,
+            f.m_pq,
+            8,
+            METRIC_L2,
+            /*throw_after=*/100,
+            /*init_throw=*/true);
+    throwing.own_fields = false;
+    throwing.is_trained = true;
+    throwing.ntotal = real_ivfpq->ntotal;
+    throwing.pq = real_ivfpq->pq;
+    throwing.code_size = real_ivfpq->code_size;
+    throwing.nprobe = real_ivfpq->nprobe;
+    throwing.replace_invlists(real_ivfpq->invlists, false);
+
+    auto* old = f.index.storage;
+    f.index.storage = &throwing;
+    f.index.own_fields = false;
+
+    std::vector<float> distances(1);
+    std::vector<idx_t> labels(1);
+    EXPECT_THROW(
+            f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()),
+            std::exception);
+
+    f.index.storage = old;
+    f.index.own_fields = true;
+}
+
+// ---------------------------------------------------------------------------
+// IndexHNSW2Level::search() mixed path loop-body throw:
+// set_query() throws inside OMP for loop, after search_preassigned
+// and OMP init have succeeded.
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, hnsw_2level_mixed_loop) {
+    HNSW2LevelFixture f;
+
+    auto* real_ivfpq = dynamic_cast<IndexIVFPQ*>(f.index.storage);
+    ASSERT_NE(real_ivfpq, nullptr);
+
+    CountedThrowIVFPQ throwing(
+            &f.quantizer,
+            f.d,
+            f.nlist,
+            f.m_pq,
+            8,
+            METRIC_L2,
+            /*throw_after=*/100,
+            /*init_throw=*/false);
+    throwing.own_fields = false;
+    throwing.is_trained = true;
+    throwing.ntotal = real_ivfpq->ntotal;
+    throwing.pq = real_ivfpq->pq;
+    throwing.code_size = real_ivfpq->code_size;
+    throwing.nprobe = real_ivfpq->nprobe;
+    throwing.replace_invlists(real_ivfpq->invlists, false);
+
+    auto* old = f.index.storage;
+    f.index.storage = &throwing;
+    f.index.own_fields = false;
+
+    std::vector<float> distances(1);
+    std::vector<idx_t> labels(1);
+    EXPECT_THROW(
+            f.index.search(1, f.xq.data(), 1, distances.data(), labels.data()),
+            std::exception);
+
+    f.index.storage = old;
+    f.index.own_fields = true;
+}