jbonofre commented on code in PR #1362: URL: https://github.com/apache/arrow-java/pull/1362#discussion_r4465455965
########## memory/memory-core/src/test/java/org/apache/arrow/memory/TestBufferLedger.java: ########## @@ -0,0 +1,210 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.memory; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CyclicBarrier; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +/** Tests for the reference counting in {@link BufferLedger}. */ +public class TestBufferLedger { + + /** + * Regression test for https://github.com/apache/arrow-java/issues/906: retain on a closed buffer + * must not revive it. + */ + @Test + public void testRetainAfterCloseThrowsEveryTime() { + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf buf = allocator.buffer(10); + buf.close(); + assertEquals(0, buf.refCnt()); + + for (int i = 0; i < 3; i++) { + assertThrows(IllegalArgumentException.class, () -> buf.getReferenceManager().retain()); + assertEquals(0, buf.refCnt()); + assertThrows(IllegalArgumentException.class, () -> buf.getReferenceManager().retain(5)); + assertEquals(0, buf.refCnt()); + } + } + } + + /** + * After a failed retain the released memory stays released: it is neither usable nor freed again. + */ + @Test + public void testFailedRetainDoesNotReviveReleasedMemory() { + try (BufferAllocator allocator = new RootAllocator()) { + ArrowBuf buf = allocator.buffer(10); + buf.close(); + assertEquals(0, allocator.getAllocatedMemory()); + + assertThrows(IllegalArgumentException.class, () -> buf.getReferenceManager().retain()); + assertThrows(IllegalArgumentException.class, () -> buf.getReferenceManager().retain()); + + // still closed: no access to memory that may already belong to another buffer + assertThrows(IllegalStateException.class, () -> buf.getByte(0)); + // and a stray release cannot free the chunk a second time + assertThrows(IllegalStateException.class, () -> buf.getReferenceManager().release()); Review Comment: This stray `release()` throws, but `decrement()` has already applied `addAndGet(-1)`, so `buf.refCnt()` is now -1. `ArrowBuf.ensureAccessible()` only rejects `refCnt() == 0`, so after this line `getByte()` and `setByte()` pass the check and touch freed memory. It's the same revival this PR fixes for `retain`, through `release` this time. Adding this after the line should fail today: ```java assertEquals(0, buf.refCnt()); assertThrows(IllegalStateException.class, () -> buf.getByte(0)); ``` The fix mirrors your `retain` change: in `decrement()`, check that `bufRefCnt >= decrement` before substracting. It already runs under the `allocationManager` lock, so it doesn't need a CAS loop. ########## memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java: ########## @@ -101,6 +101,12 @@ private BufferLedger associate(final BufferAllocator allocator, final boolean re "A buffer can only be associated between two allocators that share the same root"); synchronized (this) { + // Once the last ledger is released the memory is freed (owningLedger is null), and a new + // association would revive it (https://github.com/apache/arrow-java/issues/906). The final + // release frees the memory under this lock, so it can't happen between this check and the + // association. The constructor's own association runs before owningLedger is set. + Preconditions.checkState( + !retain || owningLedger != null, "The ArrowBuf's memory has already been released"); Review Comment: This check treats `owningLedge == null` as "memory freed", but `release()` only clears it at the end (line 177), after `releaseBytes()`, `release0()` and `listener.onRelease()`. If any of those throws (a user `AllocationListener.onRelease`, for example), the memory is already freed while `owningLedger` still points at the old ledger. This check then passes, and `retain(buf, other)` or `transferOwnership(buf, other)` returns a live buffer over freed memory, which is the #906 scenario again. `transferOwnership` would also call `releaseBytes` a second time on the old allocator. Could we clear it before freeing? In `release()`: ```java owningLedger = null; oldAllocator.releaseBytes(getSize()); release0(); oldAllocator.getListener().onRelease(getSize()); ``` None of those three calls read `owningLedger`. If `release0()` throws, we leak memory instead of allowing a use-after-free, which is the safer failure. A test with a listener whose `onRelease` throws would cover this. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
