Fix IDSelector leak via SearchParameters.sel setter (#5208)

Summary:
Pull Request resolved: #5208

A user-reported leak (#2996 follow-up) reproduces on faiss 1.14.1 when an `IDSelectorBatch` is assigned to a `SearchParameters` field after construction. Each iteration of the reporter's loop leaks ~160 MB; over a long-running serving process this grows unbounded.

# Root cause

The Python wrapper for `IDSelectorBatch` carries a SWIG `thisown` flag that controls who calls C++ `delete`. The SWIG-generated property setter for `SearchParameters.sel` flips `thisown` to 0 on the assumption that the enclosing C++ object will take responsibility, but `SearchParameters`'s C++ destructor does not free `sel` (the field is borrowed by design). Result: the C++ `IDSelectorBatch` (a 40 MB sorted `std::vector<idx_t>` for the reporter's case) plus the retained numpy buffer are orphaned forever.

`handle_SearchParameters` in `class_wrappers.py` already protected the kwargs construction path (`SearchParameters(sel=x)`) by wrapping the assignment in `RememberSwigOwnership` and `add_to_referenced_objects`. The bare-attribute path was unprotected.

# Fix

Override `__setattr__` on the base `SearchParameters` class so the same ownership protection runs for both `SearchParameters(sel=x)` and `params.sel = x`. The override filters via `hasattr(v, "thisown")` so only SWIG-wrapped C++ objects go through the dance; SWIG internals like `self.this` (a `SwigPyObject` without `thisown`), bookkeeping lists, and plain Python values bypass it.

`replacement_init` is slimmed to a pure delegator: the protection moves into `replacement_setattr` so a single source of truth handles both code paths. Keeping both layers would double-wrap on kwargs construction.

The filter changes from a denylist (`type(v) not in (int, float, bool, str)`) to a duck-type check (`hasattr(v, "thisown")`). The new filter is stricter: numpy arrays, lists, and `None` no longer get appended. Safe in practice because `SearchParameters` SWIG fields are typed as C++ pointers and only legitimately accept SWIG wrappers; numpy buffers passed indirectly (e.g., to `IDSelectorBatch`) are kept alive by the inner wrapper's own `add_to_referenced_objects`.

The duck-type filter is field-name-agnostic, so other `SearchParameters` subclass fields with the same C++ shape (`Foo* field = nullptr` with no destructor freeing `field`) benefit from the same protection automatically. Examples include `SearchParametersPreTransform.index_params` and `SearchParametersIVF.quantizer_params`.

# Per-field ref tracking

A naive implementation that uses the existing `add_to_referenced_objects` helper (which appends to a list) leaks under a different access pattern: `params = SearchParameters(); for ...: params.sel = ...` accumulates one entry per reassignment because nothing drops the prior ref. This is a common pattern in long-lived servers that initialize a `SearchParameters` once and rebind `params.sel` per request.

Fixed by using a per-field dict `self._sp_field_refs` keyed by attribute name. Reassigning the same field replaces the entry in the dict, so the prior wrapper loses its Python ref and the C++ object is freed. The else branch of `replacement_setattr` also drops any prior ref when the field is set to None or to a non-SWIG value, so explicit clear (`params.sel = None`) releases the C++ object immediately.

# Subclass guard

`__setattr__` is installed only once per class hierarchy. SWIG generates per-class `__init__` slots but does NOT generate per-class `__setattr__`, so subclasses (`SearchParametersIVF`, `SearchParametersHNSW`, ...) inherit the override via MRO. Without the `_protected_setattr` sentinel, calling `handle_SearchParameters(SearchParametersIVF)` would capture the inherited `replacement_setattr` as `parent_setattr`, then install a second `replacement_setattr` on top - chained calls would fire field tracking twice. The guard uses `getattr` (which walks the MRO) rather than `__dict__` lookup (own-attribute only) so subclasses correctly see the base's marker.

`replacement_init` does not need this guard because each SWIG class generates its own `__init__` in its own `__dict__`; wrapping always captures SWIG's fresh init, never an inherited replacement.

# Prior attempts and why this approach is targeted

This bug has been attempted twice before via the SWIG layer:

- #3139 (Apr 2024, abandoned)
- #3810 (Sept 2024, changes requested)

Both added a `%typemap(in) SWIGTYPE *` in `swigfaiss.swig` that strips the `$disown` flag globally, so SWIG never transfers ownership for any pointer field assignment anywhere in the bindings. The minimal change is one typemap.

The trade-off proved untenable. PR 3810 surfaced the smoking gun while testing: `test_graph_based.test_io_no_storage` started crashing with heap-use-after-free at `index.storage = faiss.clone_index(index.storage)` with `own_fields = True`. The author worked around it by introducing a temp variable and flipping `own_fields = False` - but that flip is itself a behavior change to a documented FAISS contract that downstream users build on.

The right framing is not "SWIG's default ownership transfer is always wrong" (the typemap assumption) but "SWIG's default ownership transfer is wrong specifically for `SearchParameters` because that class lacks a destructor for its borrowed pointer fields." This diff scopes the change to that class hierarchy; every other SWIG-managed field keeps its existing semantics. No `swigfaiss.swig` change, no surprising downstream breakage.

The memory regression test pattern (`np.arange(5_000_000)` × N iterations, measure `faiss.get_mem_usage_kb`) was independently validated by PR 3810's `test_ownership_2` and is reused here. ASAN cannot catch this leak because the C++ object remains reachable from a Python wrapper that simply has no Python-side references - it is a Python-level reference leak, not a C++-level dangling pointer.

Reviewed By: mnorris11

Differential Revision: D104750178

fbshipit-source-id: 41f865028d0dc863fe560d1fd1bbe979dc8e8e8f

Author: alibeklfc

faiss/python/class_wrappers.py

@@ -1316,25 +1316,43 @@ def __exit__(self, *ignored):
 
 
 def handle_SearchParameters(the_class):
-    """ this wrapper is to enable initializations of the form
-    SearchParametersXX(a=3, b=SearchParamsYY)
-    This also requires the enclosing class to keep a reference on the
-    sub-object, since the C++ code assumes the object ownership is
-    handled externally.
+    """Protect SearchParameters from leaking SWIG-owned sub-objects assigned
+    via either kwargs construction (SearchParametersXX(sel=x)) or bare
+    attribute assignment (params.sel = x).
     """
     the_class.original_init = the_class.__init__
 
     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)
-            if type(v) not in (int, float, bool, str):
-                add_to_referenced_objects(self, v)
+            setattr(self, k, v)
 
     the_class.__init__ = replacement_init
 
+    # Install __setattr__ once per hierarchy; subclasses inherit via MRO.
+    if getattr(the_class, "_protected_setattr", False):
+        return
+    parent_setattr = the_class.__setattr__
+
+    def replacement_setattr(self, k, v):
+        # Per-field ref dict. Reassigning the same field drops the prior
+        # ref instead of accumulating, so a long-lived SearchParameters
+        # with repeated `params.sel = ...` does not leak.
+        if v is not None and hasattr(v, "thisown"):
+            if not hasattr(self, "_sp_field_refs"):
+                parent_setattr(self, "_sp_field_refs", {})
+            with RememberSwigOwnership(v):
+                parent_setattr(self, k, v)
+            self._sp_field_refs[k] = v
+        else:
+            parent_setattr(self, k, v)
+            if hasattr(self, "_sp_field_refs"):
+                self._sp_field_refs.pop(k, None)
+
+    the_class.__setattr__ = replacement_setattr
+    the_class._protected_setattr = True
+
 
 def handle_IDSelectorSubset(the_class, class_owns, force_int64=True):
     the_class.original_init = the_class.__init__

tests/test_referenced_objects.py

@@ -84,6 +84,77 @@ def test_shards(self):
         gc.collect()
         index.search(xb, 10)
 
+    def test_SearchParameters_setattr_no_leak(self):
+        # Regression: assigning IDSelectorBatch to params.sel after
+        # construction must not leak. Before the fix, the SWIG property
+        # setter flipped IDSelectorBatch.thisown to 0 but
+        # SearchParameters never freed it -> ~80 MB orphaned per iter.
+        n_ids = 5_000_000
+        n_iters = 10
+
+        def body():
+            params = faiss.SearchParameters()
+            params.sel = faiss.IDSelectorBatch(np.arange(n_ids))
+            del params
+
+        body()  # warmup: prime allocator
+        gc.collect()
+        start_kb = faiss.get_mem_usage_kb()
+        for _ in range(n_iters):
+            body()
+            gc.collect()
+        growth_mb = (faiss.get_mem_usage_kb() - start_kb) / 1024.0
+        # One leaked IDSelectorBatch is ~80 MB; 10 iters would leak
+        # ~800 MB. 200 MB cap leaves generous headroom for allocator
+        # noise but is well below the leak signal.
+        self.assertLess(growth_mb, 200)
+
+    def test_SearchParameters_subclass_no_double_wrap(self):
+        # Regression: SWIG does not generate per-class __setattr__, so
+        # the ownership-protecting override must be installed only on
+        # the base SearchParameters; otherwise subclasses inherit and
+        # re-wrap, doubling the refcount on the assigned selector.
+        params = faiss.SearchParametersIVF()
+        sel = faiss.IDSelectorBatch(np.arange(10))
+        before = sys.getrefcount(sel)
+        params.sel = sel
+        delta = sys.getrefcount(sel) - before
+        self.assertEqual(delta, 1)
+
+    def test_SearchParameters_reassignment_no_accumulation(self):
+        # Reassigning the same field on a long-lived SearchParameters
+        # must release the prior ref. Without per-field tracking the
+        # protection list grows unbounded - reproducing the original
+        # leak shape under a different access pattern.
+        params = faiss.SearchParameters()
+        sel1 = faiss.IDSelectorBatch(np.arange(10))
+        sel2 = faiss.IDSelectorBatch(np.arange(10))
+        before1 = sys.getrefcount(sel1)
+        params.sel = sel1
+        params.sel = sel2
+        # sel1 ref dropped when sel2 took its slot
+        self.assertEqual(sys.getrefcount(sel1), before1)
+
+    def test_SearchParameters_setattr_else_branch(self):
+        # Else branch coverage: setting a field to None drops the prior
+        # SWIG ref; assigning two different fields keeps both alive
+        # independently so dropping one does not affect the other.
+        params = faiss.SearchParametersIVF()
+        sel = faiss.IDSelectorBatch(np.arange(10))
+        quant = faiss.SearchParameters()
+        sel_before = sys.getrefcount(sel)
+        quant_before = sys.getrefcount(quant)
+
+        params.sel = sel
+        params.quantizer_params = quant
+        self.assertEqual(sys.getrefcount(sel) - sel_before, 1)
+        self.assertEqual(sys.getrefcount(quant) - quant_before, 1)
+
+        # Drop sel via None; quantizer_params untouched.
+        params.sel = None
+        self.assertEqual(sys.getrefcount(sel), sel_before)
+        self.assertEqual(sys.getrefcount(quant) - quant_before, 1)
+
 
 dbin = 32
 xtbin = np.random.randint(256, size=(100, int(dbin / 8))).astype('uint8')