Fix IDSelector memory leak with SWIG, add unit tests (#3704)

faiss/python/class_wrappers.py

@@ -1170,28 +1170,6 @@ def add_to_referenced_objects(self, ref):
     else:
         self.referenced_objects.append(ref)
 
-class RememberSwigOwnership:
-    """
-    SWIG's seattr transfers ownership of SWIG wrapped objects to the class
-    (btw this seems to contradict https://www.swig.org/Doc1.3/Python.html#Python_nn22
-    31.4.2)
-    This interferes with how we manage ownership: with the referenced_objects
-    table. Therefore, we reset the thisown field in this context manager.
-    """
-
-    def __init__(self, obj):
-        self.obj = obj
-
-    def __enter__(self):
-        if hasattr(self.obj, "thisown"):
-            self.old_thisown = self.obj.thisown
-        else:
-            self.old_thisown = None
-
-    def __exit__(self, *ignored):
-        if self.old_thisown is not None:
-            self.obj.thisown = self.old_thisown
-
 
 def handle_SearchParameters(the_class):
     """ this wrapper is to enable initializations of the form
@@ -1206,8 +1184,7 @@ def replacement_init(self, **args):
         self.original_init()
         for k, v in args.items():
             assert hasattr(self, k)
-            with RememberSwigOwnership(v):
-                setattr(self, k, v)
+            setattr(self, k, v)
             if type(v) not in (int, float, bool, str):
                 add_to_referenced_objects(self, v)
 

faiss/python/swigfaiss.swig

@@ -46,6 +46,20 @@ typedef uint64_t size_t;
  * the language includes for their respective matrix libraries.
  *******************************************************************/
 
+/********************************************************
+ * Don't transfer ownership when assigning to a struct field
+ ********************************************************/
+
+// https://github.com/swig/swig/issues/2709
+// override default behavior
+%typemap(in, noblock=1) SWIGTYPE * (void  *argp = 0, int res = 0) {
+  res = SWIG_ConvertPtr($input, &argp,$descriptor, /* $disown | */ %convertptr_flags);
+  if (!SWIG_IsOK(res)) {
+    %argument_fail(res, "$type", $symname, $argnum);
+  }
+  $1 = %reinterpret_cast(argp, $ltype);
+}
+
 %{
 
 

tests/test_graph_based.py

@@ -154,8 +154,12 @@ def test_io_no_storage(self):
             faiss.serialize_index(index2))
 
         # add storage afterwards
-        index.storage = faiss.clone_index(index.storage)
-        index.own_fields = True
+        # Trying to call index.storage.get_distance_computer() after cloning
+        # will crash, unless we assign to an intermediate variable.
+        # We also get a heap-use-after-free unless we set own_fields to False.
+        temp = faiss.clone_index(index.storage)
+        index.storage = temp
+        temp.this.own(False)  # Alternatively set index.own_fields = False
 
         Dnew, Inew = index.search(self.xq, 5)
         np.testing.assert_array_equal(Dnew, Dref)

tests/test_search_params.py

@@ -397,6 +397,24 @@ def test_ownership(self):
         self.assertTrue(sel1.this.own())
         self.assertTrue(sel2.this.own())
 
+    def test_ownership_2(self):
+        res = []
+        for _ in range(20):
+            params = faiss.SearchParameters()
+            subset = np.arange(0, 5000000)
+            params.sel = faiss.IDSelectorBatch(subset)
+            mem_usage = faiss.get_mem_usage_kb() / 1024**2
+            res.append(round(mem_usage, 2))
+        print(res)
+        self.assertTrue(res[-1] < 0.9)
+
+    def test_ownership_3(self):
+        subset = np.arange(0, 5000000)
+        sel = faiss.IDSelectorBatch(subset)
+        assert sel.this.own()    # True: correct
+        _ = faiss.SearchParameters(sel=sel)
+        assert sel.this.own()   # False: why???
+
 
 class TestSelectorCallback(unittest.TestCase):