From fb1dbaff481a1582696f8565a284ac8d04e497cf Mon Sep 17 00:00:00 2001 From: Jason Lowe Date: Thu, 22 Feb 2024 16:26:11 -0600 Subject: [PATCH] Java: Add leak tracking for Scalar instances Signed-off-by: Jason Lowe --- java/src/main/java/ai/rapids/cudf/MemoryCleaner.java | 7 ++++++- java/src/main/java/ai/rapids/cudf/Scalar.java | 4 +++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java b/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java index 032b075bab7..4614ce24024 100644 --- a/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java +++ b/java/src/main/java/ai/rapids/cudf/MemoryCleaner.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -261,6 +261,11 @@ static void register(ColumnVector vec, Cleaner cleaner) { all.put(cleaner.id, new CleanerWeakReference(vec, cleaner, collected, true)); } + static void register(Scalar s, Cleaner cleaner) { + // It is now registered... + all.put(cleaner.id, new CleanerWeakReference(s, cleaner, collected, true)); + } + static void register(HostColumnVectorCore vec, Cleaner cleaner) { // It is now registered... all.put(cleaner.id, new CleanerWeakReference(vec, cleaner, collected, false)); diff --git a/java/src/main/java/ai/rapids/cudf/Scalar.java b/java/src/main/java/ai/rapids/cudf/Scalar.java index 70538ab082f..286b5c208c9 100644 --- a/java/src/main/java/ai/rapids/cudf/Scalar.java +++ b/java/src/main/java/ai/rapids/cudf/Scalar.java @@ -1,6 +1,6 @@ /* * - * Copyright (c) 2019-2022, NVIDIA CORPORATION. + * Copyright (c) 2019-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -524,6 +524,7 @@ private static ColumnVector buildNullColumnVector(HostColumnVector.DataType host Scalar(DType type, long scalarHandle) { this.type = type; this.offHeap = new OffHeapState(scalarHandle); + MemoryCleaner.register(this, offHeap); incRefCount(); } @@ -536,6 +537,7 @@ public synchronized Scalar incRefCount() { offHeap.logRefCountDebug("INC AFTER CLOSE " + this); throw new IllegalStateException("Scalar is already closed"); } + offHeap.addRef(); ++refCount; return this; }