From 798427c019d1596ff720fc5785d5fcacf31a9b7c Mon Sep 17 00:00:00 2001
From: Kumar Saurabh Arora <kuarora@meta.com>
Date: Fri, 22 Mar 2024 14:58:14 -0700
Subject: [PATCH] Handling FaissException in few destructors of ResultHandler.h
 (#3311)

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

**Context**
[Issue 2948](https://github.com/facebookresearch/faiss/issues/2948) highlights potential issue of calling allocation on result handler which may throw exception but it is not handled.

**In this diff**,
I observed two calls where we may potentially call allocation in ResultHandler.h and handled FaissException.
1/ partial result when finalized in ~SingleResultHandler
2/ partial result when merged in ~RangeSearchBlockResultHandler

Reviewed By: junjieqi

Differential Revision: D55258213

fbshipit-source-id: 259be472e73619b2fcb0ea480d6d3486affeafdf
---
 faiss/impl/ResultHandler.h | 23 ++++++++++++++++++++---
 1 file changed, 20 insertions(+), 3 deletions(-)

diff --git a/faiss/impl/ResultHandler.h b/faiss/impl/ResultHandler.h
index 270de8dcd6..713fe8e49f 100644
--- a/faiss/impl/ResultHandler.h
+++ b/faiss/impl/ResultHandler.h
@@ -12,8 +12,10 @@
 #pragma once
 
 #include <faiss/impl/AuxIndexStructures.h>
+#include <faiss/impl/FaissException.h>
 #include <faiss/utils/Heap.h>
 #include <faiss/utils/partitioning.h>
+#include <iostream>
 
 namespace faiss {
 
@@ -504,7 +506,15 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C> {
         void end() {}
 
         ~SingleResultHandler() {
-            pres.finalize();
+            try {
+                // finalize the partial result
+                pres.finalize();
+            } catch (const faiss::FaissException& e) {
+                // Do nothing if allocation fails in finalizing partial results.
+#ifndef NDEBUG
+                std::cerr << e.what() << std::endl;
+#endif
+            }
         }
     };
 
@@ -559,8 +569,15 @@ struct RangeSearchBlockResultHandler : BlockResultHandler<C> {
     }
 
     ~RangeSearchBlockResultHandler() {
-        if (partial_results.size() > 0) {
-            RangeSearchPartialResult::merge(partial_results);
+        try {
+            if (partial_results.size() > 0) {
+                RangeSearchPartialResult::merge(partial_results);
+            }
+        } catch (const faiss::FaissException& e) {
+            // Do nothing if allocation fails in merge.
+#ifndef NDEBUG
+            std::cerr << e.what() << std::endl;
+#endif
         }
     }
 };