Bulk allocate objects for nmslib index creation

CHANGELOG.md

@@ -16,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x](https://github.com/opensearch-project/k-NN/compare/2.6...2.x)
 ### Features
 ### Enhancements
+* Bulk allocate objects for nmslib index creation to avoid malloc fragmentation ([#773](https://github.com/opensearch-project/k-NN/pull/773))
 ### Bug Fixes
 ### Infrastructure
 * Adding filter type to filtering release configs ([#792](https://github.com/opensearch-project/k-NN/pull/792))

jni/CMakeLists.txt

@@ -66,7 +66,7 @@ target_include_directories(${TARGET_LIB_COMMON} PRIVATE ${CMAKE_CURRENT_SOURCE_D
 set_target_properties(${TARGET_LIB_COMMON} PROPERTIES SUFFIX ${LIB_EXT})
 set_target_properties(${TARGET_LIB_COMMON} PROPERTIES POSITION_INDEPENDENT_CODE ON)
 
-if (WIN32)
+if (NOT "${WIN32}" STREQUAL "")
 # Use RUNTIME_OUTPUT_DIRECTORY, to build the target library (opensearchknn_common) in the specified directory at runtime.
     set_target_properties(${TARGET_LIB_COMMON} PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
 else()
@@ -152,7 +152,7 @@ if (${CONFIG_FAISS} STREQUAL ON OR ${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} S
     set_target_properties(${TARGET_LIB_FAISS} PROPERTIES SUFFIX ${LIB_EXT})
     set_target_properties(${TARGET_LIB_FAISS} PROPERTIES POSITION_INDEPENDENT_CODE ON)
 
-    if (WIN32)
+    if (NOT "${WIN32}" STREQUAL "")
     # Use RUNTIME_OUTPUT_DIRECTORY, to build the target library (opensearchknn_faiss) in the specified directory at runtime.
         set_target_properties(${TARGET_LIB_FAISS} PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/release)
     else()
@@ -167,7 +167,7 @@ endif ()
 # --------------------------------- TESTS -----------------------------------
 # Windows : Comment the TESTS for now because the tests are failing(failing to build jni_tests.exe) if we are building our target libraries as SHARED libraries.
 # TODO: Fix the failing JNI TESTS on Windows
-if (!WIN32)
+if ("${WIN32}" STREQUAL "")
     if (${CONFIG_ALL} STREQUAL ON OR ${CONFIG_TEST} STREQUAL ON)
         # Reference - https://crascit.com/2015/07/25/cmake-gtest/
         configure_file(CMakeLists.txt.in googletest-download/CMakeLists.txt)

jni/src/nmslib_wrapper.cpp

@@ -27,6 +27,10 @@
 
 std::string TranslateSpaceType(const std::string& spaceType);
 
+// We do not use label functionality of nmslib so we pass default label. Setting as a const allows us to avoid a few
+// allocations
+const similarity::LabelType DEFAULT_LABEL = -1;
+
 void knn_jni::nmslib_wrapper::CreateIndex(knn_jni::JNIUtilInterface * jniUtil, JNIEnv * env, jintArray idsJ,
                                           jobjectArray vectorsJ, jstring indexPathJ, jobject parametersJ) {
 
@@ -96,24 +100,49 @@ void knn_jni::nmslib_wrapper::CreateIndex(knn_jni::JNIUtilInterface * jniUtil, J
 
     // Read dataset
     similarity::ObjectVector dataset;
+    dataset.reserve(numVectors);
     int* idsCpp;
     try {
         // Read in data set
         idsCpp = jniUtil->GetIntArrayElements(env, idsJ, nullptr);
 
         float* floatArrayCpp;
         jfloatArray floatArrayJ;
+        size_t vectorSizeInBytes = dim*sizeof(float);
+
+        // Allocate a large buffer that will contain all the vectors. Allocating the objects in one large buffer as
+        // opposed to individually will prevent heap fragmentation. We have observed that allocating individual
+        // objects causes RSS to rise throughout the lifetime of a process
+        // (see https://github.com/opensearch-project/k-NN/issues/772 and
+        // https://github.com/opensearch-project/k-NN/issues/72). This is because, in typical systems, small
+        // allocations will reside on some kind of heap managed by an allocator. Once freed, the allocator does not
+        // always return the memory to the OS. If the heap gets fragmented, this will cause the allocator
+        // to ask for more memory, causing RSS to grow. On large allocations (> 128 kb), most allocators will
+        // internally use mmap. Once freed, unmap will be called, which will immediately return memory to the OS
+        // which in turn prevents RSS from growing out of control. Wrap with a smart pointer so that buffer will be
+        // freed once variable goes out of scope. For reference, the code that specifies the layout of the buffer can be
+        // found: https://github.com/nmslib/nmslib/blob/v2.1.1/similarity_search/include/object.h#L61-L75
+        std::unique_ptr<char[]> objectBuffer(new char[(similarity::ID_SIZE + similarity::LABEL_SIZE + similarity::DATALENGTH_SIZE + vectorSizeInBytes) * numVectors]);
+        char* ptr = objectBuffer.get();
         for (int i = 0; i < numVectors; i++) {
-            floatArrayJ = (jfloatArray)jniUtil->GetObjectArrayElement(env, vectorsJ, i);
+            dataset.push_back(new similarity::Object(ptr));
+
+            memcpy(ptr, &idsCpp[i], similarity::ID_SIZE);
+            ptr += similarity::ID_SIZE;
+            memcpy(ptr, &DEFAULT_LABEL, similarity::LABEL_SIZE);
+            ptr += similarity::LABEL_SIZE;
+            memcpy(ptr, &vectorSizeInBytes, similarity::DATALENGTH_SIZE);
+            ptr += similarity::DATALENGTH_SIZE;
 
+            floatArrayJ = (jfloatArray)jniUtil->GetObjectArrayElement(env, vectorsJ, i);
             if (dim != jniUtil->GetJavaFloatArrayLength(env, floatArrayJ)) {
                 throw std::runtime_error("Dimension of vectors is inconsistent");
             }
 
             floatArrayCpp = jniUtil->GetFloatArrayElements(env, floatArrayJ, nullptr);
-
-            dataset.push_back(new similarity::Object(idsCpp[i], -1, dim*sizeof(float), floatArrayCpp));
+            memcpy(ptr, floatArrayCpp, vectorSizeInBytes);
             jniUtil->ReleaseFloatArrayElements(env, floatArrayJ, floatArrayCpp, JNI_ABORT);
+            ptr += vectorSizeInBytes;
         }
         jniUtil->ReleaseIntArrayElements(env, idsJ, idsCpp, JNI_ABORT);