[xgb] Fix xgb intern to close replaced array

deepjavalibrary · Dec 10, 2024 · ddfae0e · ddfae0e
1 parent e4d184d
commit ddfae0e
Showing 1 changed file with 27 additions and 6 deletions.
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
@@ -88,14 +88,35 @@ public ByteBuffer toByteBuffer(boolean tryDirect) {
     /** {@inheritDoc} */
     @Override
     public void intern(NDArray replaced) {
-        if (handle != null && handle.get() != 0L) {
-            long pointer = handle.getAndSet(0L);
-            JniUtils.deleteDMatrix(pointer);
+        if (replaced == null) {
+            throw new IllegalArgumentException("The replaced NDArray cannot be null.");
+        }
+        if (!(replaced instanceof XgbNDArray)) {
+            throw new IllegalArgumentException("The replaced NDArray must be an instance of XgbNDArray.");
         }
         XgbNDArray array = (XgbNDArray) replaced;
-        data = array.data;
-        handle = array.handle;
-        format = array.format;
+
+        synchronized (this) {
+            if (handle != null && handle.get() != 0L) {
+                long pointer = handle.getAndSet(0L);
+                JniUtils.deleteDMatrix(pointer);
+            }
+
+            data = array.data;
+            format = array.format;
+
+            if (array.handle != null && array.handle.get() != 0L) {
+                if (handle == null) {
+                    handle = new AtomicLong();
+                }
+                handle.set(array.handle.getAndSet(0L));
+            }
+        }
+
+        array.data = null;
+        array.handle = null;
+        array.format = null;
+        array.close();
     }
 
     /** {@inheritDoc} */