Fix OMP exception safety in IndexNNDescent and IndexNSG search (facebookresearch#5106)

Summary:

Apply the same OpenMP exception-safety fix from IndexFlatCodes::search to IndexNNDescent::search and IndexNSG::search. Both methods have identical OMP structure: a parallel region constructs per-thread VisitedTable and DistanceComputer objects, then a worksharing loop calls the graph search function. Exceptions thrown inside either the constructor or the loop body (e.g. from corrupt deserialized graph state) call std::terminate because OpenMP does not allow exceptions to escape worksharing constructs.

Wrap the per-thread setup and the per-iteration loop body in try/catch blocks that capture exceptions via std::exception_ptr and re-throw on the main thread after the parallel region completes.

Potential exception sources in the OMP region:
- VisitedTable(ntotal): std::bad_alloc if ntotal is corrupt/huge. ntotal is validated as >= 0 during deserialization but has no upper bound.
- storage_distance_computer(storage): may throw from get_distance_computer() if the storage index has corrupt state.
- nndescent.search() / nsg.search(): throws FaissException if the graph has not been built (has_built check), or std::bad_alloc from internal vector allocations with corrupt search_L.

Reviewed By: mnorris11

Differential Revision: D101031002

Author: scsiguy

faiss/IndexNNDescent.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/IndexNNDescent.h>
 
+#include <atomic>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
@@ -119,22 +120,36 @@ void IndexNNDescent::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
         {
-            VisitedTable vt(ntotal);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
+            std::unique_ptr<DistanceComputer> dis;
+            std::unique_ptr<VisitedTable> vt;
+            try {
+                vt = std::make_unique<VisitedTable>(ntotal);
+                dis.reset(storage_distance_computer(storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for
             for (idx_t i = i0; i < i1; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
-
-                nndescent.search(*dis, k, idxi, simi, vt);
+                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);
+
+                    nndescent.search(*dis, k, idxi, simi, *vt);
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 

faiss/IndexNSG.cpp

@@ -9,6 +9,7 @@
 
 #include <faiss/IndexNSG.h>
 
+#include <atomic>
 #include <cinttypes>
 #include <memory>
 
@@ -74,24 +75,39 @@ void IndexNSG::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
         {
-            VisitedTable vt(ntotal, nsg.use_visited_hashset);
-
-            std::unique_ptr<DistanceComputer> dis(
-                    storage_distance_computer(storage));
+            std::unique_ptr<DistanceComputer> dis;
+            std::unique_ptr<VisitedTable> vt;
+            try {
+                vt = std::make_unique<VisitedTable>(
+                        ntotal, nsg.use_visited_hashset);
+                dis.reset(storage_distance_computer(storage));
+            } catch (...) {
+                omp_capture_exception(ex, [&] { interrupt = true; });
+            }
 
 #pragma omp for
             for (idx_t i = i0; i < i1; i++) {
-                idx_t* idxi = labels + i * k;
-                float* simi = distances + i * k;
-                dis->set_query(x + i * d);
+                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);
 
-                nsg.search(*dis, k, idxi, simi, vt);
+                    nsg.search(*dis, k, idxi, simi, *vt);
 
-                vt.advance();
+                    vt->advance();
+                } catch (...) {
+                    omp_capture_exception(ex, [&] { interrupt = true; });
+                }
             }
         }
+        omp_rethrow_if_exception(ex);
         InterruptCallback::check();
     }
 

tests/test_omp_exception_safety.cpp

@@ -16,6 +16,8 @@
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexFlatCodes.h>
 #include <faiss/IndexIVFFlat.h>
+#include <faiss/IndexNNDescent.h>
+#include <faiss/IndexNSG.h>
 #include <faiss/impl/AuxIndexStructures.h>
 #include <faiss/impl/FaissException.h>
 #include <faiss/invlists/InvertedLists.h>
@@ -333,3 +335,49 @@ TEST(OMPExceptionSafety, flatcodes_search) {
             index.search(1, xq.data(), 1, distances.data(), labels.data()),
             std::runtime_error);
 }
+
+// ---------------------------------------------------------------------------
+// IndexNNDescent::search: exception in OMP worker propagates to caller.
+// Constructing with has_built=false triggers a FaissException inside the
+// worksharing loop body.
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, nndescent_search) {
+    int d = 4;
+    auto storage = std::make_unique<IndexFlatL2>(d);
+    std::vector<float> xb(d, 1.0f);
+    storage->add(1, xb.data());
+
+    IndexNNDescent index(storage.get(), 4);
+    index.ntotal = 1;
+    // has_built defaults to false, so nndescent.search() will throw.
+
+    std::vector<float> xq(d, 0.0f);
+    std::vector<float> distances(1);
+    std::vector<idx_t> labels(1);
+
+    EXPECT_THROW(
+            index.search(1, xq.data(), 1, distances.data(), labels.data()),
+            FaissException);
+}
+
+// ---------------------------------------------------------------------------
+// IndexNSG::search: exception in OMP worker propagates to caller
+// ---------------------------------------------------------------------------
+TEST(OMPExceptionSafety, nsg_search) {
+    int d = 4;
+    auto storage = std::make_unique<IndexFlatL2>(d);
+    std::vector<float> xb(d, 1.0f);
+    storage->add(1, xb.data());
+
+    IndexNSG index(storage.get(), 4);
+    index.ntotal = 1;
+    // nsg graph is not built, so nsg.search() will throw.
+
+    std::vector<float> xq(d, 0.0f);
+    std::vector<float> distances(1);
+    std::vector<idx_t> labels(1);
+
+    EXPECT_THROW(
+            index.search(1, xq.data(), 1, distances.data(), labels.data()),
+            FaissException);
+}