diff options
Diffstat (limited to 'src/arrow/java/memory')
82 files changed, 14294 insertions, 0 deletions
diff --git a/src/arrow/java/memory/memory-core/pom.xml b/src/arrow/java/memory/memory-core/pom.xml new file mode 100644 index 000000000..eb1226a9e --- /dev/null +++ b/src/arrow/java/memory/memory-core/pom.xml @@ -0,0 +1,40 @@ +<?xml version="1.0" encoding="UTF-8"?> +<!-- 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. --> +<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> + <parent> + <artifactId>arrow-memory</artifactId> + <groupId>org.apache.arrow</groupId> + <version>6.0.1</version> + </parent> + <modelVersion>4.0.0</modelVersion> + + <artifactId>arrow-memory-core</artifactId> + + <name>Arrow Memory - Core</name> + <description>Core off-heap memory management libraries for Arrow ValueVectors.</description> + + <dependencies> + <dependency> + <groupId>com.google.code.findbugs</groupId> + <artifactId>jsr305</artifactId> + </dependency> + <dependency> + <groupId>org.slf4j</groupId> + <artifactId>slf4j-api</artifactId> + </dependency> + <dependency> + <groupId>org.immutables</groupId> + <artifactId>value</artifactId> + </dependency> + </dependencies> + +</project> diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java new file mode 100644 index 000000000..42dac7b8c --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/Accountant.java @@ -0,0 +1,308 @@ +/* + * 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 java.util.concurrent.atomic.AtomicLong; + +import javax.annotation.concurrent.ThreadSafe; + +import org.apache.arrow.util.Preconditions; + +/** + * Provides a concurrent way to manage account for memory usage without locking. Used as basis + * for Allocators. All + * operations are threadsafe (except for close). + */ +@ThreadSafe +class Accountant implements AutoCloseable { + + /** + * The parent allocator. + */ + protected final Accountant parent; + + private final String name; + + /** + * The amount of memory reserved for this allocator. Releases below this amount of memory will + * not be returned to the + * parent Accountant until this Accountant is closed. + */ + protected final long reservation; + + private final AtomicLong peakAllocation = new AtomicLong(); + + /** + * Maximum local memory that can be held. This can be externally updated. Changing it won't + * cause past memory to + * change but will change responses to future allocation efforts + */ + private final AtomicLong allocationLimit = new AtomicLong(); + + /** + * Currently allocated amount of memory. + */ + private final AtomicLong locallyHeldMemory = new AtomicLong(); + + public Accountant(Accountant parent, String name, long reservation, long maxAllocation) { + Preconditions.checkNotNull(name, "name must not be null"); + Preconditions.checkArgument(reservation >= 0, "The initial reservation size must be non-negative."); + Preconditions.checkArgument(maxAllocation >= 0, "The maximum allocation limit must be non-negative."); + Preconditions.checkArgument(reservation <= maxAllocation, + "The initial reservation size must be <= the maximum allocation."); + Preconditions.checkArgument(reservation == 0 || parent != null, "The root accountant can't reserve memory."); + + this.parent = parent; + this.name = name; + this.reservation = reservation; + this.allocationLimit.set(maxAllocation); + + if (reservation != 0) { + // we will allocate a reservation from our parent. + final AllocationOutcome outcome = parent.allocateBytes(reservation); + if (!outcome.isOk()) { + throw new OutOfMemoryException(String.format( + "Failure trying to allocate initial reservation for Allocator. " + + "Attempted to allocate %d bytes.", reservation), outcome.getDetails()); + } + } + } + + /** + * Attempt to allocate the requested amount of memory. Either completely succeeds or completely + * fails. If it fails, no changes are made to accounting. + * + * @param size The amount of memory to reserve in bytes. + * @return the status and details of allocation at each allocator in the chain. + */ + AllocationOutcome allocateBytes(long size) { + AllocationOutcome.Status status = allocateBytesInternal(size); + if (status.isOk()) { + return AllocationOutcome.SUCCESS_INSTANCE; + } else { + // Try again, but with details this time. + // Populating details only on failures avoids performance overhead in the common case (success case). + AllocationOutcomeDetails details = new AllocationOutcomeDetails(); + status = allocateBytesInternal(size, details); + return new AllocationOutcome(status, details); + } + } + + private AllocationOutcome.Status allocateBytesInternal(long size, AllocationOutcomeDetails details) { + final AllocationOutcome.Status status = allocate(size, + true /*incomingUpdatePeek*/, false /*forceAllocation*/, details); + if (!status.isOk()) { + releaseBytes(size); + } + return status; + } + + private AllocationOutcome.Status allocateBytesInternal(long size) { + return allocateBytesInternal(size, null /*details*/); + } + + private void updatePeak() { + final long currentMemory = locallyHeldMemory.get(); + while (true) { + + final long previousPeak = peakAllocation.get(); + if (currentMemory > previousPeak) { + if (!peakAllocation.compareAndSet(previousPeak, currentMemory)) { + // peak allocation changed underneath us. try again. + continue; + } + } + + // we either succeeded to set peak allocation or we weren't above the previous peak, exit. + return; + } + } + + + /** + * Increase the accounting. Returns whether the allocation fit within limits. + * + * @param size to increase + * @return Whether the allocation fit within limits. + */ + public boolean forceAllocate(long size) { + final AllocationOutcome.Status outcome = allocate(size, true, true, null); + return outcome.isOk(); + } + + /** + * Internal method for allocation. This takes a forced approach to allocation to ensure that we + * manage reservation + * boundary issues consistently. Allocation is always done through the entire tree. The two + * options that we influence + * are whether the allocation should be forced and whether or not the peak memory allocation + * should be updated. If at + * some point during allocation escalation we determine that the allocation is no longer + * possible, we will continue to + * do a complete and consistent allocation but we will stop updating the peak allocation. We do + * this because we know + * that we will be directly unwinding this allocation (and thus never actually making the + * allocation). If force + * allocation is passed, then we continue to update the peak limits since we now know that this + * allocation will occur + * despite our moving past one or more limits. + * + * @param size The size of the allocation. + * @param incomingUpdatePeak Whether we should update the local peak for this allocation. + * @param forceAllocation Whether we should force the allocation. + * @return The outcome of the allocation. + */ + private AllocationOutcome.Status allocate(final long size, final boolean incomingUpdatePeak, + final boolean forceAllocation, AllocationOutcomeDetails details) { + final long newLocal = locallyHeldMemory.addAndGet(size); + final long beyondReservation = newLocal - reservation; + final boolean beyondLimit = newLocal > allocationLimit.get(); + final boolean updatePeak = forceAllocation || (incomingUpdatePeak && !beyondLimit); + + if (details != null) { + // Add details if required (used in exceptions and debugging). + boolean allocationFailed = true; + long allocatedLocal = 0; + if (!beyondLimit) { + allocatedLocal = size - Math.min(beyondReservation, size); + allocationFailed = false; + } + details.pushEntry(this, newLocal - size, size, allocatedLocal, allocationFailed); + } + + AllocationOutcome.Status parentOutcome = AllocationOutcome.Status.SUCCESS; + if (beyondReservation > 0 && parent != null) { + // we need to get memory from our parent. + final long parentRequest = Math.min(beyondReservation, size); + parentOutcome = parent.allocate(parentRequest, updatePeak, forceAllocation, details); + } + + final AllocationOutcome.Status finalOutcome; + if (beyondLimit) { + finalOutcome = AllocationOutcome.Status.FAILED_LOCAL; + } else { + finalOutcome = parentOutcome.isOk() ? AllocationOutcome.Status.SUCCESS + : AllocationOutcome.Status.FAILED_PARENT; + } + + if (updatePeak) { + updatePeak(); + } + + return finalOutcome; + } + + public void releaseBytes(long size) { + // reduce local memory. all memory released above reservation should be released up the tree. + final long newSize = locallyHeldMemory.addAndGet(-size); + + Preconditions.checkArgument(newSize >= 0, "Accounted size went negative."); + + final long originalSize = newSize + size; + if (originalSize > reservation && parent != null) { + // we deallocated memory that we should release to our parent. + final long possibleAmountToReleaseToParent = originalSize - reservation; + final long actualToReleaseToParent = Math.min(size, possibleAmountToReleaseToParent); + parent.releaseBytes(actualToReleaseToParent); + } + } + + public boolean isOverLimit() { + return getAllocatedMemory() > getLimit() || (parent != null && parent.isOverLimit()); + } + + /** + * Close this Accountant. This will release any reservation bytes back to a parent Accountant. + */ + @Override + public void close() { + // return memory reservation to parent allocator. + if (parent != null) { + parent.releaseBytes(reservation); + } + } + + /** + * Return the name of the accountant. + * + * @return name of accountant + */ + public String getName() { + return name; + } + + /** + * Return the current limit of this Accountant. + * + * @return Limit in bytes. + */ + public long getLimit() { + return allocationLimit.get(); + } + + /** + * Return the initial reservation. + * + * @return reservation in bytes. + */ + public long getInitReservation() { + return reservation; + } + + /** + * Set the maximum amount of memory that can be allocated in the this Accountant before failing + * an allocation. + * + * @param newLimit The limit in bytes. + */ + public void setLimit(long newLimit) { + allocationLimit.set(newLimit); + } + + /** + * Return the current amount of allocated memory that this Accountant is managing accounting + * for. Note this does not + * include reservation memory that hasn't been allocated. + * + * @return Currently allocate memory in bytes. + */ + public long getAllocatedMemory() { + return locallyHeldMemory.get(); + } + + /** + * The peak memory allocated by this Accountant. + * + * @return The peak allocated memory in bytes. + */ + public long getPeakMemoryAllocation() { + return peakAllocation.get(); + } + + public long getHeadroom() { + long localHeadroom = allocationLimit.get() - locallyHeldMemory.get(); + if (parent == null) { + return localHeadroom; + } + + // Amount of reserved memory left on top of what parent has + long reservedHeadroom = Math.max(0, reservation - locallyHeldMemory.get()); + return Math.min(localHeadroom, parent.getHeadroom() + reservedHeadroom); + } + +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationListener.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationListener.java new file mode 100644 index 000000000..ff2b25dfa --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationListener.java @@ -0,0 +1,85 @@ +/* + * 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; + +/** + * An allocation listener being notified for allocation/deallocation + * + * <p>It might be called from multiple threads if the allocator hierarchy shares a listener, in which + * case, the provider should take care of making the implementation thread-safe. + */ +public interface AllocationListener { + + AllocationListener NOOP = new AllocationListener() {}; + + /** + * Called each time a new buffer has been requested. + * + * <p>An exception can be safely thrown by this method to terminate the allocation. + * + * @param size the buffer size being allocated + */ + default void onPreAllocation(long size) {} + + /** + * Called each time a new buffer has been allocated. + * + * <p>An exception cannot be thrown by this method. + * + * @param size the buffer size being allocated + */ + default void onAllocation(long size) {} + + /** + * Informed each time a buffer is released from allocation. + * + * <p>An exception cannot be thrown by this method. + * @param size The size of the buffer being released. + */ + default void onRelease(long size) {} + + + /** + * Called whenever an allocation failed, giving the caller a chance to create some space in the + * allocator (either by freeing some resource, or by changing the limit), and, if successful, + * allowing the allocator to retry the allocation. + * + * @param size the buffer size that was being allocated + * @param outcome the outcome of the failed allocation. Carries information of what failed + * @return true, if the allocation can be retried; false if the allocation should fail + */ + default boolean onFailedAllocation(long size, AllocationOutcome outcome) { + return false; + } + + /** + * Called immediately after a child allocator was added to the parent allocator. + * + * @param parentAllocator The parent allocator to which a child was added + * @param childAllocator The child allocator that was just added + */ + default void onChildAdded(BufferAllocator parentAllocator, BufferAllocator childAllocator) {} + + /** + * Called immediately after a child allocator was removed from the parent allocator. + * + * @param parentAllocator The parent allocator from which a child was removed + * @param childAllocator The child allocator that was just removed + */ + default void onChildRemoved(BufferAllocator parentAllocator, BufferAllocator childAllocator) {} +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java new file mode 100644 index 000000000..5f8ab1244 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java @@ -0,0 +1,221 @@ +/* + * 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 java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.util.Preconditions; + +/** + * The abstract base class of AllocationManager. + * + * <p>Manages the relationship between one or more allocators and a particular UDLE. Ensures that + * one allocator owns the + * memory that multiple allocators may be referencing. Manages a BufferLedger between each of its + * associated allocators. + * + * <p>The only reason that this isn't package private is we're forced to put ArrowBuf in Netty's + * package which need access + * to these objects or methods. + * + * <p>Threading: AllocationManager manages thread-safety internally. Operations within the context + * of a single BufferLedger + * are lockless in nature and can be leveraged by multiple threads. Operations that cross the + * context of two ledgers + * will acquire a lock on the AllocationManager instance. Important note, there is one + * AllocationManager per + * UnsafeDirectLittleEndian buffer allocation. As such, there will be thousands of these in a + * typical query. The + * contention of acquiring a lock on AllocationManager should be very low. + */ +public abstract class AllocationManager { + + private static final AtomicLong MANAGER_ID_GENERATOR = new AtomicLong(0); + + private final BufferAllocator root; + private final long allocatorManagerId = MANAGER_ID_GENERATOR.incrementAndGet(); + // ARROW-1627 Trying to minimize memory overhead caused by previously used IdentityHashMap + // see JIRA for details + private final LowCostIdentityHashMap<BufferAllocator, BufferLedger> map = new LowCostIdentityHashMap<>(); + private final long amCreationTime = System.nanoTime(); + + // The ReferenceManager created at the time of creation of this AllocationManager + // is treated as the owning reference manager for the underlying chunk of memory + // managed by this allocation manager + private volatile BufferLedger owningLedger; + private volatile long amDestructionTime = 0; + + protected AllocationManager(BufferAllocator accountingAllocator) { + Preconditions.checkNotNull(accountingAllocator); + accountingAllocator.assertOpen(); + + this.root = accountingAllocator.getRoot(); + + // we do a no retain association since our creator will want to retrieve the newly created + // ledger and will create a reference count at that point + this.owningLedger = associate(accountingAllocator, false); + } + + BufferLedger getOwningLedger() { + return owningLedger; + } + + void setOwningLedger(final BufferLedger ledger) { + this.owningLedger = ledger; + } + + /** + * Associate the existing underlying buffer with a new allocator. This will increase the + * reference count on the corresponding buffer ledger by 1 + * + * @param allocator The target allocator to associate this buffer with. + * @return The reference manager (new or existing) that associates the underlying + * buffer to this new ledger. + */ + BufferLedger associate(final BufferAllocator allocator) { + return associate(allocator, true); + } + + private BufferLedger associate(final BufferAllocator allocator, final boolean retain) { + allocator.assertOpen(); + Preconditions.checkState(root == allocator.getRoot(), + "A buffer can only be associated between two allocators that share the same root"); + + synchronized (this) { + BufferLedger ledger = map.get(allocator); + if (ledger != null) { + if (retain) { + // bump the ref count for the ledger + ledger.increment(); + } + return ledger; + } + + ledger = new BufferLedger(allocator, this); + + if (retain) { + // the new reference manager will have a ref count of 1 + ledger.increment(); + } + + // store the mapping for <allocator, reference manager> + BufferLedger oldLedger = map.put(ledger); + Preconditions.checkState(oldLedger == null, + "Detected inconsistent state: A reference manager already exists for this allocator"); + + if (allocator instanceof BaseAllocator) { + // needed for debugging only: keep a pointer to reference manager inside allocator + // to dump state, verify allocator state etc + ((BaseAllocator) allocator).associateLedger(ledger); + } + return ledger; + } + } + + /** + * The way that a particular ReferenceManager (BufferLedger) communicates back to the + * AllocationManager that it no longer needs to hold a reference to a particular + * piece of memory. Reference manager needs to hold a lock to invoke this method + * It is called when the shared refcount of all the ArrowBufs managed by the + * calling ReferenceManager drops to 0. + */ + void release(final BufferLedger ledger) { + final BufferAllocator allocator = ledger.getAllocator(); + allocator.assertOpen(); + + // remove the <BaseAllocator, BufferLedger> mapping for the allocator + // of calling BufferLedger + Preconditions.checkState(map.containsKey(allocator), + "Expecting a mapping for allocator and reference manager"); + final BufferLedger oldLedger = map.remove(allocator); + + BufferAllocator oldAllocator = oldLedger.getAllocator(); + if (oldAllocator instanceof BaseAllocator) { + // needed for debug only: tell the allocator that AllocationManager is removing a + // reference manager associated with this particular allocator + ((BaseAllocator) oldAllocator).dissociateLedger(oldLedger); + } + + if (oldLedger == owningLedger) { + // the release call was made by the owning reference manager + if (map.isEmpty()) { + // the only <allocator, reference manager> mapping was for the owner + // which now has been removed, it implies we can safely destroy the + // underlying memory chunk as it is no longer being referenced + oldAllocator.releaseBytes(getSize()); + // free the memory chunk associated with the allocation manager + release0(); + oldAllocator.getListener().onRelease(getSize()); + amDestructionTime = System.nanoTime(); + owningLedger = null; + } else { + // since the refcount dropped to 0 for the owning reference manager and allocation + // manager will no longer keep a mapping for it, we need to change the owning + // reference manager to whatever the next available <allocator, reference manager> + // mapping exists. + BufferLedger newOwningLedger = map.getNextValue(); + // we'll forcefully transfer the ownership and not worry about whether we + // exceeded the limit since this consumer can't do anything with this. + oldLedger.transferBalance(newOwningLedger); + } + } else { + // the release call was made by a non-owning reference manager, so after remove there have + // to be 1 or more <allocator, reference manager> mappings + Preconditions.checkState(map.size() > 0, + "The final removal of reference manager should be connected to owning reference manager"); + } + } + + /** + * Return the size of underlying chunk of memory managed by this Allocation Manager. + * + * <p>The underlying memory chunk managed can be different from the original requested size. + * + * @return size of underlying memory chunk + */ + public abstract long getSize(); + + /** + * Return the absolute memory address pointing to the fist byte of underlying memory chunk. + */ + protected abstract long memoryAddress(); + + /** + * Release the underlying memory chunk. + */ + protected abstract void release0(); + + /** + * A factory interface for creating {@link AllocationManager}. + * One may extend this interface to use a user-defined AllocationManager implementation. + */ + public interface Factory { + + /** + * Create an {@link AllocationManager}. + * + * @param accountingAllocator The allocator that are expected to be associated with newly created AllocationManager. + * Currently it is always equivalent to "this" + * @param size Size (in bytes) of memory managed by the AllocationManager + * @return The created AllocationManager used by this allocator + */ + AllocationManager create(BufferAllocator accountingAllocator, long size); + + ArrowBuf empty(); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationOutcome.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationOutcome.java new file mode 100644 index 000000000..2977775e6 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationOutcome.java @@ -0,0 +1,97 @@ +/* + * 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 java.util.Optional; + +/** + * Describes the type of outcome that occurred when trying to account for allocation of memory. + */ +public class AllocationOutcome { + private final Status status; + private final AllocationOutcomeDetails details; + static final AllocationOutcome SUCCESS_INSTANCE = new AllocationOutcome(Status.SUCCESS); + + AllocationOutcome(Status status, AllocationOutcomeDetails details) { + this.status = status; + this.details = details; + } + + AllocationOutcome(Status status) { + this(status, null); + } + + /** + * Get the status of the allocation. + * @return status code. + */ + public Status getStatus() { + return status; + } + + /** + * Get additional details of the allocation (like the status at each allocator in the hierarchy). + * @return details of allocation + */ + public Optional<AllocationOutcomeDetails> getDetails() { + return Optional.ofNullable(details); + } + + /** + * Returns true if the allocation was a success. + * @return true if allocation was successful, false otherwise. + */ + public boolean isOk() { + return status.isOk(); + } + + /** + * Allocation status code. + */ + public enum Status { + /** + * Allocation succeeded. + */ + SUCCESS(true), + + /** + * Allocation succeeded but only because the allocator was forced to move beyond a limit. + */ + FORCED_SUCCESS(true), + + /** + * Allocation failed because the local allocator's limits were exceeded. + */ + FAILED_LOCAL(false), + + /** + * Allocation failed because a parent allocator's limits were exceeded. + */ + FAILED_PARENT(false); + + private final boolean ok; + + Status(boolean ok) { + this.ok = ok; + } + + public boolean isOk() { + return ok; + } + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationOutcomeDetails.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationOutcomeDetails.java new file mode 100644 index 000000000..6499ce84b --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationOutcomeDetails.java @@ -0,0 +1,132 @@ +/* + * 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 java.util.ArrayDeque; +import java.util.Deque; + +/** + * Captures details of allocation for each accountant in the hierarchical chain. + */ +public class AllocationOutcomeDetails { + Deque<Entry> allocEntries; + + AllocationOutcomeDetails() { + allocEntries = new ArrayDeque<>(); + } + + void pushEntry(Accountant accountant, long totalUsedBeforeAllocation, long requestedSize, + long allocatedSize, boolean allocationFailed) { + + Entry top = allocEntries.peekLast(); + if (top != null && top.allocationFailed) { + // if the allocation has already failed, stop saving the entries. + return; + } + + allocEntries.addLast(new Entry(accountant, totalUsedBeforeAllocation, requestedSize, + allocatedSize, allocationFailed)); + } + + /** + * Get the allocator that caused the failure. + * @return the allocator that caused failure, null if there was no failure. + */ + public BufferAllocator getFailedAllocator() { + Entry top = allocEntries.peekLast(); + if (top != null && top.allocationFailed && (top.accountant instanceof BufferAllocator)) { + return (BufferAllocator) top.accountant; + } else { + return null; + } + } + + @Override + public String toString() { + StringBuilder sb = new StringBuilder(); + sb.append("Allocation outcome details:\n"); + allocEntries.forEach(sb::append); + return sb.toString(); + } + + /** + * Outcome of the allocation request at one accountant in the hierarchy. + */ + public static class Entry { + private final Accountant accountant; + + // Remember allocator attributes at the time of the request. + private final long limit; + private final long used; + + // allocation outcome + private final long requestedSize; + private final long allocatedSize; + private final boolean allocationFailed; + + Entry(Accountant accountant, long totalUsedBeforeAllocation, long requestedSize, + long allocatedSize, boolean allocationFailed) { + this.accountant = accountant; + this.limit = accountant.getLimit(); + this.used = totalUsedBeforeAllocation; + + this.requestedSize = requestedSize; + this.allocatedSize = allocatedSize; + this.allocationFailed = allocationFailed; + } + + public Accountant getAccountant() { + return accountant; + } + + public long getLimit() { + return limit; + } + + public long getUsed() { + return used; + } + + public long getRequestedSize() { + return requestedSize; + } + + public long getAllocatedSize() { + return allocatedSize; + } + + public boolean isAllocationFailed() { + return allocationFailed; + } + + @Override + public String toString() { + return new StringBuilder() + .append("allocator[" + accountant.getName() + "]") + .append(" reservation: " + accountant.getInitReservation()) + .append(" limit: " + limit) + .append(" used: " + used) + .append(" requestedSize: " + requestedSize) + .append(" allocatedSize: " + allocatedSize) + .append(" localAllocationStatus: " + (allocationFailed ? "fail" : "success")) + .append("\n") + .toString(); + } + } + +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationReservation.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationReservation.java new file mode 100644 index 000000000..4331eb20c --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationReservation.java @@ -0,0 +1,88 @@ +/* + * 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; + +/** + * Supports cumulative allocation reservation. Clients may increase the size of the reservation + * repeatedly until they + * call for an allocation of the current total size. The reservation can only be used once, and + * will throw an exception + * if it is used more than once. + * <p> + * For the purposes of airtight memory accounting, the reservation must be close()d whether it is + * used or not. + * This is not threadsafe. + * </p> + */ +public interface AllocationReservation extends AutoCloseable { + + /** + * Add to the current reservation. + * + * <p>Adding may fail if the allocator is not allowed to consume any more space.</p> + * + * @param nBytes the number of bytes to add + * @return true if the addition is possible, false otherwise + * @throws IllegalStateException if called after buffer() is used to allocate the reservation + */ + boolean add(int nBytes); + + /** + * Requests a reservation of additional space. + * + * <p>The implementation of the allocator's inner class provides this.</p> + * + * @param nBytes the amount to reserve + * @return true if the reservation can be satisfied, false otherwise + */ + boolean reserve(int nBytes); + + /** + * Allocate a buffer whose size is the total of all the add()s made. + * + * <p>The allocation request can still fail, even if the amount of space + * requested is available, if the allocation cannot be made contiguously.</p> + * + * @return the buffer, or null, if the request cannot be satisfied + * @throws IllegalStateException if called called more than once + */ + ArrowBuf allocateBuffer(); + + /** + * Get the current size of the reservation (the sum of all the add()s). + * + * @return size of the current reservation + */ + int getSize(); + + /** + * Return whether or not the reservation has been used. + * + * @return whether or not the reservation has been used + */ + boolean isUsed(); + + /** + * Return whether or not the reservation has been closed. + * + * @return whether or not the reservation has been closed + */ + boolean isClosed(); + + void close(); +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocatorClosedException.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocatorClosedException.java new file mode 100644 index 000000000..39c2ef82e --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocatorClosedException.java @@ -0,0 +1,35 @@ +/* + * 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; + +/** + * Exception thrown when a closed BufferAllocator is used. Note + * this is an unchecked exception. + */ +@SuppressWarnings("serial") +public class AllocatorClosedException extends RuntimeException { + + /** + * Constructs a new allocator closed exception with a given message. + * + * @param message string associated with the cause + */ + public AllocatorClosedException(String message) { + super(message); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java new file mode 100644 index 000000000..d7827073e --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java @@ -0,0 +1,1202 @@ +/* + * 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.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.ReadOnlyBufferException; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.memory.BaseAllocator.Verbosity; +import org.apache.arrow.memory.util.CommonUtil; +import org.apache.arrow.memory.util.HistoricalLog; +import org.apache.arrow.memory.util.MemoryUtil; +import org.apache.arrow.util.Preconditions; + +/** + * ArrowBuf serves as a facade over underlying memory by providing + * several access APIs to read/write data into a chunk of direct + * memory. All the accounting, ownership and reference management + * is done by {@link ReferenceManager} and ArrowBuf can work + * with a custom user provided implementation of ReferenceManager + * <p> + * Two important instance variables of an ArrowBuf: + * (1) address - starting virtual address in the underlying memory + * chunk that this ArrowBuf has access to + * (2) length - length (in bytes) in the underlying memory chunk + * that this ArrowBuf has access to + * </p> + * <p> + * The management (allocation, deallocation, reference counting etc) for + * the memory chunk is not done by ArrowBuf. + * Default implementation of ReferenceManager, allocation is in + * {@link BaseAllocator}, {@link BufferLedger} and {@link AllocationManager} + * </p> + */ +public final class ArrowBuf implements AutoCloseable { + + private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ArrowBuf.class); + + private static final int SHORT_SIZE = Short.BYTES; + private static final int INT_SIZE = Integer.BYTES; + private static final int FLOAT_SIZE = Float.BYTES; + private static final int DOUBLE_SIZE = Double.BYTES; + private static final int LONG_SIZE = Long.BYTES; + + private static final AtomicLong idGenerator = new AtomicLong(0); + private static final int LOG_BYTES_PER_ROW = 10; + private final long id = idGenerator.incrementAndGet(); + private final ReferenceManager referenceManager; + private final BufferManager bufferManager; + private final long addr; + private long readerIndex; + private long writerIndex; + private final HistoricalLog historicalLog = BaseAllocator.DEBUG ? + new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH, "ArrowBuf[%d]", id) : null; + private volatile long length; + + /** + * Constructs a new ArrowBuf. + * + * @param referenceManager The memory manager to track memory usage and reference count of this buffer + * @param length The byte length of this buffer + */ + public ArrowBuf( + final ReferenceManager referenceManager, + final BufferManager bufferManager, + final long length, + final long memoryAddress) { + this.referenceManager = referenceManager; + this.bufferManager = bufferManager; + this.addr = memoryAddress; + this.length = length; + this.readerIndex = 0; + this.writerIndex = 0; + if (BaseAllocator.DEBUG) { + historicalLog.recordEvent("create()"); + } + } + + public int refCnt() { + return referenceManager.getRefCount(); + } + + /** + * Allows a function to determine whether not reading a particular string of bytes is valid. + * + * <p>Will throw an exception if the memory is not readable for some reason. Only doesn't + * something in the case that + * AssertionUtil.BOUNDS_CHECKING_ENABLED is true. + * + * @param start The starting position of the bytes to be read. + * @param end The exclusive endpoint of the bytes to be read. + */ + public void checkBytes(long start, long end) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + checkIndexD(start, end - start); + } + } + + /** + * For get/set operations, reference count should be >= 1. + */ + private void ensureAccessible() { + if (this.refCnt() == 0) { + throw new IllegalStateException("Ref count should be >= 1 for accessing the ArrowBuf"); + } + } + + /** + * Get reference manager for this ArrowBuf. + * @return user provided implementation of {@link ReferenceManager} + */ + public ReferenceManager getReferenceManager() { + return referenceManager; + } + + public long capacity() { + return length; + } + + /** + * Adjusts the capacity of this buffer. Size increases are NOT supported. + * + * @param newCapacity Must be in in the range [0, length). + */ + public synchronized ArrowBuf capacity(long newCapacity) { + + if (newCapacity == length) { + return this; + } + + Preconditions.checkArgument(newCapacity >= 0); + + if (newCapacity < length) { + length = newCapacity; + return this; + } + + throw new UnsupportedOperationException("Buffers don't support resizing that increases the size."); + } + + /** + * Returns the byte order of elements in this buffer. + */ + public ByteOrder order() { + return ByteOrder.nativeOrder(); + } + + /** + * Returns the number of bytes still available to read in this buffer. + */ + public long readableBytes() { + Preconditions.checkState(writerIndex >= readerIndex, + "Writer index cannot be less than reader index"); + return writerIndex - readerIndex; + } + + /** + * Returns the number of bytes still available to write into this buffer before capacity is reached. + */ + public long writableBytes() { + return capacity() - writerIndex; + } + + /** + * Returns a slice of only the readable bytes in the buffer. + */ + public ArrowBuf slice() { + return slice(readerIndex, readableBytes()); + } + + /** + * Returns a slice (view) starting at <code>index</code> with the given <code>length</code>. + */ + public ArrowBuf slice(long index, long length) { + + Preconditions.checkPositionIndex(index, this.length); + Preconditions.checkPositionIndex(index + length, this.length); + + /* + * Re the behavior of reference counting, see http://netty.io/wiki/reference-counted-objects + * .html#wiki-h3-5, which + * explains that derived buffers share their reference count with their parent + */ + final ArrowBuf newBuf = referenceManager.deriveBuffer(this, index, length); + newBuf.writerIndex(length); + return newBuf; + } + + /** + * Make a nio byte buffer from this arrowbuf. + */ + public ByteBuffer nioBuffer() { + return nioBuffer(readerIndex, checkedCastToInt(readableBytes())); + } + + + /** + * Make a nio byte buffer from this ArrowBuf. + */ + public ByteBuffer nioBuffer(long index, int length) { + chk(index, length); + return getDirectBuffer(index, length); + } + + private ByteBuffer getDirectBuffer(long index, int length) { + long address = addr(index); + return MemoryUtil.directBuffer(address, length); + } + + public long memoryAddress() { + return this.addr; + } + + @Override + public String toString() { + return String.format("ArrowBuf[%d], address:%d, length:%d", id, memoryAddress(), length); + } + + @Override + public int hashCode() { + return System.identityHashCode(this); + } + + @Override + public boolean equals(Object obj) { + // identity equals only. + return this == obj; + } + + /* + * IMPORTANT NOTE + * The data getters and setters work with a caller provided + * index. This index is 0 based and since ArrowBuf has access + * to a portion of underlying chunk of memory starting at + * some address, we convert the given relative index into + * absolute index as memory address + index. + * + * Example: + * + * Let's say we have an underlying chunk of memory of length 64 bytes + * Now let's say we have an ArrowBuf that has access to the chunk + * from offset 4 for length of 16 bytes. + * + * If the starting virtual address of chunk is MAR, then memory + * address of this ArrowBuf is MAR + offset -- this is what is stored + * in variable addr. See the BufferLedger and AllocationManager code + * for the implementation of ReferenceManager that manages a + * chunk of memory and creates ArrowBuf with access to a range of + * bytes within the chunk (or the entire chunk) + * + * So now to get/set data, we will do => addr + index + * This logic is put in method addr(index) and is frequently + * used in get/set data methods to compute the absolute + * byte address for get/set operation in the underlying chunk + * + * @param index the index at which we the user wants to read/write + * @return the absolute address within the memory + */ + private long addr(long index) { + return addr + index; + } + + + + /*-------------------------------------------------* + | Following are a set of fast path data set and | + | get APIs to write/read data from ArrowBuf | + | at a given index (0 based relative to this | + | ArrowBuf and not relative to the underlying | + | memory chunk). | + | | + *-------------------------------------------------*/ + + + + /** + * Helper function to do bounds checking at a particular + * index for particular length of data. + * @param index index (0 based relative to this ArrowBuf) + * @param length provided length of data for get/set + */ + private void chk(long index, long length) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + checkIndexD(index, length); + } + } + + private void checkIndexD(long index, long fieldLength) { + // check reference count + ensureAccessible(); + // check bounds + Preconditions.checkArgument(fieldLength >= 0, "expecting non-negative data length"); + if (index < 0 || index > capacity() - fieldLength) { + if (BaseAllocator.DEBUG) { + historicalLog.logHistory(logger); + } + throw new IndexOutOfBoundsException(String.format( + "index: %d, length: %d (expected: range(0, %d))", index, fieldLength, capacity())); + } + } + + /** + * Get long value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return 8 byte long value + */ + public long getLong(long index) { + chk(index, LONG_SIZE); + return MemoryUtil.UNSAFE.getLong(addr(index)); + } + + /** + * Set long value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setLong(long index, long value) { + chk(index, LONG_SIZE); + MemoryUtil.UNSAFE.putLong(addr(index), value); + } + + /** + * Get float value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return 4 byte float value + */ + public float getFloat(long index) { + return Float.intBitsToFloat(getInt(index)); + } + + /** + * Set float value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setFloat(long index, float value) { + chk(index, FLOAT_SIZE); + MemoryUtil.UNSAFE.putInt(addr(index), Float.floatToRawIntBits(value)); + } + + /** + * Get double value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return 8 byte double value + */ + public double getDouble(long index) { + return Double.longBitsToDouble(getLong(index)); + } + + /** + * Set double value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setDouble(long index, double value) { + chk(index, DOUBLE_SIZE); + MemoryUtil.UNSAFE.putLong(addr(index), Double.doubleToRawLongBits(value)); + } + + /** + * Get char value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return 2 byte char value + */ + public char getChar(long index) { + return (char) getShort(index); + } + + /** + * Set char value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setChar(long index, int value) { + chk(index, SHORT_SIZE); + MemoryUtil.UNSAFE.putShort(addr(index), (short) value); + } + + /** + * Get int value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return 4 byte int value + */ + public int getInt(long index) { + chk(index, INT_SIZE); + return MemoryUtil.UNSAFE.getInt(addr(index)); + } + + /** + * Set int value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setInt(long index, int value) { + chk(index, INT_SIZE); + MemoryUtil.UNSAFE.putInt(addr(index), value); + } + + /** + * Get short value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return 2 byte short value + */ + public short getShort(long index) { + chk(index, SHORT_SIZE); + return MemoryUtil.UNSAFE.getShort(addr(index)); + } + + /** + * Set short value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setShort(long index, int value) { + setShort(index, (short) value); + } + + /** + * Set short value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setShort(long index, short value) { + chk(index, SHORT_SIZE); + MemoryUtil.UNSAFE.putShort(addr(index), value); + } + + /** + * Set byte value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setByte(long index, int value) { + chk(index, 1); + MemoryUtil.UNSAFE.putByte(addr(index), (byte) value); + } + + /** + * Set byte value at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be written + * @param value value to write + */ + public void setByte(long index, byte value) { + chk(index, 1); + MemoryUtil.UNSAFE.putByte(addr(index), value); + } + + /** + * Get byte value stored at a particular index in the + * underlying memory chunk this ArrowBuf has access to. + * @param index index (0 based relative to this ArrowBuf) + * where the value will be read from + * @return byte value + */ + public byte getByte(long index) { + chk(index, 1); + return MemoryUtil.UNSAFE.getByte(addr(index)); + } + + + + /*--------------------------------------------------* + | Following are another set of data set APIs | + | that directly work with writerIndex | + | | + *--------------------------------------------------*/ + + + + /** + * Helper function to do bound checking w.r.t writerIndex + * by checking if we can set "length" bytes of data at the + * writerIndex in this ArrowBuf. + * @param length provided length of data for set + */ + private void ensureWritable(final int length) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + Preconditions.checkArgument(length >= 0, "expecting non-negative length"); + // check reference count + this.ensureAccessible(); + // check bounds + if (length > writableBytes()) { + throw new IndexOutOfBoundsException( + String.format("writerIndex(%d) + length(%d) exceeds capacity(%d)", writerIndex, length, capacity())); + } + } + } + + /** + * Helper function to do bound checking w.r.t readerIndex + * by checking if we can read "length" bytes of data at the + * readerIndex in this ArrowBuf. + * @param length provided length of data for get + */ + private void ensureReadable(final int length) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + Preconditions.checkArgument(length >= 0, "expecting non-negative length"); + // check reference count + this.ensureAccessible(); + // check bounds + if (length > readableBytes()) { + throw new IndexOutOfBoundsException( + String.format("readerIndex(%d) + length(%d) exceeds writerIndex(%d)", readerIndex, length, writerIndex)); + } + } + } + + /** + * Read the byte at readerIndex. + * @return byte value + */ + public byte readByte() { + ensureReadable(1); + final byte b = getByte(readerIndex); + ++readerIndex; + return b; + } + + /** + * Read dst.length bytes at readerIndex into dst byte array + * @param dst byte array where the data will be written + */ + public void readBytes(byte[] dst) { + Preconditions.checkArgument(dst != null, "expecting valid dst bytearray"); + ensureReadable(dst.length); + getBytes(readerIndex, dst, 0, checkedCastToInt(dst.length)); + } + + /** + * Set the provided byte value at the writerIndex. + * @param value value to set + */ + public void writeByte(byte value) { + ensureWritable(1); + MemoryUtil.UNSAFE.putByte(addr(writerIndex), value); + ++writerIndex; + } + + /** + * Set the lower order byte for the provided value at + * the writerIndex. + * @param value value to be set + */ + public void writeByte(int value) { + ensureWritable(1); + MemoryUtil.UNSAFE.putByte(addr(writerIndex), (byte) value); + ++writerIndex; + } + + /** + * Write the bytes from given byte array into this + * ArrowBuf starting at writerIndex. + * @param src src byte array + */ + public void writeBytes(byte[] src) { + Preconditions.checkArgument(src != null, "expecting valid src array"); + writeBytes(src, 0, src.length); + } + + /** + * Write the bytes from given byte array starting at srcIndex + * into this ArrowBuf starting at writerIndex. + * @param src src byte array + * @param srcIndex index in the byte array where the copy will being from + * @param length length of data to copy + */ + public void writeBytes(byte[] src, int srcIndex, int length) { + ensureWritable(length); + setBytes(writerIndex, src, srcIndex, length); + writerIndex += length; + } + + /** + * Set the provided int value as short at the writerIndex. + * @param value value to set + */ + public void writeShort(int value) { + ensureWritable(SHORT_SIZE); + MemoryUtil.UNSAFE.putShort(addr(writerIndex), (short) value); + writerIndex += SHORT_SIZE; + } + + /** + * Set the provided int value at the writerIndex. + * @param value value to set + */ + public void writeInt(int value) { + ensureWritable(INT_SIZE); + MemoryUtil.UNSAFE.putInt(addr(writerIndex), value); + writerIndex += INT_SIZE; + } + + /** + * Set the provided long value at the writerIndex. + * @param value value to set + */ + public void writeLong(long value) { + ensureWritable(LONG_SIZE); + MemoryUtil.UNSAFE.putLong(addr(writerIndex), value); + writerIndex += LONG_SIZE; + } + + /** + * Set the provided float value at the writerIndex. + * @param value value to set + */ + public void writeFloat(float value) { + ensureWritable(FLOAT_SIZE); + MemoryUtil.UNSAFE.putInt(addr(writerIndex), Float.floatToRawIntBits(value)); + writerIndex += FLOAT_SIZE; + } + + /** + * Set the provided double value at the writerIndex. + * @param value value to set + */ + public void writeDouble(double value) { + ensureWritable(DOUBLE_SIZE); + MemoryUtil.UNSAFE.putLong(addr(writerIndex), Double.doubleToRawLongBits(value)); + writerIndex += DOUBLE_SIZE; + } + + + /*--------------------------------------------------* + | Following are another set of data set/get APIs | + | that read and write stream of bytes from/to byte | + | arrays, ByteBuffer, ArrowBuf etc | + | | + *--------------------------------------------------*/ + + /** + * Determine if the requested {@code index} and {@code length} will fit within {@code capacity}. + * @param index The starting index. + * @param length The length which will be utilized (starting from {@code index}). + * @param capacity The capacity that {@code index + length} is allowed to be within. + * @return {@code true} if the requested {@code index} and {@code length} will fit within {@code capacity}. + * {@code false} if this would result in an index out of bounds exception. + */ + private static boolean isOutOfBounds(long index, long length, long capacity) { + return (index | length | (index + length) | (capacity - (index + length))) < 0; + } + + private void checkIndex(long index, long fieldLength) { + // check reference count + this.ensureAccessible(); + // check bounds + if (isOutOfBounds(index, fieldLength, this.capacity())) { + throw new IndexOutOfBoundsException(String.format("index: %d, length: %d (expected: range(0, %d))", + index, fieldLength, this.capacity())); + } + } + + /** + * Copy data from this ArrowBuf at a given index in into destination + * byte array. + * @param index starting index (0 based relative to the portion of memory) + * this ArrowBuf has access to + * @param dst byte array to copy the data into + */ + public void getBytes(long index, byte[] dst) { + getBytes(index, dst, 0, dst.length); + } + + /** + * Copy data from this ArrowBuf at a given index into destination byte array. + * @param index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param dst byte array to copy the data into + * @param dstIndex starting index in dst byte array to copy into + * @param length length of data to copy from this ArrowBuf + */ + public void getBytes(long index, byte[] dst, int dstIndex, int length) { + // bound check for this ArrowBuf where the data will be copied from + checkIndex(index, length); + // null check + Preconditions.checkArgument(dst != null, "expecting a valid dst byte array"); + // bound check for dst byte array where the data will be copied to + if (isOutOfBounds(dstIndex, length, dst.length)) { + // not enough space to copy "length" bytes into dst array from dstIndex onwards + throw new IndexOutOfBoundsException("Not enough space to copy data into destination" + dstIndex); + } + if (length != 0) { + // copy "length" bytes from this ArrowBuf starting at addr(index) address + // into dst byte array at dstIndex onwards + MemoryUtil.UNSAFE.copyMemory(null, addr(index), dst, MemoryUtil.BYTE_ARRAY_BASE_OFFSET + dstIndex, length); + } + } + + /** + * Copy data from a given byte array into this ArrowBuf starting at + * a given index. + * @param index starting index (0 based relative to the portion of memory) + * this ArrowBuf has access to + * @param src byte array to copy the data from + */ + public void setBytes(long index, byte[] src) { + setBytes(index, src, 0, src.length); + } + + /** + * Copy data from a given byte array starting at the given source index into + * this ArrowBuf at a given index. + * @param index index (0 based relative to the portion of memory this ArrowBuf + * has access to) + * @param src src byte array to copy the data from + * @param srcIndex index in the byte array where the copy will start from + * @param length length of data to copy from byte array + */ + public void setBytes(long index, byte[] src, int srcIndex, long length) { + // bound check for this ArrowBuf where the data will be copied into + checkIndex(index, length); + // null check + Preconditions.checkArgument(src != null, "expecting a valid src byte array"); + // bound check for src byte array where the data will be copied from + if (isOutOfBounds(srcIndex, length, src.length)) { + // not enough space to copy "length" bytes into dst array from dstIndex onwards + throw new IndexOutOfBoundsException("Not enough space to copy data from byte array" + srcIndex); + } + if (length > 0) { + // copy "length" bytes from src byte array at the starting index (srcIndex) + // into this ArrowBuf starting at address "addr(index)" + MemoryUtil.UNSAFE.copyMemory(src, MemoryUtil.BYTE_ARRAY_BASE_OFFSET + srcIndex, null, addr(index), length); + } + } + + /** + * Copy data from this ArrowBuf at a given index into the destination + * ByteBuffer. + * @param index index (0 based relative to the portion of memory this ArrowBuf + * has access to) + * @param dst dst ByteBuffer where the data will be copied into + */ + public void getBytes(long index, ByteBuffer dst) { + // bound check for this ArrowBuf where the data will be copied from + checkIndex(index, dst.remaining()); + // dst.remaining() bytes of data will be copied into dst ByteBuffer + if (dst.remaining() != 0) { + // address in this ArrowBuf where the copy will begin from + final long srcAddress = addr(index); + if (dst.isDirect()) { + if (dst.isReadOnly()) { + throw new ReadOnlyBufferException(); + } + // copy dst.remaining() bytes of data from this ArrowBuf starting + // at address srcAddress into the dst ByteBuffer starting at + // address dstAddress + final long dstAddress = MemoryUtil.getByteBufferAddress(dst) + dst.position(); + MemoryUtil.UNSAFE.copyMemory(null, srcAddress, null, dstAddress, dst.remaining()); + // after copy, bump the next write position for the dst ByteBuffer + dst.position(dst.position() + dst.remaining()); + } else if (dst.hasArray()) { + // copy dst.remaining() bytes of data from this ArrowBuf starting + // at address srcAddress into the dst ByteBuffer starting at + // index dstIndex + final int dstIndex = dst.arrayOffset() + dst.position(); + MemoryUtil.UNSAFE.copyMemory( + null, srcAddress, dst.array(), MemoryUtil.BYTE_ARRAY_BASE_OFFSET + dstIndex, dst.remaining()); + // after copy, bump the next write position for the dst ByteBuffer + dst.position(dst.position() + dst.remaining()); + } else { + throw new UnsupportedOperationException("Copy from this ArrowBuf to ByteBuffer is not supported"); + } + } + } + + /** + * Copy data into this ArrowBuf at a given index onwards from + * a source ByteBuffer. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param src src ByteBuffer where the data will be copied from + */ + public void setBytes(long index, ByteBuffer src) { + // bound check for this ArrowBuf where the data will be copied into + checkIndex(index, src.remaining()); + // length of data to copy + int length = src.remaining(); + // address in this ArrowBuf where the data will be copied to + long dstAddress = addr(index); + if (length != 0) { + if (src.isDirect()) { + // copy src.remaining() bytes of data from src ByteBuffer starting at + // address srcAddress into this ArrowBuf starting at address dstAddress + final long srcAddress = MemoryUtil.getByteBufferAddress(src) + src.position(); + MemoryUtil.UNSAFE.copyMemory(null, srcAddress, null, dstAddress, length); + // after copy, bump the next read position for the src ByteBuffer + src.position(src.position() + length); + } else if (src.hasArray()) { + // copy src.remaining() bytes of data from src ByteBuffer starting at + // index srcIndex into this ArrowBuf starting at address dstAddress + final int srcIndex = src.arrayOffset() + src.position(); + MemoryUtil.UNSAFE.copyMemory( + src.array(), MemoryUtil.BYTE_ARRAY_BASE_OFFSET + srcIndex, null, dstAddress, length); + // after copy, bump the next read position for the src ByteBuffer + src.position(src.position() + length); + } else { + final ByteOrder originalByteOrder = src.order(); + src.order(order()); + try { + // copy word at a time + while (length - 128 >= LONG_SIZE) { + for (int x = 0; x < 16; x++) { + MemoryUtil.UNSAFE.putLong(dstAddress, src.getLong()); + length -= LONG_SIZE; + dstAddress += LONG_SIZE; + } + } + while (length >= LONG_SIZE) { + MemoryUtil.UNSAFE.putLong(dstAddress, src.getLong()); + length -= LONG_SIZE; + dstAddress += LONG_SIZE; + } + // copy last byte + while (length > 0) { + MemoryUtil.UNSAFE.putByte(dstAddress, src.get()); + --length; + ++dstAddress; + } + } finally { + src.order(originalByteOrder); + } + } + } + } + + /** + * Copy data into this ArrowBuf at a given index onwards from + * a source ByteBuffer starting at a given srcIndex for a certain + * length. + * @param index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param src src ByteBuffer where the data will be copied from + * @param srcIndex starting index in the src ByteBuffer where the data copy + * will start from + * @param length length of data to copy from src ByteBuffer + */ + public void setBytes(long index, ByteBuffer src, int srcIndex, int length) { + // bound check for this ArrowBuf where the data will be copied into + checkIndex(index, length); + if (src.isDirect()) { + // copy length bytes of data from src ByteBuffer starting at address + // srcAddress into this ArrowBuf at address dstAddress + final long srcAddress = MemoryUtil.getByteBufferAddress(src) + srcIndex; + final long dstAddress = addr(index); + MemoryUtil.UNSAFE.copyMemory(null, srcAddress, null, dstAddress, length); + } else { + if (srcIndex == 0 && src.capacity() == length) { + // copy the entire ByteBuffer from start to end of length + setBytes(index, src); + } else { + ByteBuffer newBuf = src.duplicate(); + newBuf.position(srcIndex); + newBuf.limit(srcIndex + length); + setBytes(index, newBuf); + } + } + } + + /** + * Copy a given length of data from this ArrowBuf starting at a given index + * into a dst ArrowBuf at dstIndex. + * @param index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param dst dst ArrowBuf where the data will be copied into + * @param dstIndex index (0 based relative to the portion of memory + * dst ArrowBuf has access to) + * @param length length of data to copy + */ + public void getBytes(long index, ArrowBuf dst, long dstIndex, int length) { + // bound check for this ArrowBuf where the data will be copied from + checkIndex(index, length); + // bound check for this ArrowBuf where the data will be copied into + Preconditions.checkArgument(dst != null, "expecting a valid ArrowBuf"); + // bound check for dst ArrowBuf + if (isOutOfBounds(dstIndex, length, dst.capacity())) { + throw new IndexOutOfBoundsException(String.format("index: %d, length: %d (expected: range(0, %d))", + dstIndex, length, dst.capacity())); + } + if (length != 0) { + // copy length bytes of data from this ArrowBuf starting at + // address srcAddress into dst ArrowBuf starting at address + // dstAddress + final long srcAddress = addr(index); + final long dstAddress = dst.memoryAddress() + (long) dstIndex; + MemoryUtil.UNSAFE.copyMemory(null, srcAddress, null, dstAddress, length); + } + } + + /** + * Copy data from src ArrowBuf starting at index srcIndex into this + * ArrowBuf at given index. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param src src ArrowBuf where the data will be copied from + * @param srcIndex starting index in the src ArrowBuf where the copy + * will begin from + * @param length length of data to copy from src ArrowBuf + */ + public void setBytes(long index, ArrowBuf src, long srcIndex, long length) { + // bound check for this ArrowBuf where the data will be copied into + checkIndex(index, length); + // null check + Preconditions.checkArgument(src != null, "expecting a valid ArrowBuf"); + // bound check for src ArrowBuf + if (isOutOfBounds(srcIndex, length, src.capacity())) { + throw new IndexOutOfBoundsException(String.format("index: %d, length: %d (expected: range(0, %d))", + index, length, src.capacity())); + } + if (length != 0) { + // copy length bytes of data from src ArrowBuf starting at + // address srcAddress into this ArrowBuf starting at address + // dstAddress + final long srcAddress = src.memoryAddress() + srcIndex; + final long dstAddress = addr(index); + MemoryUtil.UNSAFE.copyMemory(null, srcAddress, null, dstAddress, length); + } + } + + /** + * Copy readableBytes() number of bytes from src ArrowBuf + * starting from its readerIndex into this ArrowBuf starting + * at the given index. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param src src ArrowBuf where the data will be copied from + */ + public void setBytes(long index, ArrowBuf src) { + // null check + Preconditions.checkArgument(src != null, "expecting valid ArrowBuf"); + final long length = src.readableBytes(); + // bound check for this ArrowBuf where the data will be copied into + checkIndex(index, length); + final long srcAddress = src.memoryAddress() + src.readerIndex; + final long dstAddress = addr(index); + MemoryUtil.UNSAFE.copyMemory(null, srcAddress, null, dstAddress, length); + src.readerIndex(src.readerIndex + length); + } + + /** + * Copy a certain length of bytes from given InputStream + * into this ArrowBuf at the provided index. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param in src stream to copy from + * @param length length of data to copy + * @return number of bytes copied from stream into ArrowBuf + * @throws IOException on failing to read from stream + */ + public int setBytes(long index, InputStream in, int length) throws IOException { + Preconditions.checkArgument(in != null, "expecting valid input stream"); + checkIndex(index, length); + int readBytes = 0; + if (length > 0) { + byte[] tmp = new byte[length]; + // read the data from input stream into tmp byte array + readBytes = in.read(tmp); + if (readBytes > 0) { + // copy readBytes length of data from the tmp byte array starting + // at srcIndex 0 into this ArrowBuf starting at address addr(index) + MemoryUtil.UNSAFE.copyMemory(tmp, MemoryUtil.BYTE_ARRAY_BASE_OFFSET, null, addr(index), readBytes); + } + } + return readBytes; + } + + /** + * Copy a certain length of bytes from this ArrowBuf at a given + * index into the given OutputStream. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param out dst stream to copy data into + * @param length length of data to copy + * @throws IOException on failing to write to stream + */ + public void getBytes(long index, OutputStream out, int length) throws IOException { + Preconditions.checkArgument(out != null, "expecting valid output stream"); + checkIndex(index, length); + if (length > 0) { + // copy length bytes of data from this ArrowBuf starting at + // address addr(index) into the tmp byte array starting at index 0 + byte[] tmp = new byte[length]; + MemoryUtil.UNSAFE.copyMemory(null, addr(index), tmp, MemoryUtil.BYTE_ARRAY_BASE_OFFSET, length); + // write the copied data to output stream + out.write(tmp); + } + } + + @Override + public void close() { + referenceManager.release(); + } + + /** + * Returns the possible memory consumed by this ArrowBuf in the worse case scenario. + * (not shared, connected to larger underlying buffer of allocated memory) + * @return Size in bytes. + */ + public long getPossibleMemoryConsumed() { + return referenceManager.getSize(); + } + + /** + * Return that is Accounted for by this buffer (and its potentially shared siblings within the + * context of the associated allocator). + * @return Size in bytes. + */ + public long getActualMemoryConsumed() { + return referenceManager.getAccountedSize(); + } + + /** + * Return the buffer's byte contents in the form of a hex dump. + * + * @param start the starting byte index + * @param length how many bytes to log + * @return A hex dump in a String. + */ + public String toHexString(final long start, final int length) { + final long roundedStart = (start / LOG_BYTES_PER_ROW) * LOG_BYTES_PER_ROW; + + final StringBuilder sb = new StringBuilder("buffer byte dump\n"); + long index = roundedStart; + for (long nLogged = 0; nLogged < length; nLogged += LOG_BYTES_PER_ROW) { + sb.append(String.format(" [%05d-%05d]", index, index + LOG_BYTES_PER_ROW - 1)); + for (int i = 0; i < LOG_BYTES_PER_ROW; ++i) { + try { + final byte b = getByte(index++); + sb.append(String.format(" 0x%02x", b)); + } catch (IndexOutOfBoundsException ioob) { + sb.append(" <ioob>"); + } + } + sb.append('\n'); + } + return sb.toString(); + } + + /** + * Get the integer id assigned to this ArrowBuf for debugging purposes. + * @return integer id + */ + public long getId() { + return id; + } + + /** + * Prints information of this buffer into <code>sb</code> at the given + * indentation and verbosity level. + * + * <p>It will include history if BaseAllocator.DEBUG is true and + * the verbosity.includeHistoricalLog are true. + * + */ + public void print(StringBuilder sb, int indent, Verbosity verbosity) { + CommonUtil.indent(sb, indent).append(toString()); + + if (BaseAllocator.DEBUG && verbosity.includeHistoricalLog) { + sb.append("\n"); + historicalLog.buildHistory(sb, indent + 1, verbosity.includeStackTraces); + } + } + + /** + * Get the index at which the next byte will be read from. + * @return reader index + */ + public long readerIndex() { + return readerIndex; + } + + /** + * Get the index at which next byte will be written to. + * @return writer index + */ + public long writerIndex() { + return writerIndex; + } + + /** + * Set the reader index for this ArrowBuf. + * @param readerIndex new reader index + * @return this ArrowBuf + */ + public ArrowBuf readerIndex(long readerIndex) { + this.readerIndex = readerIndex; + return this; + } + + /** + * Set the writer index for this ArrowBuf. + * @param writerIndex new writer index + * @return this ArrowBuf + */ + public ArrowBuf writerIndex(long writerIndex) { + this.writerIndex = writerIndex; + return this; + } + + /** + * Zero-out the bytes in this ArrowBuf starting at + * the given index for the given length. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param length length of bytes to zero-out + * @return this ArrowBuf + */ + public ArrowBuf setZero(long index, long length) { + if (length != 0) { + this.checkIndex(index, length); + MemoryUtil.UNSAFE.setMemory(this.addr + index, length, (byte) 0); + } + return this; + } + + /** + * Sets all bits to one in the specified range. + * @param index index index (0 based relative to the portion of memory + * this ArrowBuf has access to) + * @param length length of bytes to set. + * @return this ArrowBuf + */ + public ArrowBuf setOne(int index, int length) { + if (length != 0) { + this.checkIndex(index, length); + MemoryUtil.UNSAFE.setMemory(this.addr + index, length, (byte) 0xff); + } + return this; + } + + /** + * Returns <code>this</code> if size is less then {@link #capacity()}, otherwise + * delegates to {@link BufferManager#replace(ArrowBuf, long)} to get a new buffer. + */ + public ArrowBuf reallocIfNeeded(final long size) { + Preconditions.checkArgument(size >= 0, "reallocation size must be non-negative"); + if (this.capacity() >= size) { + return this; + } + if (bufferManager != null) { + return bufferManager.replace(this, size); + } else { + throw new UnsupportedOperationException( + "Realloc is only available in the context of operator's UDFs"); + } + } + + public ArrowBuf clear() { + this.readerIndex = this.writerIndex = 0; + return this; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java new file mode 100644 index 000000000..8d21cef7a --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BaseAllocator.java @@ -0,0 +1,951 @@ +/* + * 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 java.util.Collection; +import java.util.Collections; +import java.util.HashSet; +import java.util.IdentityHashMap; +import java.util.Map; +import java.util.Set; + +import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; +import org.apache.arrow.memory.rounding.RoundingPolicy; +import org.apache.arrow.memory.util.AssertionUtil; +import org.apache.arrow.memory.util.CommonUtil; +import org.apache.arrow.memory.util.HistoricalLog; +import org.apache.arrow.util.Preconditions; +import org.immutables.value.Value; + +/** + * A base-class that implements all functionality of {@linkplain BufferAllocator}s. + * + * <p>The class is abstract to enforce usage of {@linkplain RootAllocator}/{@linkplain ChildAllocator} + * facades. + */ +abstract class BaseAllocator extends Accountant implements BufferAllocator { + + public static final String DEBUG_ALLOCATOR = "arrow.memory.debug.allocator"; + public static final int DEBUG_LOG_LENGTH = 6; + public static final boolean DEBUG; + private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(BaseAllocator.class); + + // Initialize this before DEFAULT_CONFIG as DEFAULT_CONFIG will eventually initialize the allocation manager, + // which in turn allocates an ArrowBuf, which requires DEBUG to have been properly initialized + static { + // the system property takes precedence. + String propValue = System.getProperty(DEBUG_ALLOCATOR); + if (propValue != null) { + DEBUG = Boolean.parseBoolean(propValue); + } else { + DEBUG = AssertionUtil.isAssertionsEnabled(); + } + logger.info("Debug mode " + (DEBUG ? "enabled." : "disabled.")); + } + + public static final Config DEFAULT_CONFIG = ImmutableConfig.builder().build(); + + // Package exposed for sharing between AllocatorManger and BaseAllocator objects + private final String name; + private final RootAllocator root; + private final Object DEBUG_LOCK = DEBUG ? new Object() : null; + private final AllocationListener listener; + private final BaseAllocator parentAllocator; + private final Map<BaseAllocator, Object> childAllocators; + private final ArrowBuf empty; + // members used purely for debugging + private final IdentityHashMap<BufferLedger, Object> childLedgers; + private final IdentityHashMap<Reservation, Object> reservations; + private final HistoricalLog historicalLog; + private final RoundingPolicy roundingPolicy; + private final AllocationManager.Factory allocationManagerFactory; + + private volatile boolean isClosed = false; // the allocator has been closed + + /** + * Initialize an allocator. + * + * @param parentAllocator parent allocator. null if defining a root allocator + * @param name name of this allocator + * @param config configuration including other options of this allocator + * + * @see Config + */ + protected BaseAllocator( + final BaseAllocator parentAllocator, + final String name, + final Config config) throws OutOfMemoryException { + super(parentAllocator, name, config.getInitReservation(), config.getMaxAllocation()); + + this.listener = config.getListener(); + this.allocationManagerFactory = config.getAllocationManagerFactory(); + + if (parentAllocator != null) { + this.root = parentAllocator.root; + empty = parentAllocator.empty; + } else if (this instanceof RootAllocator) { + this.root = (RootAllocator) this; + empty = createEmpty(); + } else { + throw new IllegalStateException("An parent allocator must either carry a root or be the " + + "root."); + } + + this.parentAllocator = parentAllocator; + this.name = name; + + this.childAllocators = Collections.synchronizedMap(new IdentityHashMap<>()); + + if (DEBUG) { + reservations = new IdentityHashMap<>(); + childLedgers = new IdentityHashMap<>(); + historicalLog = new HistoricalLog(DEBUG_LOG_LENGTH, "allocator[%s]", name); + hist("created by \"%s\", owned = %d", name, this.getAllocatedMemory()); + } else { + reservations = null; + historicalLog = null; + childLedgers = null; + } + this.roundingPolicy = config.getRoundingPolicy(); + } + + @Override + public AllocationListener getListener() { + return listener; + } + + @Override + public BaseAllocator getParentAllocator() { + return parentAllocator; + } + + @Override + public Collection<BufferAllocator> getChildAllocators() { + synchronized (childAllocators) { + return new HashSet<>(childAllocators.keySet()); + } + } + + private static String createErrorMsg(final BufferAllocator allocator, final long rounded, final long requested) { + if (rounded != requested) { + return String.format( + "Unable to allocate buffer of size %d (rounded from %d) due to memory limit. Current " + + "allocation: %d", rounded, requested, allocator.getAllocatedMemory()); + } else { + return String.format( + "Unable to allocate buffer of size %d due to memory limit. Current " + + "allocation: %d", rounded, allocator.getAllocatedMemory()); + } + } + + public static boolean isDebug() { + return DEBUG; + } + + @Override + public void assertOpen() { + if (AssertionUtil.ASSERT_ENABLED) { + if (isClosed) { + throw new IllegalStateException("Attempting operation on allocator when allocator is closed.\n" + + toVerboseString()); + } + } + } + + @Override + public String getName() { + return name; + } + + @Override + public ArrowBuf getEmpty() { + return empty; + } + + /** + * For debug/verification purposes only. Allows an AllocationManager to tell the allocator that + * we have a new ledger + * associated with this allocator. + */ + void associateLedger(BufferLedger ledger) { + assertOpen(); + if (DEBUG) { + synchronized (DEBUG_LOCK) { + childLedgers.put(ledger, null); + } + } + } + + /** + * For debug/verification purposes only. Allows an AllocationManager to tell the allocator that + * we are removing a + * ledger associated with this allocator + */ + void dissociateLedger(BufferLedger ledger) { + assertOpen(); + if (DEBUG) { + synchronized (DEBUG_LOCK) { + if (!childLedgers.containsKey(ledger)) { + throw new IllegalStateException("Trying to remove a child ledger that doesn't exist."); + } + childLedgers.remove(ledger); + } + } + } + + /** + * Track when a ChildAllocator of this BaseAllocator is closed. Used for debugging purposes. + * + * @param childAllocator The child allocator that has been closed. + */ + private void childClosed(final BaseAllocator childAllocator) { + assertOpen(); + + if (DEBUG) { + Preconditions.checkArgument(childAllocator != null, "child allocator can't be null"); + + synchronized (DEBUG_LOCK) { + final Object object = childAllocators.remove(childAllocator); + if (object == null) { + childAllocator.historicalLog.logHistory(logger); + throw new IllegalStateException("Child allocator[" + childAllocator.name + + "] not found in parent allocator[" + name + "]'s childAllocators"); + } + } + } else { + childAllocators.remove(childAllocator); + } + listener.onChildRemoved(this, childAllocator); + } + + @Override + public ArrowBuf buffer(final long initialRequestSize) { + assertOpen(); + + return buffer(initialRequestSize, null); + } + + private ArrowBuf createEmpty() { + return allocationManagerFactory.empty(); + } + + @Override + public ArrowBuf buffer(final long initialRequestSize, BufferManager manager) { + assertOpen(); + + Preconditions.checkArgument(initialRequestSize >= 0, "the requested size must be non-negative"); + + if (initialRequestSize == 0) { + return getEmpty(); + } + + // round the request size according to the rounding policy + final long actualRequestSize = roundingPolicy.getRoundedSize(initialRequestSize); + + listener.onPreAllocation(actualRequestSize); + + AllocationOutcome outcome = this.allocateBytes(actualRequestSize); + if (!outcome.isOk()) { + if (listener.onFailedAllocation(actualRequestSize, outcome)) { + // Second try, in case the listener can do something about it + outcome = this.allocateBytes(actualRequestSize); + } + if (!outcome.isOk()) { + throw new OutOfMemoryException(createErrorMsg(this, actualRequestSize, + initialRequestSize), outcome.getDetails()); + } + } + + boolean success = false; + try { + ArrowBuf buffer = bufferWithoutReservation(actualRequestSize, manager); + success = true; + listener.onAllocation(actualRequestSize); + return buffer; + } catch (OutOfMemoryError e) { + throw e; + } finally { + if (!success) { + releaseBytes(actualRequestSize); + } + } + } + + /** + * Used by usual allocation as well as for allocating a pre-reserved buffer. + * Skips the typical accounting associated with creating a new buffer. + */ + private ArrowBuf bufferWithoutReservation( + final long size, + BufferManager bufferManager) throws OutOfMemoryException { + assertOpen(); + + final AllocationManager manager = newAllocationManager(size); + final BufferLedger ledger = manager.associate(this); // +1 ref cnt (required) + final ArrowBuf buffer = ledger.newArrowBuf(size, bufferManager); + + // make sure that our allocation is equal to what we expected. + Preconditions.checkArgument(buffer.capacity() == size, + "Allocated capacity %d was not equal to requested capacity %d.", buffer.capacity(), size); + + return buffer; + } + + private AllocationManager newAllocationManager(long size) { + return newAllocationManager(this, size); + } + + + private AllocationManager newAllocationManager(BaseAllocator accountingAllocator, long size) { + return allocationManagerFactory.create(accountingAllocator, size); + } + + @Override + public BufferAllocator getRoot() { + return root; + } + + @Override + public BufferAllocator newChildAllocator( + final String name, + final long initReservation, + final long maxAllocation) { + return newChildAllocator(name, this.listener, initReservation, maxAllocation); + } + + @Override + public BufferAllocator newChildAllocator( + final String name, + final AllocationListener listener, + final long initReservation, + final long maxAllocation) { + assertOpen(); + + final ChildAllocator childAllocator = + new ChildAllocator(this, name, configBuilder() + .listener(listener) + .initReservation(initReservation) + .maxAllocation(maxAllocation) + .roundingPolicy(roundingPolicy) + .allocationManagerFactory(allocationManagerFactory) + .build()); + + if (DEBUG) { + synchronized (DEBUG_LOCK) { + childAllocators.put(childAllocator, childAllocator); + historicalLog.recordEvent("allocator[%s] created new child allocator[%s]", name, + childAllocator.getName()); + } + } else { + childAllocators.put(childAllocator, childAllocator); + } + this.listener.onChildAdded(this, childAllocator); + + return childAllocator; + } + + @Override + public AllocationReservation newReservation() { + assertOpen(); + + return new Reservation(); + } + + @Override + public synchronized void close() { + /* + * Some owners may close more than once because of complex cleanup and shutdown + * procedures. + */ + if (isClosed) { + return; + } + + isClosed = true; + + StringBuilder outstandingChildAllocators = new StringBuilder(); + if (DEBUG) { + synchronized (DEBUG_LOCK) { + verifyAllocator(); + + // are there outstanding child allocators? + if (!childAllocators.isEmpty()) { + for (final BaseAllocator childAllocator : childAllocators.keySet()) { + if (childAllocator.isClosed) { + logger.warn(String.format( + "Closed child allocator[%s] on parent allocator[%s]'s child list.\n%s", + childAllocator.name, name, toString())); + } + } + + throw new IllegalStateException( + String.format("Allocator[%s] closed with outstanding child allocators.\n%s", name, + toString())); + } + + // are there outstanding buffers? + final int allocatedCount = childLedgers.size(); + if (allocatedCount > 0) { + throw new IllegalStateException( + String.format("Allocator[%s] closed with outstanding buffers allocated (%d).\n%s", + name, allocatedCount, toString())); + } + + if (reservations.size() != 0) { + throw new IllegalStateException( + String.format("Allocator[%s] closed with outstanding reservations (%d).\n%s", name, + reservations.size(), + toString())); + } + + } + } else { + if (!childAllocators.isEmpty()) { + outstandingChildAllocators.append("Outstanding child allocators : \n"); + synchronized (childAllocators) { + for (final BaseAllocator childAllocator : childAllocators.keySet()) { + outstandingChildAllocators.append(String.format(" %s", childAllocator.toString())); + } + } + } + } + + // Is there unaccounted-for outstanding allocation? + final long allocated = getAllocatedMemory(); + if (allocated > 0) { + if (parent != null && reservation > allocated) { + parent.releaseBytes(reservation - allocated); + } + String msg = String.format("Memory was leaked by query. Memory leaked: (%d)\n%s%s", allocated, + outstandingChildAllocators.toString(), toString()); + logger.error(msg); + throw new IllegalStateException(msg); + } + + // we need to release our memory to our parent before we tell it we've closed. + super.close(); + + // Inform our parent allocator that we've closed + if (parentAllocator != null) { + parentAllocator.childClosed(this); + } + + if (DEBUG) { + historicalLog.recordEvent("closed"); + logger.debug(String.format("closed allocator[%s].", name)); + } + + + } + + @Override + public String toString() { + final Verbosity verbosity = logger.isTraceEnabled() ? Verbosity.LOG_WITH_STACKTRACE + : Verbosity.BASIC; + final StringBuilder sb = new StringBuilder(); + print(sb, 0, verbosity); + return sb.toString(); + } + + /** + * Provide a verbose string of the current allocator state. Includes the state of all child + * allocators, along with + * historical logs of each object and including stacktraces. + * + * @return A Verbose string of current allocator state. + */ + @Override + public String toVerboseString() { + final StringBuilder sb = new StringBuilder(); + print(sb, 0, Verbosity.LOG_WITH_STACKTRACE); + return sb.toString(); + } + + private void hist(String noteFormat, Object... args) { + historicalLog.recordEvent(noteFormat, args); + } + + /** + * Verifies the accounting state of the allocator. Only works for DEBUG. + * + * @throws IllegalStateException when any problems are found + */ + void verifyAllocator() { + final IdentityHashMap<AllocationManager, BaseAllocator> seen = new IdentityHashMap<>(); + verifyAllocator(seen); + } + + /** + * Verifies the accounting state of the allocator (Only works for DEBUG) + * This overload is used for recursive calls, allowing for checking + * that ArrowBufs are unique across all allocators that are checked. + * + * @param buffersSeen a map of buffers that have already been seen when walking a tree of + * allocators + * @throws IllegalStateException when any problems are found + */ + private void verifyAllocator( + final IdentityHashMap<AllocationManager, BaseAllocator> buffersSeen) { + // The remaining tests can only be performed if we're in debug mode. + if (!DEBUG) { + return; + } + + synchronized (DEBUG_LOCK) { + final long allocated = getAllocatedMemory(); + + // verify my direct descendants + final Set<BaseAllocator> childSet = childAllocators.keySet(); + for (final BaseAllocator childAllocator : childSet) { + childAllocator.verifyAllocator(buffersSeen); + } + + /* + * Verify my relationships with my descendants. + * + * The sum of direct child allocators' owned memory must be <= my allocated memory; my + * allocated memory also + * includes ArrowBuf's directly allocated by me. + */ + long childTotal = 0; + for (final BaseAllocator childAllocator : childSet) { + childTotal += Math.max(childAllocator.getAllocatedMemory(), childAllocator.reservation); + } + if (childTotal > getAllocatedMemory()) { + historicalLog.logHistory(logger); + logger.debug("allocator[" + name + "] child event logs BEGIN"); + for (final BaseAllocator childAllocator : childSet) { + childAllocator.historicalLog.logHistory(logger); + } + logger.debug("allocator[" + name + "] child event logs END"); + throw new IllegalStateException( + "Child allocators own more memory (" + childTotal + ") than their parent (name = " + + name + " ) has allocated (" + getAllocatedMemory() + ')'); + } + + // Furthermore, the amount I've allocated should be that plus buffers I've allocated. + long bufferTotal = 0; + + final Set<BufferLedger> ledgerSet = childLedgers.keySet(); + for (final BufferLedger ledger : ledgerSet) { + if (!ledger.isOwningLedger()) { + continue; + } + + final AllocationManager am = ledger.getAllocationManager(); + /* + * Even when shared, ArrowBufs are rewrapped, so we should never see the same instance + * twice. + */ + final BaseAllocator otherOwner = buffersSeen.get(am); + if (otherOwner != null) { + throw new IllegalStateException("This allocator's ArrowBuf already owned by another " + + "allocator"); + } + buffersSeen.put(am, this); + + bufferTotal += am.getSize(); + } + + // Preallocated space has to be accounted for + final Set<Reservation> reservationSet = reservations.keySet(); + long reservedTotal = 0; + for (final Reservation reservation : reservationSet) { + if (!reservation.isUsed()) { + reservedTotal += reservation.getSize(); + } + } + + if (bufferTotal + reservedTotal + childTotal != getAllocatedMemory()) { + final StringBuilder sb = new StringBuilder(); + sb.append("allocator["); + sb.append(name); + sb.append("]\nallocated: "); + sb.append(Long.toString(allocated)); + sb.append(" allocated - (bufferTotal + reservedTotal + childTotal): "); + sb.append(Long.toString(allocated - (bufferTotal + reservedTotal + childTotal))); + sb.append('\n'); + + if (bufferTotal != 0) { + sb.append("buffer total: "); + sb.append(Long.toString(bufferTotal)); + sb.append('\n'); + dumpBuffers(sb, ledgerSet); + } + + if (childTotal != 0) { + sb.append("child total: "); + sb.append(Long.toString(childTotal)); + sb.append('\n'); + + for (final BaseAllocator childAllocator : childSet) { + sb.append("child allocator["); + sb.append(childAllocator.name); + sb.append("] owned "); + sb.append(Long.toString(childAllocator.getAllocatedMemory())); + sb.append('\n'); + } + } + + if (reservedTotal != 0) { + sb.append(String.format("reserved total : %d bytes.", reservedTotal)); + for (final Reservation reservation : reservationSet) { + reservation.historicalLog.buildHistory(sb, 0, true); + sb.append('\n'); + } + } + + logger.debug(sb.toString()); + + final long allocated2 = getAllocatedMemory(); + + if (allocated2 != allocated) { + throw new IllegalStateException(String.format( + "allocator[%s]: allocated t1 (%d) + allocated t2 (%d). Someone released memory while in verification.", + name, allocated, allocated2)); + + } + throw new IllegalStateException(String.format( + "allocator[%s]: buffer space (%d) + prealloc space (%d) + child space (%d) != allocated (%d)", + name, bufferTotal, reservedTotal, childTotal, allocated)); + } + } + } + + void print(StringBuilder sb, int level, Verbosity verbosity) { + + CommonUtil.indent(sb, level) + .append("Allocator(") + .append(name) + .append(") ") + .append(reservation) + .append('/') + .append(getAllocatedMemory()) + .append('/') + .append(getPeakMemoryAllocation()) + .append('/') + .append(getLimit()) + .append(" (res/actual/peak/limit)") + .append('\n'); + + if (DEBUG) { + CommonUtil.indent(sb, level + 1).append(String.format("child allocators: %d\n", childAllocators.size())); + for (BaseAllocator child : childAllocators.keySet()) { + child.print(sb, level + 2, verbosity); + } + + CommonUtil.indent(sb, level + 1).append(String.format("ledgers: %d\n", childLedgers.size())); + for (BufferLedger ledger : childLedgers.keySet()) { + ledger.print(sb, level + 2, verbosity); + } + + final Set<Reservation> reservations = this.reservations.keySet(); + CommonUtil.indent(sb, level + 1).append(String.format("reservations: %d\n", reservations.size())); + for (final Reservation reservation : reservations) { + if (verbosity.includeHistoricalLog) { + reservation.historicalLog.buildHistory(sb, level + 3, true); + } + } + + } + + } + + private void dumpBuffers(final StringBuilder sb, final Set<BufferLedger> ledgerSet) { + for (final BufferLedger ledger : ledgerSet) { + if (!ledger.isOwningLedger()) { + continue; + } + final AllocationManager am = ledger.getAllocationManager(); + sb.append("UnsafeDirectLittleEndian[identityHashCode == "); + sb.append(Integer.toString(System.identityHashCode(am))); + sb.append("] size "); + sb.append(Long.toString(am.getSize())); + sb.append('\n'); + } + } + + /** + * Enum for logging verbosity. + */ + public enum Verbosity { + BASIC(false, false), // only include basic information + LOG(true, false), // include basic + LOG_WITH_STACKTRACE(true, true) // + ; + + public final boolean includeHistoricalLog; + public final boolean includeStackTraces; + + Verbosity(boolean includeHistoricalLog, boolean includeStackTraces) { + this.includeHistoricalLog = includeHistoricalLog; + this.includeStackTraces = includeStackTraces; + } + } + + /** + * Returns a default {@link Config} instance. + * + * @see ImmutableConfig.Builder + */ + public static Config defaultConfig() { + return DEFAULT_CONFIG; + + } + + /** + * Returns a builder class for configuring BaseAllocator's options. + */ + public static ImmutableConfig.Builder configBuilder() { + return ImmutableConfig.builder(); + } + + @Override + public RoundingPolicy getRoundingPolicy() { + return roundingPolicy; + } + + /** + * Config class of {@link BaseAllocator}. + */ + @Value.Immutable + abstract static class Config { + /** + * Factory for creating {@link AllocationManager} instances. + */ + @Value.Default + AllocationManager.Factory getAllocationManagerFactory() { + return DefaultAllocationManagerOption.getDefaultAllocationManagerFactory(); + } + + /** + * Listener callback. Must be non-null. + */ + @Value.Default + AllocationListener getListener() { + return AllocationListener.NOOP; + } + + /** + * Initial reservation size (in bytes) for this allocator. + */ + @Value.Default + long getInitReservation() { + return 0; + } + + /** + * Max allocation size (in bytes) for this allocator, allocations past this limit fail. + * Can be modified after construction. + */ + @Value.Default + long getMaxAllocation() { + return Long.MAX_VALUE; + } + + /** + * The policy for rounding the buffer size. + */ + @Value.Default + RoundingPolicy getRoundingPolicy() { + return DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY; + } + } + + /** + * Implementation of {@link AllocationReservation} that supports + * history tracking under {@linkplain #DEBUG} is true. + */ + public class Reservation implements AllocationReservation { + + private final HistoricalLog historicalLog; + private int nBytes = 0; + private boolean used = false; + private boolean closed = false; + + /** + * Creates a new reservation. + * + * <p>If {@linkplain #DEBUG} is true this will capture a historical + * log of events relevant to this Reservation. + */ + public Reservation() { + if (DEBUG) { + historicalLog = new HistoricalLog("Reservation[allocator[%s], %d]", name, System + .identityHashCode(this)); + historicalLog.recordEvent("created"); + synchronized (DEBUG_LOCK) { + reservations.put(this, this); + } + } else { + historicalLog = null; + } + } + + @Override + public boolean add(final int nBytes) { + assertOpen(); + + Preconditions.checkArgument(nBytes >= 0, "nBytes(%d) < 0", nBytes); + Preconditions.checkState(!closed, "Attempt to increase reservation after reservation has been closed"); + Preconditions.checkState(!used, "Attempt to increase reservation after reservation has been used"); + + // we round up to next power of two since all reservations are done in powers of two. This + // may overestimate the + // preallocation since someone may perceive additions to be power of two. If this becomes a + // problem, we can look + // at + // modifying this behavior so that we maintain what we reserve and what the user asked for + // and make sure to only + // round to power of two as necessary. + final int nBytesTwo = CommonUtil.nextPowerOfTwo(nBytes); + if (!reserve(nBytesTwo)) { + return false; + } + + this.nBytes += nBytesTwo; + return true; + } + + @Override + public ArrowBuf allocateBuffer() { + assertOpen(); + + Preconditions.checkState(!closed, "Attempt to allocate after closed"); + Preconditions.checkState(!used, "Attempt to allocate more than once"); + + final ArrowBuf arrowBuf = allocate(nBytes); + used = true; + return arrowBuf; + } + + @Override + public int getSize() { + return nBytes; + } + + @Override + public boolean isUsed() { + return used; + } + + @Override + public boolean isClosed() { + return closed; + } + + @Override + public void close() { + assertOpen(); + + if (closed) { + return; + } + + if (DEBUG) { + if (!isClosed()) { + final Object object; + synchronized (DEBUG_LOCK) { + object = reservations.remove(this); + } + if (object == null) { + final StringBuilder sb = new StringBuilder(); + print(sb, 0, Verbosity.LOG_WITH_STACKTRACE); + logger.debug(sb.toString()); + throw new IllegalStateException(String.format("Didn't find closing reservation[%d]", + System.identityHashCode(this))); + } + + historicalLog.recordEvent("closed"); + } + } + + if (!used) { + releaseReservation(nBytes); + } + + closed = true; + } + + @Override + public boolean reserve(int nBytes) { + assertOpen(); + + final AllocationOutcome outcome = BaseAllocator.this.allocateBytes(nBytes); + + if (DEBUG) { + historicalLog.recordEvent("reserve(%d) => %s", nBytes, Boolean.toString(outcome.isOk())); + } + + return outcome.isOk(); + } + + /** + * Allocate a buffer of the requested size. + * + * <p>The implementation of the allocator's inner class provides this. + * + * @param nBytes the size of the buffer requested + * @return the buffer, or null, if the request cannot be satisfied + */ + private ArrowBuf allocate(int nBytes) { + assertOpen(); + + boolean success = false; + + /* + * The reservation already added the requested bytes to the allocators owned and allocated + * bytes via reserve(). + * This ensures that they can't go away. But when we ask for the buffer here, that will add + * to the allocated bytes + * as well, so we need to return the same number back to avoid double-counting them. + */ + try { + final ArrowBuf arrowBuf = BaseAllocator.this.bufferWithoutReservation(nBytes, null); + + listener.onAllocation(nBytes); + if (DEBUG) { + historicalLog.recordEvent("allocate() => %s", String.format("ArrowBuf[%d]", arrowBuf + .getId())); + } + success = true; + return arrowBuf; + } finally { + if (!success) { + releaseBytes(nBytes); + } + } + } + + /** + * Return the reservation back to the allocator without having used it. + * + * @param nBytes the size of the reservation + */ + private void releaseReservation(int nBytes) { + assertOpen(); + + releaseBytes(nBytes); + + if (DEBUG) { + historicalLog.recordEvent("releaseReservation(%d)", nBytes); + } + } + + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java new file mode 100644 index 000000000..bbf7ff34d --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BoundsChecking.java @@ -0,0 +1,63 @@ +/* + * 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; + +/** + * Configuration class to determine if bounds checking should be turned on or off. + * + * <p> + * Bounds checking is on by default. You can disable it by setting either the system property or + * the environmental variable to "true". The system property can be "arrow.enable_unsafe_memory_access" + * or "drill.enable_unsafe_memory_access". The latter is deprecated. The environmental variable is named + * "ARROW_ENABLE_UNSAFE_MEMORY_ACCESS". + * When both the system property and the environmental variable are set, the system property takes precedence. + * </p> + */ +public class BoundsChecking { + + public static final boolean BOUNDS_CHECKING_ENABLED; + static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(BoundsChecking.class); + + static { + String envProperty = System.getenv("ARROW_ENABLE_UNSAFE_MEMORY_ACCESS"); + String oldProperty = System.getProperty("drill.enable_unsafe_memory_access"); + if (oldProperty != null) { + logger.warn("\"drill.enable_unsafe_memory_access\" has been renamed to \"arrow.enable_unsafe_memory_access\""); + logger.warn("\"arrow.enable_unsafe_memory_access\" can be set to: " + + " true (to not check) or false (to check, default)"); + } + String newProperty = System.getProperty("arrow.enable_unsafe_memory_access"); + + // The priority of determining the unsafe flag: + // 1. The system properties take precedence over the environmental variable. + // 2. The new system property takes precedence over the new system property. + String unsafeFlagValue = newProperty; + if (unsafeFlagValue == null) { + unsafeFlagValue = oldProperty; + } + if (unsafeFlagValue == null) { + unsafeFlagValue = envProperty; + } + + BOUNDS_CHECKING_ENABLED = !"true".equals(unsafeFlagValue); + } + + private BoundsChecking() { + } + +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java new file mode 100644 index 000000000..e59349c64 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferAllocator.java @@ -0,0 +1,238 @@ +/* + * 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 java.util.Collection; + +import org.apache.arrow.memory.rounding.DefaultRoundingPolicy; +import org.apache.arrow.memory.rounding.RoundingPolicy; + +/** + * Wrapper class to deal with byte buffer allocation. Ensures users only use designated methods. + */ +public interface BufferAllocator extends AutoCloseable { + + /** + * Allocate a new or reused buffer of the provided size. Note that the buffer may technically + * be larger than the + * requested size for rounding purposes. However, the buffer's capacity will be set to the + * configured size. + * + * @param size The size in bytes. + * @return a new ArrowBuf, or null if the request can't be satisfied + * @throws OutOfMemoryException if buffer cannot be allocated + */ + ArrowBuf buffer(long size); + + /** + * Allocate a new or reused buffer of the provided size. Note that the buffer may technically + * be larger than the + * requested size for rounding purposes. However, the buffer's capacity will be set to the + * configured size. + * + * @param size The size in bytes. + * @param manager A buffer manager to manage reallocation. + * @return a new ArrowBuf, or null if the request can't be satisfied + * @throws OutOfMemoryException if buffer cannot be allocated + */ + ArrowBuf buffer(long size, BufferManager manager); + + /** + * Get the root allocator of this allocator. If this allocator is already a root, return + * this directly. + * + * @return The root allocator + */ + BufferAllocator getRoot(); + + /** + * Create a new child allocator. + * + * @param name the name of the allocator. + * @param initReservation the initial space reservation (obtained from this allocator) + * @param maxAllocation maximum amount of space the new allocator can allocate + * @return the new allocator, or null if it can't be created + */ + BufferAllocator newChildAllocator(String name, long initReservation, long maxAllocation); + + /** + * Create a new child allocator. + * + * @param name the name of the allocator. + * @param listener allocation listener for the newly created child + * @param initReservation the initial space reservation (obtained from this allocator) + * @param maxAllocation maximum amount of space the new allocator can allocate + * @return the new allocator, or null if it can't be created + */ + BufferAllocator newChildAllocator( + String name, + AllocationListener listener, + long initReservation, + long maxAllocation); + + /** + * Close and release all buffers generated from this buffer pool. + * + * <p>When assertions are on, complains if there are any outstanding buffers; to avoid + * that, release all buffers before the allocator is closed.</p> + */ + @Override + void close(); + + /** + * Returns the amount of memory currently allocated from this allocator. + * + * @return the amount of memory currently allocated + */ + long getAllocatedMemory(); + + /** + * Return the current maximum limit this allocator imposes. + * + * @return Limit in number of bytes. + */ + long getLimit(); + + /** + * Return the initial reservation. + * + * @return reservation in bytes. + */ + long getInitReservation(); + + /** + * Set the maximum amount of memory this allocator is allowed to allocate. + * + * @param newLimit The new Limit to apply to allocations + */ + void setLimit(long newLimit); + + /** + * Returns the peak amount of memory allocated from this allocator. + * + * @return the peak amount of memory allocated + */ + long getPeakMemoryAllocation(); + + /** + * Returns the amount of memory that can probably be allocated at this moment + * without exceeding this or any parents allocation maximum. + * + * @return Headroom in bytes + */ + long getHeadroom(); + + /** + * Forcibly allocate bytes. Returns whether the allocation fit within limits. + * + * @param size to increase + * @return Whether the allocation fit within limits. + */ + boolean forceAllocate(long size); + + + /** + * Release bytes from this allocator. + * + * @param size to release + */ + void releaseBytes(long size); + + /** + * Returns the allocation listener used by this allocator. + * + * @return the {@link AllocationListener} instance. Or {@link AllocationListener#NOOP} by default if no listener + * is configured when this allocator was created. + */ + AllocationListener getListener(); + + /** + * Returns the parent allocator. + * + * @return parent allocator + */ + BufferAllocator getParentAllocator(); + + /** + * Returns the set of child allocators. + * + * @return set of child allocators + */ + Collection<BufferAllocator> getChildAllocators(); + + /** + * Create an allocation reservation. A reservation is a way of building up + * a request for a buffer whose size is not known in advance. See + * + * @return the newly created reservation + * @see AllocationReservation + */ + AllocationReservation newReservation(); + + /** + * Get a reference to the empty buffer associated with this allocator. Empty buffers are + * special because we don't + * worry about them leaking or managing reference counts on them since they don't actually + * point to any memory. + * + * @return the empty buffer + */ + ArrowBuf getEmpty(); + + /** + * Return the name of this allocator. This is a human readable name that can help debugging. + * Typically provides + * coordinates about where this allocator was created + * + * @return the name of the allocator + */ + String getName(); + + /** + * Return whether or not this allocator (or one if its parents) is over its limits. In the case + * that an allocator is + * over its limit, all consumers of that allocator should aggressively try to address the + * overlimit situation. + * + * @return whether or not this allocator (or one if its parents) is over its limits + */ + boolean isOverLimit(); + + /** + * Return a verbose string describing this allocator. If in DEBUG mode, this will also include + * relevant stacktraces + * and historical logs for underlying objects + * + * @return A very verbose description of the allocator hierarchy. + */ + String toVerboseString(); + + /** + * Asserts (using java assertions) that the provided allocator is currently open. If assertions + * are disabled, this is + * a no-op. + */ + void assertOpen(); + + /** + * Gets the rounding policy of the allocator. + */ + default RoundingPolicy getRoundingPolicy() { + return DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java new file mode 100644 index 000000000..48b3e183d --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferLedger.java @@ -0,0 +1,525 @@ +/* + * 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 java.util.IdentityHashMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.memory.util.CommonUtil; +import org.apache.arrow.memory.util.HistoricalLog; +import org.apache.arrow.util.Preconditions; + +/** + * The reference manager that binds an {@link AllocationManager} to + * {@link BufferAllocator} and a set of {@link ArrowBuf}. The set of + * ArrowBufs managed by this reference manager share a common + * fate (same reference count). + */ +public class BufferLedger implements ValueWithKeyIncluded<BufferAllocator>, ReferenceManager { + private final IdentityHashMap<ArrowBuf, Object> buffers = + BaseAllocator.DEBUG ? new IdentityHashMap<>() : null; + private static final AtomicLong LEDGER_ID_GENERATOR = new AtomicLong(0); + // unique ID assigned to each ledger + private final long ledgerId = LEDGER_ID_GENERATOR.incrementAndGet(); + private final AtomicInteger bufRefCnt = new AtomicInteger(0); // start at zero so we can + // manage request for retain + // correctly + private final long lCreationTime = System.nanoTime(); + private final BufferAllocator allocator; + private final AllocationManager allocationManager; + private final HistoricalLog historicalLog = + BaseAllocator.DEBUG ? new HistoricalLog(BaseAllocator.DEBUG_LOG_LENGTH, + "BufferLedger[%d]", 1) : null; + private volatile long lDestructionTime = 0; + + BufferLedger(final BufferAllocator allocator, final AllocationManager allocationManager) { + this.allocator = allocator; + this.allocationManager = allocationManager; + } + + boolean isOwningLedger() { + return this == allocationManager.getOwningLedger(); + } + + public BufferAllocator getKey() { + return allocator; + } + + /** + * Get the buffer allocator associated with this reference manager. + * @return buffer allocator + */ + @Override + public BufferAllocator getAllocator() { + return allocator; + } + + /** + * Get this ledger's reference count. + * @return reference count + */ + @Override + public int getRefCount() { + return bufRefCnt.get(); + } + + /** + * Increment the ledger's reference count for the associated + * underlying memory chunk. All ArrowBufs managed by this ledger + * will share the ref count. + */ + void increment() { + bufRefCnt.incrementAndGet(); + } + + /** + * Decrement the ledger's reference count by 1 for the associated underlying + * memory chunk. If the reference count drops to 0, it implies that + * no ArrowBufs managed by this reference manager need access to the memory + * chunk. In that case, the ledger should inform the allocation manager + * about releasing its ownership for the chunk. Whether or not the memory + * chunk will be released is something that {@link AllocationManager} will + * decide since tracks the usage of memory chunk across multiple reference + * managers and allocators. + * @return true if the new ref count has dropped to 0, false otherwise + */ + @Override + public boolean release() { + return release(1); + } + + /** + * Decrement the ledger's reference count for the associated underlying + * memory chunk. If the reference count drops to 0, it implies that + * no ArrowBufs managed by this reference manager need access to the memory + * chunk. In that case, the ledger should inform the allocation manager + * about releasing its ownership for the chunk. Whether or not the memory + * chunk will be released is something that {@link AllocationManager} will + * decide since tracks the usage of memory chunk across multiple reference + * managers and allocators. + * @param decrement amount to decrease the reference count by + * @return true if the new ref count has dropped to 0, false otherwise + */ + @Override + public boolean release(int decrement) { + Preconditions.checkState(decrement >= 1, + "ref count decrement should be greater than or equal to 1"); + // decrement the ref count + final int refCnt = decrement(decrement); + if (BaseAllocator.DEBUG) { + historicalLog.recordEvent("release(%d). original value: %d", + decrement, refCnt + decrement); + } + // the new ref count should be >= 0 + Preconditions.checkState(refCnt >= 0, "RefCnt has gone negative"); + return refCnt == 0; + } + + /** + * Decrement the ledger's reference count for the associated underlying + * memory chunk. If the reference count drops to 0, it implies that + * no ArrowBufs managed by this reference manager need access to the memory + * chunk. In that case, the ledger should inform the allocation manager + * about releasing its ownership for the chunk. Whether or not the memory + * chunk will be released is something that {@link AllocationManager} will + * decide since tracks the usage of memory chunk across multiple reference + * managers and allocators. + * + * @param decrement amount to decrease the reference count by + * @return the new reference count + */ + private int decrement(int decrement) { + allocator.assertOpen(); + final int outcome; + synchronized (allocationManager) { + outcome = bufRefCnt.addAndGet(-decrement); + if (outcome == 0) { + lDestructionTime = System.nanoTime(); + // refcount of this reference manager has dropped to 0 + // inform the allocation manager that this reference manager + // no longer holds references to underlying memory + allocationManager.release(this); + } + } + return outcome; + } + + /** + * Increment the ledger's reference count for associated + * underlying memory chunk by 1. + */ + @Override + public void retain() { + retain(1); + } + + /** + * Increment the ledger's reference count for associated + * underlying memory chunk by the given amount. + * + * @param increment amount to increase the reference count by + */ + @Override + public void retain(int increment) { + Preconditions.checkArgument(increment > 0, "retain(%s) argument is not positive", increment); + if (BaseAllocator.DEBUG) { + historicalLog.recordEvent("retain(%d)", increment); + } + final int originalReferenceCount = bufRefCnt.getAndAdd(increment); + Preconditions.checkArgument(originalReferenceCount > 0); + } + + /** + * Derive a new ArrowBuf from a given source ArrowBuf. The new derived + * ArrowBuf will share the same reference count as rest of the ArrowBufs + * associated with this ledger. This operation is typically used for + * slicing -- creating new ArrowBufs from a compound ArrowBuf starting at + * a particular index in the underlying memory and having access to a + * particular length (in bytes) of data in memory chunk. + * <p> + * This method is also used as a helper for transferring ownership and retain to target + * allocator. + * </p> + * @param sourceBuffer source ArrowBuf + * @param index index (relative to source ArrowBuf) new ArrowBuf should be + * derived from + * @param length length (bytes) of data in underlying memory that derived buffer will + * have access to in underlying memory + * @return derived buffer + */ + @Override + public ArrowBuf deriveBuffer(final ArrowBuf sourceBuffer, long index, long length) { + /* + * Usage type 1 for deriveBuffer(): + * Used for slicing where index represents a relative index in the source ArrowBuf + * as the slice start point. This is why we need to add the source buffer offset + * to compute the start virtual address of derived buffer within the + * underlying chunk. + * + * Usage type 2 for deriveBuffer(): + * Used for retain(target allocator) and transferOwnership(target allocator) + * where index is 0 since these operations simply create a new ArrowBuf associated + * with another combination of allocator buffer ledger for the same underlying memory + */ + + // the memory address stored inside ArrowBuf is its starting virtual + // address in the underlying memory chunk from the point it has + // access. so it is already accounting for the offset of the source buffer + // we can simply add the index to get the starting address of new buffer. + final long derivedBufferAddress = sourceBuffer.memoryAddress() + index; + + // create new ArrowBuf + final ArrowBuf derivedBuf = new ArrowBuf( + this, + null, + length, // length (in bytes) in the underlying memory chunk for this new ArrowBuf + derivedBufferAddress // starting byte address in the underlying memory for this new ArrowBuf + ); + + // logging + if (BaseAllocator.DEBUG) { + historicalLog.recordEvent( + "ArrowBuf(BufferLedger, BufferAllocator[%s], " + + "UnsafeDirectLittleEndian[identityHashCode == " + + "%d](%s)) => ledger hc == %d", + allocator.getName(), System.identityHashCode(derivedBuf), derivedBuf.toString(), + System.identityHashCode(this)); + + synchronized (buffers) { + buffers.put(derivedBuf, null); + } + } + + return derivedBuf; + } + + /** + * Used by an allocator to create a new ArrowBuf. This is provided + * as a helper method for the allocator when it allocates a new memory chunk + * using a new instance of allocation manager and creates a new reference manager + * too. + * + * @param length The length in bytes that this ArrowBuf will provide access to. + * @param manager An optional BufferManager argument that can be used to manage expansion of + * this ArrowBuf + * @return A new ArrowBuf that shares references with all ArrowBufs associated + * with this BufferLedger + */ + ArrowBuf newArrowBuf(final long length, final BufferManager manager) { + allocator.assertOpen(); + + // the start virtual address of the ArrowBuf will be same as address of memory chunk + final long startAddress = allocationManager.memoryAddress(); + + // create ArrowBuf + final ArrowBuf buf = new ArrowBuf(this, manager, length, startAddress); + + // logging + if (BaseAllocator.DEBUG) { + historicalLog.recordEvent( + "ArrowBuf(BufferLedger, BufferAllocator[%s], " + + "UnsafeDirectLittleEndian[identityHashCode == " + "%d](%s)) => ledger hc == %d", + allocator.getName(), System.identityHashCode(buf), buf.toString(), + System.identityHashCode(this)); + + synchronized (buffers) { + buffers.put(buf, null); + } + } + + return buf; + } + + /** + * Create a new ArrowBuf that is associated with an alternative allocator for the purposes of + * memory ownership and accounting. This has no impact on the reference counting for the current + * ArrowBuf except in the situation where the passed in Allocator is the same as the current buffer. + * <p> + * This operation has no impact on the reference count of this ArrowBuf. The newly created + * ArrowBuf with either have a reference count of 1 (in the case that this is the first time this + * memory is being associated with the target allocator or in other words allocation manager currently + * doesn't hold a mapping for the target allocator) or the current value of the reference count for + * the target allocator-reference manager combination + 1 in the case that the provided allocator + * already had an association to this underlying memory. + * </p> + * + * @param srcBuffer source ArrowBuf + * @param target The target allocator to create an association with. + * @return A new ArrowBuf which shares the same underlying memory as the provided ArrowBuf. + */ + @Override + public ArrowBuf retain(final ArrowBuf srcBuffer, BufferAllocator target) { + + if (BaseAllocator.DEBUG) { + historicalLog.recordEvent("retain(%s)", target.getName()); + } + + // the call to associate will return the corresponding reference manager (buffer ledger) for + // the target allocator. if the allocation manager didn't already have a mapping + // for the target allocator, it will create one and return the new reference manager with a + // reference count of 1. Thus the newly created buffer in this case will have a ref count of 1. + // alternatively, if there was already a mapping for <buffer allocator, ref manager> in + // allocation manager, the ref count of the new buffer will be targetrefmanager.refcount() + 1 + // and this will be true for all the existing buffers currently managed by targetrefmanager + final BufferLedger targetRefManager = allocationManager.associate(target); + // create a new ArrowBuf to associate with new allocator and target ref manager + final long targetBufLength = srcBuffer.capacity(); + ArrowBuf targetArrowBuf = targetRefManager.deriveBuffer(srcBuffer, 0, targetBufLength); + targetArrowBuf.readerIndex(srcBuffer.readerIndex()); + targetArrowBuf.writerIndex(srcBuffer.writerIndex()); + return targetArrowBuf; + } + + /** + * Transfer any balance the current ledger has to the target ledger. In the case + * that the current ledger holds no memory, no transfer is made to the new ledger. + * + * @param targetReferenceManager The ledger to transfer ownership account to. + * @return Whether transfer fit within target ledgers limits. + */ + boolean transferBalance(final ReferenceManager targetReferenceManager) { + Preconditions.checkArgument(targetReferenceManager != null, + "Expecting valid target reference manager"); + final BufferAllocator targetAllocator = targetReferenceManager.getAllocator(); + Preconditions.checkArgument(allocator.getRoot() == targetAllocator.getRoot(), + "You can only transfer between two allocators that share the same root."); + + allocator.assertOpen(); + targetReferenceManager.getAllocator().assertOpen(); + + // if we're transferring to ourself, just return. + if (targetReferenceManager == this) { + return true; + } + + // since two balance transfers out from the allocation manager could cause incorrect + // accounting, we need to ensure + // that this won't happen by synchronizing on the allocation manager instance. + synchronized (allocationManager) { + if (allocationManager.getOwningLedger() != this) { + // since the calling reference manager is not the owning + // reference manager for the underlying memory, transfer is + // a NO-OP + return true; + } + + if (BaseAllocator.DEBUG) { + this.historicalLog.recordEvent("transferBalance(%s)", + targetReferenceManager.getAllocator().getName()); + } + + boolean overlimit = targetAllocator.forceAllocate(allocationManager.getSize()); + allocator.releaseBytes(allocationManager.getSize()); + // since the transfer can only happen from the owning reference manager, + // we need to set the target ref manager as the new owning ref manager + // for the chunk of memory in allocation manager + allocationManager.setOwningLedger((BufferLedger) targetReferenceManager); + return overlimit; + } + } + + /** + * Transfer the memory accounting ownership of this ArrowBuf to another allocator. + * This will generate a new ArrowBuf that carries an association with the underlying memory + * of this ArrowBuf. If this ArrowBuf is connected to the owning BufferLedger of this memory, + * that memory ownership/accounting will be transferred to the target allocator. If this + * ArrowBuf does not currently own the memory underlying it (and is only associated with it), + * this does not transfer any ownership to the newly created ArrowBuf. + * <p> + * This operation has no impact on the reference count of this ArrowBuf. The newly created + * ArrowBuf with either have a reference count of 1 (in the case that this is the first time + * this memory is being associated with the new allocator) or the current value of the reference + * count for the other AllocationManager/BufferLedger combination + 1 in the case that the provided + * allocator already had an association to this underlying memory. + * </p> + * <p> + * Transfers will always succeed, even if that puts the other allocator into an overlimit + * situation. This is possible due to the fact that the original owning allocator may have + * allocated this memory out of a local reservation whereas the target allocator may need to + * allocate new memory from a parent or RootAllocator. This operation is done n a mostly-lockless + * but consistent manner. As such, the overlimit==true situation could occur slightly prematurely + * to an actual overlimit==true condition. This is simply conservative behavior which means we may + * return overlimit slightly sooner than is necessary. + * </p> + * + * @param target The allocator to transfer ownership to. + * @return A new transfer result with the impact of the transfer (whether it was overlimit) as + * well as the newly created ArrowBuf. + */ + @Override + public TransferResult transferOwnership(final ArrowBuf srcBuffer, final BufferAllocator target) { + // the call to associate will return the corresponding reference manager (buffer ledger) for + // the target allocator. if the allocation manager didn't already have a mapping + // for the target allocator, it will create one and return the new reference manager with a + // reference count of 1. Thus the newly created buffer in this case will have a ref count of 1. + // alternatively, if there was already a mapping for <buffer allocator, ref manager> in + // allocation manager, the ref count of the new buffer will be targetrefmanager.refcount() + 1 + // and this will be true for all the existing buffers currently managed by targetrefmanager + final BufferLedger targetRefManager = allocationManager.associate(target); + // create a new ArrowBuf to associate with new allocator and target ref manager + final long targetBufLength = srcBuffer.capacity(); + final ArrowBuf targetArrowBuf = targetRefManager.deriveBuffer(srcBuffer, 0, targetBufLength); + targetArrowBuf.readerIndex(srcBuffer.readerIndex()); + targetArrowBuf.writerIndex(srcBuffer.writerIndex()); + final boolean allocationFit = transferBalance(targetRefManager); + return new TransferResult(allocationFit, targetArrowBuf); + } + + /** + * The outcome of a Transfer. + */ + public class TransferResult implements OwnershipTransferResult { + + // Whether this transfer fit within the target allocator's capacity. + final boolean allocationFit; + + // The newly created buffer associated with the target allocator + public final ArrowBuf buffer; + + private TransferResult(boolean allocationFit, ArrowBuf buffer) { + this.allocationFit = allocationFit; + this.buffer = buffer; + } + + @Override + public ArrowBuf getTransferredBuffer() { + return buffer; + } + + @Override + public boolean getAllocationFit() { + return allocationFit; + } + } + + /** + * Total size (in bytes) of memory underlying this reference manager. + * @return Size (in bytes) of the memory chunk + */ + @Override + public long getSize() { + return allocationManager.getSize(); + } + + /** + * How much memory is accounted for by this ledger. This is either getSize() + * if this is the owning ledger for the memory or zero in the case that this + * is not the owning ledger associated with this memory. + * @return Amount of accounted(owned) memory associated with this ledger. + */ + @Override + public long getAccountedSize() { + synchronized (allocationManager) { + if (allocationManager.getOwningLedger() == this) { + return allocationManager.getSize(); + } else { + return 0; + } + } + } + + /** + * Print the current ledger state to the provided StringBuilder. + * + * @param sb The StringBuilder to populate. + * @param indent The level of indentation to position the data. + * @param verbosity The level of verbosity to print. + */ + void print(StringBuilder sb, int indent, BaseAllocator.Verbosity verbosity) { + CommonUtil.indent(sb, indent) + .append("ledger[") + .append(ledgerId) + .append("] allocator: ") + .append(allocator.getName()) + .append("), isOwning: ") + .append(", size: ") + .append(", references: ") + .append(bufRefCnt.get()) + .append(", life: ") + .append(lCreationTime) + .append("..") + .append(lDestructionTime) + .append(", allocatorManager: [") + .append(", life: "); + + if (!BaseAllocator.DEBUG) { + sb.append("]\n"); + } else { + synchronized (buffers) { + sb.append("] holds ") + .append(buffers.size()) + .append(" buffers. \n"); + for (ArrowBuf buf : buffers.keySet()) { + buf.print(sb, indent + 2, verbosity); + sb.append('\n'); + } + } + } + } + + /** + * Get the {@link AllocationManager} used by this BufferLedger. + * + * @return The AllocationManager used by this BufferLedger. + */ + public AllocationManager getAllocationManager() { + return allocationManager; + } + +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferManager.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferManager.java new file mode 100644 index 000000000..6b622e719 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/BufferManager.java @@ -0,0 +1,53 @@ +/* + * 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; + +/** + * Manages a list of {@link ArrowBuf}s that can be reallocated as needed. Upon + * re-allocation the old buffer will be freed. Managing a list of these buffers + * prevents some parts of the system from needing to define a correct location + * to place the final call to free them. + */ +public interface BufferManager extends AutoCloseable { + + /** + * Replace an old buffer with a new version at least of the provided size. Does not copy data. + * + * @param old Old Buffer that the user is no longer going to use. + * @param newSize Size of new replacement buffer. + * @return A new version of the buffer. + */ + ArrowBuf replace(ArrowBuf old, long newSize); + + /** + * Get a managed buffer of indeterminate size. + * + * @return A buffer. + */ + ArrowBuf getManagedBuffer(); + + /** + * Get a managed buffer of at least a certain size. + * + * @param size The desired size + * @return A buffer + */ + ArrowBuf getManagedBuffer(long size); + + void close(); +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java new file mode 100644 index 000000000..79b825aa2 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/CheckAllocator.java @@ -0,0 +1,87 @@ +/* + * 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 java.io.IOException; +import java.net.URL; +import java.util.Enumeration; +import java.util.LinkedHashSet; +import java.util.Set; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Static method to ensure we have a RootAllocator on the classpath and report which one is used. + */ +final class CheckAllocator { + private static final Logger logger = LoggerFactory.getLogger(CheckAllocator.class); + private static final String ALLOCATOR_PATH = "org/apache/arrow/memory/DefaultAllocationManagerFactory.class"; + + private CheckAllocator() { + + } + + static String check() { + Set<URL> urls = scanClasspath(); + URL rootAllocator = assertOnlyOne(urls); + reportResult(rootAllocator); + return "org.apache.arrow.memory.DefaultAllocationManagerFactory"; + } + + + private static Set<URL> scanClasspath() { + // LinkedHashSet appropriate here because it preserves insertion order + // during iteration + Set<URL> allocatorPathSet = new LinkedHashSet<>(); + try { + ClassLoader allocatorClassLoader = CheckAllocator.class.getClassLoader(); + Enumeration<URL> paths; + if (allocatorClassLoader == null) { + paths = ClassLoader.getSystemResources(ALLOCATOR_PATH); + } else { + paths = allocatorClassLoader.getResources(ALLOCATOR_PATH); + } + while (paths.hasMoreElements()) { + URL path = paths.nextElement(); + allocatorPathSet.add(path); + } + } catch (IOException ioe) { + logger.error("Error getting resources from path", ioe); + } + return allocatorPathSet; + } + + private static void reportResult(URL rootAllocator) { + String path = rootAllocator.getPath(); + String subPath = path.substring(path.indexOf("memory")); + logger.info("Using DefaultAllocationManager at {}", subPath); + } + + private static URL assertOnlyOne(Set<URL> urls) { + if (urls.size() > 1) { + logger.warn("More than one DefaultAllocationManager on classpath. Choosing first found"); + } + if (urls.isEmpty()) { + throw new RuntimeException("No DefaultAllocationManager found on classpath. Can't allocate Arrow buffers." + + " Please consider adding arrow-memory-netty or arrow-memory-unsafe as a dependency."); + } + return urls.iterator().next(); + } + +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java new file mode 100644 index 000000000..67156f89d --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ChildAllocator.java @@ -0,0 +1,44 @@ +/* + * 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; + + +/** + * Child allocator class. Only slightly different from the {@see RootAllocator}, + * in that these can't be created directly, but must be obtained from + * {@see BufferAllocator#newChildAllocator(AllocatorOwner, long, long, int)}. + * + * <p>Child allocators can only be created by the root, or other children, so + * this class is package private.</p> + */ +class ChildAllocator extends BaseAllocator { + + /** + * Constructor. + * + * @param parentAllocator parent allocator -- the one creating this child + * @param name the name of this child allocator + * @param config configuration of this child allocator + */ + ChildAllocator( + BaseAllocator parentAllocator, + String name, + Config config) { + super(parentAllocator, name, config); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java new file mode 100644 index 000000000..15120c252 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerOption.java @@ -0,0 +1,133 @@ +/* + * 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 java.lang.reflect.Field; + +/** + * A class for choosing the default allocation manager. + */ +public class DefaultAllocationManagerOption { + + /** + * The environmental variable to set the default allocation manager type. + */ + public static final String ALLOCATION_MANAGER_TYPE_ENV_NAME = "ARROW_ALLOCATION_MANAGER_TYPE"; + + /** + * The system property to set the default allocation manager type. + */ + public static final String ALLOCATION_MANAGER_TYPE_PROPERTY_NAME = "arrow.allocation.manager.type"; + + static final org.slf4j.Logger LOGGER = org.slf4j.LoggerFactory.getLogger(DefaultAllocationManagerOption.class); + + /** + * The default allocation manager factory. + */ + private static AllocationManager.Factory DEFAULT_ALLOCATION_MANAGER_FACTORY = null; + + /** + * The allocation manager type. + */ + public enum AllocationManagerType { + /** + * Netty based allocation manager. + */ + Netty, + + /** + * Unsafe based allocation manager. + */ + Unsafe, + + /** + * Unknown type. + */ + Unknown, + } + + static AllocationManagerType getDefaultAllocationManagerType() { + AllocationManagerType ret = AllocationManagerType.Unknown; + + try { + String envValue = System.getenv(ALLOCATION_MANAGER_TYPE_ENV_NAME); + ret = AllocationManagerType.valueOf(envValue); + } catch (IllegalArgumentException | NullPointerException e) { + // ignore the exception, and make the allocation manager type remain unchanged + } + + // system property takes precedence + try { + String propValue = System.getProperty(ALLOCATION_MANAGER_TYPE_PROPERTY_NAME); + ret = AllocationManagerType.valueOf(propValue); + } catch (IllegalArgumentException | NullPointerException e) { + // ignore the exception, and make the allocation manager type remain unchanged + } + return ret; + } + + static AllocationManager.Factory getDefaultAllocationManagerFactory() { + if (DEFAULT_ALLOCATION_MANAGER_FACTORY != null) { + return DEFAULT_ALLOCATION_MANAGER_FACTORY; + } + AllocationManagerType type = getDefaultAllocationManagerType(); + switch (type) { + case Netty: + DEFAULT_ALLOCATION_MANAGER_FACTORY = getNettyFactory(); + break; + case Unsafe: + DEFAULT_ALLOCATION_MANAGER_FACTORY = getUnsafeFactory(); + break; + case Unknown: + LOGGER.info("allocation manager type not specified, using netty as the default type"); + DEFAULT_ALLOCATION_MANAGER_FACTORY = getFactory(CheckAllocator.check()); + break; + default: + throw new IllegalStateException("Unknown allocation manager type: " + type); + } + return DEFAULT_ALLOCATION_MANAGER_FACTORY; + } + + private static AllocationManager.Factory getFactory(String clazzName) { + try { + Field field = Class.forName(clazzName).getDeclaredField("FACTORY"); + field.setAccessible(true); + return (AllocationManager.Factory) field.get(null); + } catch (Exception e) { + throw new RuntimeException("Unable to instantiate Allocation Manager for " + clazzName, e); + } + } + + private static AllocationManager.Factory getUnsafeFactory() { + try { + return getFactory("org.apache.arrow.memory.UnsafeAllocationManager"); + } catch (RuntimeException e) { + throw new RuntimeException("Please add arrow-memory-unsafe to your classpath," + + " No DefaultAllocationManager found to instantiate an UnsafeAllocationManager", e); + } + } + + private static AllocationManager.Factory getNettyFactory() { + try { + return getFactory("org.apache.arrow.memory.NettyAllocationManager"); + } catch (RuntimeException e) { + throw new RuntimeException("Please add arrow-memory-netty to your classpath," + + " No DefaultAllocationManager found to instantiate an NettyAllocationManager", e); + } + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/LowCostIdentityHashMap.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/LowCostIdentityHashMap.java new file mode 100644 index 000000000..edfa82392 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/LowCostIdentityHashMap.java @@ -0,0 +1,336 @@ +/* + * 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 org.apache.arrow.util.Preconditions; +import org.apache.arrow.util.VisibleForTesting; + +/** + * Highly specialized IdentityHashMap that implements only partial + * Map APIs. + * It incurs low initial cost (just two elements by default). + * It assumes Value includes the Key - Implements @ValueWithKeyIncluded iface + * that provides "getKey" method. + * + * @param <K> Key type + * @param <V> Value type + */ +public class LowCostIdentityHashMap<K, V extends ValueWithKeyIncluded<K>> { + + /* + * The internal data structure to hold values. + */ + private Object[] elementData; + + /* Actual number of values. */ + private int size; + + /* + * maximum number of elements that can be put in this map before having to + * rehash. + */ + private int threshold; + + private static final int DEFAULT_MIN_SIZE = 1; + + /* Default load factor of 0.75; */ + private static final int LOAD_FACTOR = 7500; + + /** + * Creates a Map with default expected maximum size. + */ + public LowCostIdentityHashMap() { + this(DEFAULT_MIN_SIZE); + } + + /** + * Creates a Map with the specified maximum size parameter. + * + * @param maxSize + * The estimated maximum number of entries that will be put in + * this map. + */ + public LowCostIdentityHashMap(int maxSize) { + if (maxSize >= 0) { + this.size = 0; + threshold = getThreshold(maxSize); + elementData = newElementArray(computeElementArraySize()); + } else { + throw new IllegalArgumentException(); + } + } + + private int getThreshold(int maxSize) { + // assign the threshold to maxSize initially, this will change to a + // higher value if rehashing occurs. + return maxSize > 2 ? maxSize : 2; + } + + private int computeElementArraySize() { + int arraySize = (int) (((long) threshold * 10000) / LOAD_FACTOR); + // ensure arraySize is positive, the above cast from long to int type + // leads to overflow and negative arraySize if threshold is too big + return arraySize < 0 ? -arraySize : arraySize; + } + + /** + * Create a new element array. + * + * @param s + * the number of elements + * @return Reference to the element array + */ + private Object[] newElementArray(int s) { + return new Object[s]; + } + + /** + * Removes all elements from this map, leaving it empty. + * + * @see #isEmpty() + * @see #size() + */ + public void clear() { + size = 0; + for (int i = 0; i < elementData.length; i++) { + elementData[i] = null; + } + } + + /** + * Returns whether this map contains the specified key. + * + * @param key + * the key to search for. + * @return {@code true} if this map contains the specified key, + * {@code false} otherwise. + */ + public boolean containsKey(K key) { + Preconditions.checkNotNull(key); + + int index = findIndex(key, elementData); + return (elementData[index] == null) ? false : ((V) elementData[index]).getKey() == key; + } + + /** + * Returns whether this map contains the specified value. + * + * @param value + * the value to search for. + * @return {@code true} if this map contains the specified value, + * {@code false} otherwise. + */ + public boolean containsValue(V value) { + Preconditions.checkNotNull(value); + + for (int i = 0; i < elementData.length; i++) { + if (elementData[i] == value) { + return true; + } + } + return false; + } + + /** + * Returns the value of the mapping with the specified key. + * + * @param key the key. + * @return the value of the mapping with the specified key. + */ + public V get(K key) { + Preconditions.checkNotNull(key); + + int index = findIndex(key, elementData); + + return (elementData[index] == null) ? null : + (((V) elementData[index]).getKey() == key) ? (V) elementData[index] : null; + } + + /** + * Returns the index where the key is found at, or the index of the next + * empty spot if the key is not found in this table. + */ + @VisibleForTesting + int findIndex(Object key, Object[] array) { + int length = array.length; + int index = getModuloHash(key, length); + int last = (index + length - 1) % length; + while (index != last) { + if ((array[index] == null) || ((V) array[index]).getKey() == key) { + /* + * Found the key, or the next empty spot (which means key is not + * in the table) + */ + break; + } + index = (index + 1) % length; + } + return index; + } + + @VisibleForTesting + static int getModuloHash(Object key, int length) { + return ((System.identityHashCode(key) & 0x7FFFFFFF) % length); + } + + /** + * Maps the specified key to the specified value. + * + * @param value the value. + * @return the value of any previous mapping with the specified key or + * {@code null} if there was no such mapping. + */ + public V put(V value) { + Preconditions.checkNotNull(value); + K key = value.getKey(); + Preconditions.checkNotNull(key); + + int index = findIndex(key, elementData); + + // if the key doesn't exist in the table + if (elementData[index] == null || ((V) elementData[index]).getKey() != key) { + if (++size > threshold) { + rehash(); + index = findIndex(key, elementData); + } + + // insert the key and assign the value to null initially + elementData[index] = null; + } + + // insert value to where it needs to go, return the old value + Object result = elementData[index]; + elementData[index] = value; + + return (V) result; + } + + @VisibleForTesting + void rehash() { + int newlength = elementData.length * 15 / 10; + if (newlength == 0) { + newlength = 1; + } + Object[] newData = newElementArray(newlength); + for (int i = 0; i < elementData.length; i++) { + Object key = (elementData[i] == null) ? null : ((V) elementData[i]).getKey(); + if (key != null) { + // if not empty + int index = findIndex(key, newData); + newData[index] = elementData[i]; + } + } + elementData = newData; + computeMaxSize(); + } + + private void computeMaxSize() { + threshold = (int) ((long) (elementData.length) * LOAD_FACTOR / 10000); + } + + /** + * Removes the mapping with the specified key from this map. + * + * @param key the key of the mapping to remove. + * @return the value of the removed mapping, or {@code null} if no mapping + * for the specified key was found. + */ + public V remove(K key) { + Preconditions.checkNotNull(key); + + boolean hashedOk; + int index; + int next; + int hash; + Object result; + Object object; + index = next = findIndex(key, elementData); + + if (elementData[index] == null || ((V) elementData[index]).getKey() != key) { + return null; + } + + // store the value for this key + result = elementData[index]; + // clear value to allow movement of the rest of the elements + elementData[index] = null; + size--; + + // shift the following elements up if needed + // until we reach an empty spot + int length = elementData.length; + while (true) { + next = (next + 1) % length; + object = elementData[next]; + if (object == null) { + break; + } + + hash = getModuloHash(((V) object).getKey(), length); + hashedOk = hash > index; + if (next < index) { + hashedOk = hashedOk || (hash <= next); + } else { + hashedOk = hashedOk && (hash <= next); + } + if (!hashedOk) { + elementData[index] = object; + index = next; + elementData[index] = null; + } + } + + return (V) result; + } + + + + /** + * Returns whether this Map has no elements. + * + * @return {@code true} if this Map has no elements, + * {@code false} otherwise. + * @see #size() + */ + public boolean isEmpty() { + return size == 0; + } + + /** + * Returns the number of mappings in this Map. + * + * @return the number of mappings in this Map. + */ + public int size() { + return size; + } + + /** + * Special API to return next value - substitute of regular Map.values.iterator().next(). + * + * @return next available value or null if none available + */ + public V getNextValue() { + for (int i = 0; i < elementData.length; i++) { + if (elementData[i] != null) { + return (V) elementData[i]; + } + } + return null; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OutOfMemoryException.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OutOfMemoryException.java new file mode 100644 index 000000000..841ffcd8f --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OutOfMemoryException.java @@ -0,0 +1,67 @@ +/* + * 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 java.util.Optional; + +/** + * Indicates memory could not be allocated for Arrow buffers. + * + * <p>This is different from {@linkplain OutOfMemoryError} which indicates the JVM + * is out of memory. This error indicates that static limit of one of Arrow's + * allocators (e.g. {@linkplain BaseAllocator}) has been exceeded. + */ +public class OutOfMemoryException extends RuntimeException { + + static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(OutOfMemoryException + .class); + private static final long serialVersionUID = -6858052345185793382L; + private Optional<AllocationOutcomeDetails> outcomeDetails = Optional.empty(); + + public OutOfMemoryException() { + super(); + } + + public OutOfMemoryException(String message, Throwable cause, boolean enableSuppression, boolean + writableStackTrace) { + super(message, cause, enableSuppression, writableStackTrace); + } + + public OutOfMemoryException(String message, Throwable cause) { + super(message, cause); + + } + + public OutOfMemoryException(String message) { + super(message); + } + + public OutOfMemoryException(String message, Optional<AllocationOutcomeDetails> details) { + super(message); + this.outcomeDetails = details; + } + + public OutOfMemoryException(Throwable cause) { + super(cause); + + } + + public Optional<AllocationOutcomeDetails> getOutcomeDetails() { + return outcomeDetails; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OwnershipTransferNOOP.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OwnershipTransferNOOP.java new file mode 100644 index 000000000..d4fed8448 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OwnershipTransferNOOP.java @@ -0,0 +1,39 @@ +/* + * 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; + +/** + * An {@link OwnershipTransferResult} indicating no transfer needed. + */ +public class OwnershipTransferNOOP implements OwnershipTransferResult { + private final ArrowBuf buffer; + + OwnershipTransferNOOP(final ArrowBuf buf) { + this.buffer = buf; + } + + @Override + public ArrowBuf getTransferredBuffer() { + return buffer; + } + + @Override + public boolean getAllocationFit() { + return true; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OwnershipTransferResult.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OwnershipTransferResult.java new file mode 100644 index 000000000..ef857d827 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/OwnershipTransferResult.java @@ -0,0 +1,28 @@ +/* + * 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; + +/** + * The result of transferring an {@link ArrowBuf} between {@linkplain BufferAllocator}s. + */ +public interface OwnershipTransferResult { + + boolean getAllocationFit(); + + ArrowBuf getTransferredBuffer(); +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/README.md b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/README.md new file mode 100644 index 000000000..f5f924ce8 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/README.md @@ -0,0 +1,121 @@ +<!-- + 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. +--> +# Memory: Allocation, Accounting and Management + +The memory management package contains all the memory allocation related items that Arrow uses to manage memory. + + +## Key Components +Memory management can be broken into the following main components: + +- Memory chunk allocation and fragmentation management + - `PooledByteBufAllocatorL` - A LittleEndian clone of Netty's jemalloc implementation + - `UnsafeDirectLittleEndian` - A base level memory access interface + - `LargeBuffer` - A buffer backing implementation used when working with data larger than one Netty chunk (default to 16mb) +- Memory limits & Accounting + - `Accountant` - A nestable class of lockfree memory accountors. +- Application-level memory allocation + - `BufferAllocator` - The public interface application users should be leveraging + - `BaseAllocator` - The base implementation of memory allocation, contains the meat of our the Arrow allocator implementation + - `RootAllocator` - The root allocator. Typically only one created for a JVM + - `ChildAllocator` - A child allocator that derives from the root allocator +- Buffer ownership and transfer capabilities + - `AllocationManager` - Responsible for managing the relationship between multiple allocators and a single chunk of memory + - `BufferLedger` - Responsible for allowing maintaining the relationship between an `AllocationManager`, a `BufferAllocator` and one or more individual `ArrowBuf`s +- Memory access + - `ArrowBuf` - The facade for interacting directly with a chunk of memory. + + +## Memory Management Overview +Arrow's memory model is based on the following basic concepts: + + - Memory can be allocated up to some limit. That limit could be a real limit (OS/JVM) or a locally imposed limit. + - Allocation operates in two phases: accounting then actual allocation. Allocation could fail at either point. + - Allocation failure should be recoverable. In all cases, the Allocator infrastructure should expose memory allocation failures (OS or internal limit-based) as `OutOfMemoryException`s. + - Any allocator can reserve memory when created. This memory shall be held such that this allocator will always be able to allocate that amount of memory. + - A particular application component should work to use a local allocator to understand local memory usage and better debug memory leaks. + - The same physical memory can be shared by multiple allocators and the allocator must provide an accounting paradigm for this purpose. + +## Allocator Trees + +Arrow provides a tree-based model for memory allocation. The RootAllocator is created first, then all allocators are created as children of that allocator. The RootAllocator is responsible for being the master bookkeeper for memory allocations. All other allocators are created as children of this tree. Each allocator can first determine whether it has enough local memory to satisfy a particular request. If not, the allocator can ask its parent for an additional memory allocation. + +## Reserving Memory + +Arrow provides two different ways to reserve memory: + + - BufferAllocator accounting reservations: + When a new allocator (other than the `RootAllocator`) is initialized, it can set aside memory that it will keep locally for its lifetime. This is memory that will never be released back to its parent allocator until the allocator is closed. + - `AllocationReservation` via BufferAllocator.newReservation(): Allows a short-term preallocation strategy so that a particular subsystem can ensure future memory is available to support a particular request. + +## Memory Ownership, Reference Counts and Sharing +Many BufferAllocators can reference the same piece of memory at the same time. The most common situation for this is in the case of a Broadcast Join: in this situation many downstream operators in the same Arrowbit will receive the same physical memory. Each of these operators will be operating within its own Allocator context. We therefore have multiple allocators all pointing at the same physical memory. It is the AllocationManager's responsibility to ensure that in this situation, that all memory is accurately accounted for from the Root's perspective and also to ensure that the memory is correctly released once all BufferAllocators have stopped using that memory. + +For simplicity of accounting, we treat that memory as being used by one of the BufferAllocators associated with the memory. When that allocator releases its claim on that memory, the memory ownership is then moved to another BufferLedger belonging to the same AllocationManager. Note that because a ArrowBuf.release() is what actually causes memory ownership transfer to occur, we always precede with ownership transfer (even if that violates an allocator limit). It is the responsibility of the application owning a particular allocator to frequently confirm whether the allocator is over its memory limit (BufferAllocator.isOverLimit()) and if so, attempt to aggressively release memory to ameliorate the situation. + +All ArrowBufs (direct or sliced) related to a single BufferLedger/BufferAllocator combination share the same reference count and either all will be valid or all will be invalid. + +## Object Hierarchy + +There are two main ways that someone can look at the object hierarchy for Arrow's memory management scheme. The first is a memory based perspective as below: + +### Memory Perspective +<pre> ++ AllocationManager +| +|-- UnsignedDirectLittleEndian (One per AllocationManager) +| +|-+ BufferLedger 1 ==> Allocator A (owning) +| ` - ArrowBuf 1 +|-+ BufferLedger 2 ==> Allocator B (non-owning) +| ` - ArrowBuf 2 +|-+ BufferLedger 3 ==> Allocator C (non-owning) + | - ArrowBuf 3 + | - ArrowBuf 4 + ` - ArrowBuf 5 +</pre> + +In this picture, a piece of memory is owned by an allocator manager. An allocator manager is responsible for that piece of memory no matter which allocator(s) it is working with. An allocator manager will have relationships with a piece of raw memory (via its reference to UnsignedDirectLittleEndian) as well as references to each BufferAllocator it has a relationship to. + +### Allocator Perspective +<pre> ++ RootAllocator +|-+ ChildAllocator 1 +| | - ChildAllocator 1.1 +| ` ... +| +|-+ ChildAllocator 2 +|-+ ChildAllocator 3 +| | +| |-+ BufferLedger 1 ==> AllocationManager 1 (owning) ==> UDLE +| | `- ArrowBuf 1 +| `-+ BufferLedger 2 ==> AllocationManager 2 (non-owning)==> UDLE +| `- ArrowBuf 2 +| +|-+ BufferLedger 3 ==> AllocationManager 1 (non-owning)==> UDLE +| ` - ArrowBuf 3 +|-+ BufferLedger 4 ==> AllocationManager 2 (owning) ==> UDLE + | - ArrowBuf 4 + | - ArrowBuf 5 + ` - ArrowBuf 6 +</pre> + +In this picture, a RootAllocator owns three ChildAllocators. The first ChildAllocator (ChildAllocator 1) owns a subsequent ChildAllocator. ChildAllocator has two BufferLedgers/AllocationManager references. Coincidentally, each of these AllocationManager's is also associated with the RootAllocator. In this case, one of the these AllocationManagers is owned by ChildAllocator 3 (AllocationManager 1) while the other AllocationManager (AllocationManager 2) is owned/accounted for by the RootAllocator. Note that in this scenario, ArrowBuf 1 is sharing the underlying memory as ArrowBuf 3. However the subset of that memory (e.g. through slicing) might be different. Also note that ArrowBuf 2 and ArrowBuf 4, 5 and 6 are also sharing the same underlying memory. Also note that ArrowBuf 4, 5 and 6 all share the same reference count and fate. + +## Debugging Issues +The Allocator object provides a useful set of tools to better understand the status of the allocator. If in `DEBUG` mode, the allocator and supporting classes will record additional debug tracking information to better track down memory leaks and issues. To enable DEBUG mode, either enable Java assertions with `-ea` or pass the following system property to the VM when starting `-Darrow.memory.debug.allocator=true`. The BufferAllocator also provides a `BufferAllocator.toVerboseString()` which can be used in DEBUG mode to get extensive stacktrace information and events associated with various Allocator behaviors. diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ReferenceManager.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ReferenceManager.java new file mode 100644 index 000000000..00ae274b7 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ReferenceManager.java @@ -0,0 +1,175 @@ +/* + * 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; + +/** + * Reference Manager manages one or more ArrowBufs that share the + * reference count for the underlying memory chunk. + */ +public interface ReferenceManager { + + /** + * Return the reference count. + * @return reference count + */ + int getRefCount(); + + /** + * Decrement this reference manager's reference count by 1 for the associated underlying + * memory. If the reference count drops to 0, it implies that ArrowBufs managed by this + * reference manager no longer need access to the underlying memory + * @return true if ref count has dropped to 0, false otherwise + */ + boolean release(); + + /** + * Decrement this reference manager's reference count for the associated underlying + * memory. If the reference count drops to 0, it implies that ArrowBufs managed by this + * reference manager no longer need access to the underlying memory + * @param decrement the count to decrease the reference count by + * @return the new reference count + */ + boolean release(int decrement); + + /** + * Increment this reference manager's reference count by 1 for the associated underlying + * memory. + */ + void retain(); + + /** + * Increment this reference manager's reference count by a given amount for the + * associated underlying memory. + * @param increment the count to increase the reference count by + */ + void retain(int increment); + + /** + * Create a new ArrowBuf that is associated with an alternative allocator for the purposes of + * memory ownership and accounting. This has no impact on the reference counting for the current + * ArrowBuf except in the situation where the passed in Allocator is the same as the current buffer. + * This operation has no impact on the reference count of this ArrowBuf. The newly created + * ArrowBuf with either have a reference count of 1 (in the case that this is the first time this + * memory is being associated with the target allocator or in other words allocation manager currently + * doesn't hold a mapping for the target allocator) or the current value of the reference count for + * the target allocator-reference manager combination + 1 in the case that the provided allocator + * already had an association to this underlying memory. + * + * @param srcBuffer source ArrowBuf + * @param targetAllocator The target allocator to create an association with. + * @return A new ArrowBuf which shares the same underlying memory as this ArrowBuf. + */ + ArrowBuf retain(ArrowBuf srcBuffer, BufferAllocator targetAllocator); + + /** + * Derive a new ArrowBuf from a given source ArrowBuf. The new derived + * ArrowBuf will share the same reference count as rest of the ArrowBufs + * associated with this reference manager. + * @param sourceBuffer source ArrowBuf + * @param index index (relative to source ArrowBuf) new ArrowBuf should be derived from + * @param length length (bytes) of data in underlying memory that derived buffer will + * have access to in underlying memory + * @return derived buffer + */ + ArrowBuf deriveBuffer(ArrowBuf sourceBuffer, long index, long length); + + /** + * Transfer the memory accounting ownership of this ArrowBuf to another allocator. + * This will generate a new ArrowBuf that carries an association with the underlying memory + * for the given ArrowBuf + * @param sourceBuffer source ArrowBuf + * @param targetAllocator The target allocator to create an association with + * @return {@link OwnershipTransferResult} with info on transfer result and new buffer + */ + OwnershipTransferResult transferOwnership(ArrowBuf sourceBuffer, BufferAllocator targetAllocator); + + /** + * Get the buffer allocator associated with this reference manager. + * @return buffer allocator. + */ + BufferAllocator getAllocator(); + + /** + * Total size (in bytes) of memory underlying this reference manager. + * @return Size (in bytes) of the memory chunk. + */ + long getSize(); + + /** + * Get the total accounted size (in bytes). + * @return accounted size. + */ + long getAccountedSize(); + + String NO_OP_ERROR_MESSAGE = "Operation not supported on NO_OP Reference Manager"; + + // currently used for empty ArrowBufs + ReferenceManager NO_OP = new ReferenceManager() { + @Override + public int getRefCount() { + return 1; + } + + @Override + public boolean release() { + return false; + } + + @Override + public boolean release(int decrement) { + return false; + } + + @Override + public void retain() { } + + @Override + public void retain(int increment) { } + + @Override + public ArrowBuf retain(ArrowBuf srcBuffer, BufferAllocator targetAllocator) { + return srcBuffer; + } + + @Override + public ArrowBuf deriveBuffer(ArrowBuf sourceBuffer, long index, long length) { + return sourceBuffer; + } + + @Override + public OwnershipTransferResult transferOwnership(ArrowBuf sourceBuffer, BufferAllocator targetAllocator) { + return new OwnershipTransferNOOP(sourceBuffer); + } + + @Override + public BufferAllocator getAllocator() { + return new RootAllocator(0); + } + + @Override + public long getSize() { + return 0L; + } + + @Override + public long getAccountedSize() { + return 0L; + } + + }; +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/RootAllocator.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/RootAllocator.java new file mode 100644 index 000000000..89889118c --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/RootAllocator.java @@ -0,0 +1,71 @@ +/* + * 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 org.apache.arrow.memory.rounding.DefaultRoundingPolicy; +import org.apache.arrow.memory.rounding.RoundingPolicy; +import org.apache.arrow.util.VisibleForTesting; + +/** + * A root allocator for using direct memory for Arrow Vectors/Arrays. Supports creating a + * tree of descendant child allocators to facilitate better instrumentation of memory + * allocations. + */ +public class RootAllocator extends BaseAllocator { + + public RootAllocator() { + this(AllocationListener.NOOP, Long.MAX_VALUE); + } + + public RootAllocator(final long limit) { + this(AllocationListener.NOOP, limit); + } + + public RootAllocator(final AllocationListener listener, final long limit) { + //todo fix DefaultRoundingPolicy when using Netty + this(listener, limit, DefaultRoundingPolicy.DEFAULT_ROUNDING_POLICY); + } + + /** + * Constructor. + * + * @param listener the allocation listener + * @param limit max allocation size in bytes + * @param roundingPolicy the policy for rounding the buffer size + */ + public RootAllocator(final AllocationListener listener, final long limit, RoundingPolicy roundingPolicy) { + this(configBuilder() + .listener(listener) + .maxAllocation(limit) + .roundingPolicy(roundingPolicy) + .build() + ); + } + + public RootAllocator(Config config) { + super(null, "ROOT", config); + } + + /** + * Verify the accounting state of the allocation system. + */ + @VisibleForTesting + public void verify() { + verifyAllocator(); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ValueWithKeyIncluded.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ValueWithKeyIncluded.java new file mode 100644 index 000000000..2699b6a46 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ValueWithKeyIncluded.java @@ -0,0 +1,28 @@ +/* + * 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; + +/** + * Helper interface to generify a value to be included in the map where + * key is part of the value. + * + * @param <K> The type of the key. + */ +public interface ValueWithKeyIncluded<K> { + K getKey(); +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/package-info.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/package-info.java new file mode 100644 index 000000000..5aef955a3 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/package-info.java @@ -0,0 +1,26 @@ +/* + * 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. + */ + +/** + * Memory Allocation, Account and Management + * + * See the README.md file in this directory for detailed information about Arrow's memory + * allocation subsystem. + * + */ + +package org.apache.arrow.memory; diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java new file mode 100644 index 000000000..7ba231b0c --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/DefaultRoundingPolicy.java @@ -0,0 +1,114 @@ +/* + * 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.rounding; + +import org.apache.arrow.memory.util.CommonUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * The default rounding policy. That is, if the requested size is within the chunk size, + * the rounded size will be the next power of two. Otherwise, the rounded size + * will be identical to the requested size. + */ +public class DefaultRoundingPolicy implements RoundingPolicy { + private static final Logger logger = LoggerFactory.getLogger(DefaultRoundingPolicy.class); + public final long chunkSize; + + /** + * The variables here and the static block calculates the DEFAULT_CHUNK_SIZE. + * + * <p> + * It was copied from {@link io.netty.buffer.PooledByteBufAllocator}. + * </p> + */ + private static final int MIN_PAGE_SIZE = 4096; + private static final int MAX_CHUNK_SIZE = (int) (((long) Integer.MAX_VALUE + 1) / 2); + private static final long DEFAULT_CHUNK_SIZE; + + + static { + int defaultPageSize = Integer.getInteger("org.apache.memory.allocator.pageSize", 8192); + Throwable pageSizeFallbackCause = null; + try { + validateAndCalculatePageShifts(defaultPageSize); + } catch (Throwable t) { + pageSizeFallbackCause = t; + defaultPageSize = 8192; + } + + int defaultMaxOrder = Integer.getInteger("org.apache.memory.allocator.maxOrder", 11); + Throwable maxOrderFallbackCause = null; + try { + validateAndCalculateChunkSize(defaultPageSize, defaultMaxOrder); + } catch (Throwable t) { + maxOrderFallbackCause = t; + defaultMaxOrder = 11; + } + DEFAULT_CHUNK_SIZE = validateAndCalculateChunkSize(defaultPageSize, defaultMaxOrder); + if (logger.isDebugEnabled()) { + logger.debug("-Dorg.apache.memory.allocator.pageSize: {}", defaultPageSize); + logger.debug("-Dorg.apache.memory.allocator.maxOrder: {}", defaultMaxOrder); + } + } + + private static int validateAndCalculatePageShifts(int pageSize) { + if (pageSize < MIN_PAGE_SIZE) { + throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: " + MIN_PAGE_SIZE + ")"); + } + + if ((pageSize & pageSize - 1) != 0) { + throw new IllegalArgumentException("pageSize: " + pageSize + " (expected: power of 2)"); + } + + // Logarithm base 2. At this point we know that pageSize is a power of two. + return Integer.SIZE - 1 - Integer.numberOfLeadingZeros(pageSize); + } + + private static int validateAndCalculateChunkSize(int pageSize, int maxOrder) { + if (maxOrder > 14) { + throw new IllegalArgumentException("maxOrder: " + maxOrder + " (expected: 0-14)"); + } + + // Ensure the resulting chunkSize does not overflow. + int chunkSize = pageSize; + for (int i = maxOrder; i > 0; i --) { + if (chunkSize > MAX_CHUNK_SIZE / 2) { + throw new IllegalArgumentException(String.format( + "pageSize (%d) << maxOrder (%d) must not exceed %d", pageSize, maxOrder, MAX_CHUNK_SIZE)); + } + chunkSize <<= 1; + } + return chunkSize; + } + + /** + * The singleton instance. + */ + public static final DefaultRoundingPolicy DEFAULT_ROUNDING_POLICY = new DefaultRoundingPolicy(DEFAULT_CHUNK_SIZE); + + private DefaultRoundingPolicy(long chunkSize) { + this.chunkSize = chunkSize; + } + + @Override + public long getRoundedSize(long requestSize) { + return requestSize < chunkSize ? + CommonUtil.nextPowerOfTwo(requestSize) : requestSize; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/RoundingPolicy.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/RoundingPolicy.java new file mode 100644 index 000000000..434b9a0cd --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/RoundingPolicy.java @@ -0,0 +1,26 @@ +/* + * 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.rounding; + +/** + * The policy for rounding the buffer size, to improve performance and avoid memory fragmentation. + * In particular, given a requested buffer size, the policy will determine the rounded buffer size. + */ +public interface RoundingPolicy { + long getRoundedSize(long requestSize); +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java new file mode 100644 index 000000000..d2bc4451d --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/rounding/SegmentRoundingPolicy.java @@ -0,0 +1,60 @@ +/* + * 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.rounding; + +import org.apache.arrow.util.Preconditions; + +/** + * The rounding policy that each buffer size must a multiple of the segment size. + */ +public class SegmentRoundingPolicy implements RoundingPolicy { + + /** + * The minimal segment size. + */ + public static final long MIN_SEGMENT_SIZE = 1024L; + + /** + * The segment size. It must be at least {@link SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, + * and be a power of 2. + */ + private int segmentSize; + + /** + * Constructor for the segment rounding policy. + * @param segmentSize the segment size. + * @throws IllegalArgumentException if the segment size is smaller than + * {@link SegmentRoundingPolicy#MIN_SEGMENT_SIZE}, or is not a power of 2. + */ + public SegmentRoundingPolicy(int segmentSize) { + Preconditions.checkArgument(segmentSize >= MIN_SEGMENT_SIZE, + "The segment size cannot be smaller than %s", MIN_SEGMENT_SIZE); + Preconditions.checkArgument((segmentSize & (segmentSize - 1)) == 0, + "The segment size must be a power of 2"); + this.segmentSize = segmentSize; + } + + @Override + public long getRoundedSize(long requestSize) { + return (requestSize + (segmentSize - 1)) / segmentSize * segmentSize; + } + + public int getSegmentSize() { + return segmentSize; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/ArrowBufPointer.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/ArrowBufPointer.java new file mode 100644 index 000000000..fa1cfbdb2 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/ArrowBufPointer.java @@ -0,0 +1,187 @@ +/* + * 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.util; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.memory.util.hash.SimpleHasher; +import org.apache.arrow.util.Preconditions; + +/** + * Pointer to a memory region within an {@link ArrowBuf}. + * It will be used as the basis for calculating hash code within a vector, and equality determination. + */ +public final class ArrowBufPointer { + + /** + * The hash code when the arrow buffer is null. + */ + public static final int NULL_HASH_CODE = 0; + + private ArrowBuf buf; + + private long offset; + + private long length; + + private int hashCode = NULL_HASH_CODE; + + private final ArrowBufHasher hasher; + + /** + * A flag indicating if the underlying memory region has changed. + */ + private boolean hashCodeChanged = false; + + /** + * The default constructor. + */ + public ArrowBufPointer() { + this(SimpleHasher.INSTANCE); + } + + /** + * Constructs an arrow buffer pointer with the specified hasher. + * @param hasher the hasher to use. + */ + public ArrowBufPointer(ArrowBufHasher hasher) { + Preconditions.checkNotNull(hasher); + this.hasher = hasher; + } + + /** + * Constructs an Arrow buffer pointer. + * @param buf the underlying {@link ArrowBuf}, which can be null. + * @param offset the start off set of the memory region pointed to. + * @param length the length off set of the memory region pointed to. + */ + public ArrowBufPointer(ArrowBuf buf, long offset, long length) { + this(buf, offset, length, SimpleHasher.INSTANCE); + } + + /** + * Constructs an Arrow buffer pointer. + * @param buf the underlying {@link ArrowBuf}, which can be null. + * @param offset the start off set of the memory region pointed to. + * @param length the length off set of the memory region pointed to. + * @param hasher the hasher used to calculate the hash code. + */ + public ArrowBufPointer(ArrowBuf buf, long offset, long length, ArrowBufHasher hasher) { + Preconditions.checkNotNull(hasher); + this.hasher = hasher; + set(buf, offset, length); + } + + /** + * Sets this pointer. + * @param buf the underlying {@link ArrowBuf}, which can be null. + * @param offset the start off set of the memory region pointed to. + * @param length the length off set of the memory region pointed to. + */ + public void set(ArrowBuf buf, long offset, long length) { + this.buf = buf; + this.offset = offset; + this.length = length; + + hashCodeChanged = true; + } + + /** + * Gets the underlying buffer, or null if the underlying data is invalid or null. + * @return the underlying buffer, if any, or null if the underlying data is invalid or null. + */ + public ArrowBuf getBuf() { + return buf; + } + + public long getOffset() { + return offset; + } + + public long getLength() { + return length; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + if (!hasher.equals(((ArrowBufPointer) o).hasher)) { + // note that the hasher is incorporated in equality determination + // this is to avoid problems in cases where two Arrow buffer pointers are not equal + // while having different hashers and equal hash codes. + return false; + } + + ArrowBufPointer other = (ArrowBufPointer) o; + if (buf == null || other.buf == null) { + if (buf == null && other.buf == null) { + return true; + } else { + return false; + } + } + + return ByteFunctionHelpers.equal(buf, offset, offset + length, + other.buf, other.offset, other.offset + other.length) != 0; + } + + @Override + public int hashCode() { + if (!hashCodeChanged) { + return hashCode; + } + + // re-compute the hash code + if (buf == null) { + hashCode = NULL_HASH_CODE; + } else { + hashCode = hasher.hashCode(buf, offset, length); + } + + hashCodeChanged = false; + return hashCode; + } + + /** + * Compare two arrow buffer pointers. + * The comparison is based on lexicographic order. + * @param that the other pointer to compare. + * @return 0 if the two pointers are equal; + * a positive integer if this pointer is larger; + * a negative integer if this pointer is smaller. + */ + public int compareTo(ArrowBufPointer that) { + if (this.buf == null || that.buf == null) { + if (this.buf == null && that.buf == null) { + return 0; + } else { + // null is smaller + return this.buf == null ? -1 : 1; + } + } + + return ByteFunctionHelpers.compare(this.buf, this.offset, this.offset + this.length, + that.buf, that.offset, that.offset + that.length); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/AssertionUtil.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/AssertionUtil.java new file mode 100644 index 000000000..5e5b331fa --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/AssertionUtil.java @@ -0,0 +1,40 @@ +/* + * 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.util; + +/** + * Utility class to that provides {@link #ASSERT_ENABLED} constant to determine if assertions are enabled. + */ +public class AssertionUtil { + + public static final boolean ASSERT_ENABLED; + static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(AssertionUtil.class); + + static { + boolean isAssertEnabled = false; + assert isAssertEnabled = true; + ASSERT_ENABLED = isAssertEnabled; + } + + private AssertionUtil() { + } + + public static boolean isAssertionsEnabled() { + return ASSERT_ENABLED; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/AutoCloseableLock.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/AutoCloseableLock.java new file mode 100644 index 000000000..95228cf78 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/AutoCloseableLock.java @@ -0,0 +1,43 @@ +/* + * 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.util; + +import java.util.concurrent.locks.Lock; + +/** + * Simple wrapper class that allows Locks to be released via a try-with-resources block. + */ +public class AutoCloseableLock implements AutoCloseable { + + private final Lock lock; + + public AutoCloseableLock(Lock lock) { + this.lock = lock; + } + + public AutoCloseableLock open() { + lock.lock(); + return this; + } + + @Override + public void close() { + lock.unlock(); + } + +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/ByteFunctionHelpers.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/ByteFunctionHelpers.java new file mode 100644 index 000000000..9579245ca --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/ByteFunctionHelpers.java @@ -0,0 +1,347 @@ +/* + * 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.util; + +import java.nio.ByteOrder; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BoundsChecking; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.memory.util.hash.SimpleHasher; + +/** + * Utility methods for memory comparison at a byte level. + */ +public class ByteFunctionHelpers { + static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(ByteFunctionHelpers.class); + + private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; + + private ByteFunctionHelpers() {} + + /** + * Helper function to check for equality of bytes in two ArrowBufs. + * + * @param left Left ArrowBuf for comparison + * @param lStart start offset in the buffer + * @param lEnd end offset in the buffer + * @param right Right ArrowBuf for comparison + * @param rStart start offset in the buffer + * @param rEnd end offset in the buffer + * @return 1 if equals, 0 otherwise + */ + public static int equal(final ArrowBuf left, long lStart, long lEnd, final ArrowBuf right, long rStart, long rEnd) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + left.checkBytes(lStart, lEnd); + right.checkBytes(rStart, rEnd); + } + return memEqual(left.memoryAddress(), lStart, lEnd, right.memoryAddress(), rStart, rEnd); + } + + private static int memEqual(final long laddr, long lStart, long lEnd, final long raddr, long rStart, + final long rEnd) { + + long n = lEnd - lStart; + if (n == rEnd - rStart) { + long lPos = laddr + lStart; + long rPos = raddr + rStart; + + while (n > 63) { + for (int x = 0; x < 8; x++) { + long leftLong = MemoryUtil.UNSAFE.getLong(lPos); + long rightLong = MemoryUtil.UNSAFE.getLong(rPos); + if (leftLong != rightLong) { + return 0; + } + lPos += 8; + rPos += 8; + } + n -= 64; + } + + while (n > 7) { + long leftLong = MemoryUtil.UNSAFE.getLong(lPos); + long rightLong = MemoryUtil.UNSAFE.getLong(rPos); + if (leftLong != rightLong) { + return 0; + } + lPos += 8; + rPos += 8; + n -= 8; + } + + if (n > 3) { + int leftInt = MemoryUtil.UNSAFE.getInt(lPos); + int rightInt = MemoryUtil.UNSAFE.getInt(rPos); + if (leftInt != rightInt) { + return 0; + } + lPos += 4; + rPos += 4; + n -= 4; + } + + while (n-- != 0) { + byte leftByte = MemoryUtil.UNSAFE.getByte(lPos); + byte rightByte = MemoryUtil.UNSAFE.getByte(rPos); + if (leftByte != rightByte) { + return 0; + } + lPos++; + rPos++; + } + return 1; + } else { + return 0; + } + } + + /** + * Helper function to compare a set of bytes in two ArrowBufs. + * + * <p>Function will check data before completing in the case that + * + * @param left Left ArrowBuf to compare + * @param lStart start offset in the buffer + * @param lEnd end offset in the buffer + * @param right Right ArrowBuf to compare + * @param rStart start offset in the buffer + * @param rEnd end offset in the buffer + * @return 1 if left input is greater, -1 if left input is smaller, 0 otherwise + */ + public static int compare( + final ArrowBuf left, + long lStart, + long lEnd, + final ArrowBuf right, + long rStart, + long rEnd) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + left.checkBytes(lStart, lEnd); + right.checkBytes(rStart, rEnd); + } + return memcmp(left.memoryAddress(), lStart, lEnd, right.memoryAddress(), rStart, rEnd); + } + + private static int memcmp( + final long laddr, + long lStart, + long lEnd, + final long raddr, + long rStart, + final long rEnd) { + long lLen = lEnd - lStart; + long rLen = rEnd - rStart; + long n = Math.min(rLen, lLen); + long lPos = laddr + lStart; + long rPos = raddr + rStart; + + while (n > 63) { + for (int x = 0; x < 8; x++) { + long leftLong = MemoryUtil.UNSAFE.getLong(lPos); + long rightLong = MemoryUtil.UNSAFE.getLong(rPos); + if (leftLong != rightLong) { + if (LITTLE_ENDIAN) { + return unsignedLongCompare(Long.reverseBytes(leftLong), Long.reverseBytes(rightLong)); + } else { + return unsignedLongCompare(leftLong, rightLong); + } + } + lPos += 8; + rPos += 8; + } + n -= 64; + } + + while (n > 7) { + long leftLong = MemoryUtil.UNSAFE.getLong(lPos); + long rightLong = MemoryUtil.UNSAFE.getLong(rPos); + if (leftLong != rightLong) { + if (LITTLE_ENDIAN) { + return unsignedLongCompare(Long.reverseBytes(leftLong), Long.reverseBytes(rightLong)); + } else { + return unsignedLongCompare(leftLong, rightLong); + } + } + lPos += 8; + rPos += 8; + n -= 8; + } + + if (n > 3) { + int leftInt = MemoryUtil.UNSAFE.getInt(lPos); + int rightInt = MemoryUtil.UNSAFE.getInt(rPos); + if (leftInt != rightInt) { + if (LITTLE_ENDIAN) { + return unsignedIntCompare(Integer.reverseBytes(leftInt), Integer.reverseBytes(rightInt)); + } else { + return unsignedIntCompare(leftInt, rightInt); + } + } + lPos += 4; + rPos += 4; + n -= 4; + } + + while (n-- != 0) { + byte leftByte = MemoryUtil.UNSAFE.getByte(lPos); + byte rightByte = MemoryUtil.UNSAFE.getByte(rPos); + if (leftByte != rightByte) { + return ((leftByte & 0xFF) - (rightByte & 0xFF)) > 0 ? 1 : -1; + } + lPos++; + rPos++; + } + + if (lLen == rLen) { + return 0; + } + + return lLen > rLen ? 1 : -1; + + } + + /** + * Helper function to compare a set of bytes in ArrowBuf to a ByteArray. + * + * @param left Left ArrowBuf for comparison purposes + * @param lStart start offset in the buffer + * @param lEnd end offset in the buffer + * @param right second input to be compared + * @param rStart start offset in the byte array + * @param rEnd end offset in the byte array + * @return 1 if left input is greater, -1 if left input is smaller, 0 otherwise + */ + public static int compare( + final ArrowBuf left, + int lStart, + int lEnd, + final byte[] right, + int rStart, + final int rEnd) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + left.checkBytes(lStart, lEnd); + } + return memcmp(left.memoryAddress(), lStart, lEnd, right, rStart, rEnd); + } + + + /** + * Compares the two specified {@code long} values, treating them as unsigned values between + * {@code 0} and {@code 2^64 - 1} inclusive. + * + * @param a the first unsigned {@code long} to compare + * @param b the second unsigned {@code long} to compare + * @return a negative value if {@code a} is less than {@code b}; a positive value if {@code a} is + * greater than {@code b}; or zero if they are equal + */ + public static int unsignedLongCompare(long a, long b) { + return Long.compare(a ^ Long.MIN_VALUE, b ^ Long.MIN_VALUE); + } + + public static int unsignedIntCompare(int a, int b) { + return Integer.compare(a ^ Integer.MIN_VALUE, b ^ Integer.MIN_VALUE); + } + + private static int memcmp( + final long laddr, + int lStart, + int lEnd, + final byte[] right, + int rStart, + final int rEnd) { + int lLen = lEnd - lStart; + int rLen = rEnd - rStart; + int n = Math.min(rLen, lLen); + long lPos = laddr + lStart; + int rPos = rStart; + + while (n > 7) { + long leftLong = MemoryUtil.UNSAFE.getLong(lPos); + long rightLong = MemoryUtil.UNSAFE.getLong(right, MemoryUtil.BYTE_ARRAY_BASE_OFFSET + rPos); + if (leftLong != rightLong) { + if (LITTLE_ENDIAN) { + return unsignedLongCompare(Long.reverseBytes(leftLong), Long.reverseBytes(rightLong)); + } else { + return unsignedLongCompare(leftLong, rightLong); + } + } + lPos += 8; + rPos += 8; + n -= 8; + } + + if (n > 3) { + int leftInt = MemoryUtil.UNSAFE.getInt(lPos); + int rightInt = MemoryUtil.UNSAFE.getInt(right, MemoryUtil.BYTE_ARRAY_BASE_OFFSET + rPos); + if (leftInt != rightInt) { + if (LITTLE_ENDIAN) { + return unsignedIntCompare(Integer.reverseBytes(leftInt), Integer.reverseBytes(rightInt)); + } else { + return unsignedIntCompare(leftInt, rightInt); + } + } + lPos += 4; + rPos += 4; + n -= 4; + } + + while (n-- != 0) { + byte leftByte = MemoryUtil.UNSAFE.getByte(lPos); + byte rightByte = right[rPos]; + if (leftByte != rightByte) { + return ((leftByte & 0xFF) - (rightByte & 0xFF)) > 0 ? 1 : -1; + } + lPos++; + rPos++; + } + + if (lLen == rLen) { + return 0; + } + + return lLen > rLen ? 1 : -1; + } + + /** + * Compute hashCode with the given {@link ArrowBuf} and start/end index. + */ + public static int hash(final ArrowBuf buf, long start, long end) { + + return hash(SimpleHasher.INSTANCE, buf, start, end); + } + + /** + * Compute hashCode with the given {@link ArrowBufHasher}, {@link ArrowBuf} and start/end index. + */ + public static final int hash(ArrowBufHasher hasher, final ArrowBuf buf, long start, long end) { + + if (hasher == null) { + hasher = SimpleHasher.INSTANCE; + } + + return hasher.hashCode(buf, start, end - start); + } + + /** + * Generate a new hashCode with the given current hashCode and new hashCode. + */ + public static int combineHash(int currentHash, int newHash) { + return currentHash * 31 + newHash; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/CommonUtil.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/CommonUtil.java new file mode 100644 index 000000000..ccca7b1e0 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/CommonUtil.java @@ -0,0 +1,79 @@ +/* + * 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.util; + +import java.util.Arrays; + +/** + * Utilities and static methods needed for arrow-memory. + */ +public final class CommonUtil { + + private CommonUtil() { } + + /** + * Rounds up the provided value to the nearest power of two. + * + * @param val An integer value. + * @return The closest power of two of that value. + */ + public static int nextPowerOfTwo(int val) { + if (val == 0 || val == 1) { + return val + 1; + } + int highestBit = Integer.highestOneBit(val); + if (highestBit == val) { + return val; + } else { + return highestBit << 1; + } + } + + /** + * Rounds up the provided value to the nearest power of two. + * + * @param val A long value. + * @return The closest power of two of that value. + */ + public static long nextPowerOfTwo(long val) { + if (val == 0 || val == 1) { + return val + 1; + } + long highestBit = Long.highestOneBit(val); + if (highestBit == val) { + return val; + } else { + return highestBit << 1; + } + } + + /** + * Specify an indentation amount when using a StringBuilder. + * + * @param sb StringBuilder to use + * @param indent Indentation amount + * @return the StringBuilder object with indentation applied + */ + public static StringBuilder indent(StringBuilder sb, int indent) { + final char[] indentation = new char[indent * 2]; + Arrays.fill(indentation, ' '); + sb.append(indentation); + return sb; + } +} + diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/HistoricalLog.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/HistoricalLog.java new file mode 100644 index 000000000..f02539a8a --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/HistoricalLog.java @@ -0,0 +1,178 @@ +/* + * 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.util; + +import java.util.Arrays; +import java.util.LinkedList; + +import org.slf4j.Logger; + +/** + * Utility class that can be used to log activity within a class + * for later logging and debugging. Supports recording events and + * recording the stack at the time they occur. + */ +public class HistoricalLog { + + private final LinkedList<Event> history = new LinkedList<>(); + private final String idString; // the formatted id string + private final int limit; // the limit on the number of events kept + private Event firstEvent; // the first stack trace recorded + + /** + * Constructor. The format string will be formatted and have its arguments + * substituted at the time this is called. + * + * @param idStringFormat {@link String#format} format string that can be used to identify this + * object in a log. Including some kind of unique identifier that can be + * associated with the object instance is best. + * @param args for the format string, or nothing if none are required + */ + public HistoricalLog(final String idStringFormat, Object... args) { + this(Integer.MAX_VALUE, idStringFormat, args); + } + + /** + * Constructor. The format string will be formatted and have its arguments + * substituted at the time this is called. + * + * <p>This form supports the specification of a limit that will limit the + * number of historical entries kept (which keeps down the amount of memory + * used). With the limit, the first entry made is always kept (under the + * assumption that this is the creation site of the object, which is usually + * interesting), and then up to the limit number of entries are kept after that. + * Each time a new entry is made, the oldest that is not the first is dropped. + * + * @param limit the maximum number of historical entries that will be kept, not including + * the first entry made + * @param idStringFormat {@link String#format} format string that can be used to identify this + * object in a log. Including some kind of unique identifier that can be + * associated with the object instance is best. + * @param args for the format string, or nothing if none are required + */ + public HistoricalLog(final int limit, final String idStringFormat, Object... args) { + this.limit = limit; + this.idString = String.format(idStringFormat, args); + } + + /** + * Record an event. Automatically captures the stack trace at the time this is + * called. The format string will be formatted and have its arguments substituted + * at the time this is called. + * + * @param noteFormat {@link String#format} format string that describes the event + * @param args for the format string, or nothing if none are required + */ + public synchronized void recordEvent(final String noteFormat, Object... args) { + final String note = String.format(noteFormat, args); + final Event event = new Event(note); + if (firstEvent == null) { + firstEvent = event; + } + if (history.size() == limit) { + history.removeFirst(); + } + history.add(event); + } + + /** + * Write the history of this object to the given {@link StringBuilder}. The history + * includes the identifying string provided at construction time, and all the recorded + * events with their stack traces. + * + * @param sb {@link StringBuilder} to write to + * @param includeStackTrace whether to include the stacktrace of each event in the history + */ + public void buildHistory(final StringBuilder sb, boolean includeStackTrace) { + buildHistory(sb, 0, includeStackTrace); + } + + /** + * Build the history and write it to sb. + * + * @param sb output + * @param indent starting indent (usually "") + * @param includeStackTrace whether to include the stacktrace of each event. + */ + public synchronized void buildHistory( + final StringBuilder sb, int indent, boolean includeStackTrace) { + final char[] indentation = new char[indent]; + final char[] innerIndentation = new char[indent + 2]; + Arrays.fill(indentation, ' '); + Arrays.fill(innerIndentation, ' '); + + sb.append(indentation) + .append("event log for: ") + .append(idString) + .append('\n'); + + if (firstEvent != null) { + sb.append(innerIndentation) + .append(firstEvent.time) + .append(' ') + .append(firstEvent.note) + .append('\n'); + if (includeStackTrace) { + firstEvent.stackTrace.writeToBuilder(sb, indent + 2); + } + + for (final Event event : history) { + if (event == firstEvent) { + continue; + } + sb.append(innerIndentation) + .append(" ") + .append(event.time) + .append(' ') + .append(event.note) + .append('\n'); + + if (includeStackTrace) { + event.stackTrace.writeToBuilder(sb, indent + 2); + sb.append('\n'); + } + } + } + } + + /** + * Write the history of this object to the given {@link Logger}. The history + * includes the identifying string provided at construction time, and all the recorded + * events with their stack traces. + * + * @param logger {@link Logger} to write to + */ + public void logHistory(final Logger logger) { + final StringBuilder sb = new StringBuilder(); + buildHistory(sb, 0, true); + logger.debug(sb.toString()); + } + + private static class Event { + + private final String note; // the event text + private final StackTrace stackTrace; // where the event occurred + private final long time; + + public Event(final String note) { + this.note = note; + this.time = System.nanoTime(); + stackTrace = new StackTrace(); + } + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/LargeMemoryUtil.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/LargeMemoryUtil.java new file mode 100644 index 000000000..db63bbd14 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/LargeMemoryUtil.java @@ -0,0 +1,44 @@ +/* + * 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.util; + +import org.apache.arrow.memory.BoundsChecking; + +/** Contains utilities for dealing with a 64-bit address base. */ +public final class LargeMemoryUtil { + + private LargeMemoryUtil() {} + + /** + * Casts length to an int, but raises an exception the value is outside + * the range of an int. + */ + public static int checkedCastToInt(long length) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + return Math.toIntExact(length); + } + return (int) length; + } + + /** + * Returns a min(Integer.MAX_VALUE, length). + */ + public static int capAtMaxInt(long length) { + return (int) Math.min(length, Integer.MAX_VALUE); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java new file mode 100644 index 000000000..16ef39702 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/MemoryUtil.java @@ -0,0 +1,170 @@ +/* + * 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.util; + +import java.lang.reflect.Constructor; +import java.lang.reflect.Field; +import java.lang.reflect.InvocationTargetException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.security.AccessController; +import java.security.PrivilegedAction; + +import sun.misc.Unsafe; + +/** + * Utilities for memory related operations. + */ +public class MemoryUtil { + private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(MemoryUtil.class); + + private static final Constructor<?> DIRECT_BUFFER_CONSTRUCTOR; + /** + * The unsafe object from which to access the off-heap memory. + */ + public static final Unsafe UNSAFE; + + /** + * The start offset of array data relative to the start address of the array object. + */ + public static final long BYTE_ARRAY_BASE_OFFSET; + + /** + * The offset of the address field with the {@link java.nio.ByteBuffer} object. + */ + static final long BYTE_BUFFER_ADDRESS_OFFSET; + + /** + * If the native byte order is little-endian. + */ + public static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; + + static { + try { + // try to get the unsafe object + final Object maybeUnsafe = AccessController.doPrivileged(new PrivilegedAction<Object>() { + @Override + public Object run() { + try { + final Field unsafeField = Unsafe.class.getDeclaredField("theUnsafe"); + unsafeField.setAccessible(true); + return unsafeField.get(null); + } catch (Throwable e) { + return e; + } + } + }); + + if (maybeUnsafe instanceof Throwable) { + throw (Throwable) maybeUnsafe; + } + + UNSAFE = (Unsafe) maybeUnsafe; + + // get the offset of the data inside a byte array object + BYTE_ARRAY_BASE_OFFSET = UNSAFE.arrayBaseOffset(byte[].class); + + // get the offset of the address field in a java.nio.Buffer object + Field addressField = java.nio.Buffer.class.getDeclaredField("address"); + addressField.setAccessible(true); + BYTE_BUFFER_ADDRESS_OFFSET = UNSAFE.objectFieldOffset(addressField); + + Constructor<?> directBufferConstructor; + long address = -1; + final ByteBuffer direct = ByteBuffer.allocateDirect(1); + try { + + final Object maybeDirectBufferConstructor = + AccessController.doPrivileged(new PrivilegedAction<Object>() { + @Override + public Object run() { + try { + final Constructor<?> constructor = + direct.getClass().getDeclaredConstructor(long.class, int.class); + constructor.setAccessible(true); + logger.debug("Constructor for direct buffer found and made accessible"); + return constructor; + } catch (NoSuchMethodException e) { + logger.debug("Cannot get constructor for direct buffer allocation", e); + return e; + } catch (SecurityException e) { + logger.debug("Cannot get constructor for direct buffer allocation", e); + return e; + } + } + }); + + if (maybeDirectBufferConstructor instanceof Constructor<?>) { + address = UNSAFE.allocateMemory(1); + // try to use the constructor now + try { + ((Constructor<?>) maybeDirectBufferConstructor).newInstance(address, 1); + directBufferConstructor = (Constructor<?>) maybeDirectBufferConstructor; + logger.debug("direct buffer constructor: available"); + } catch (InstantiationException | IllegalAccessException | InvocationTargetException e) { + logger.warn("unable to instantiate a direct buffer via constructor", e); + directBufferConstructor = null; + } + } else { + logger.debug( + "direct buffer constructor: unavailable", + (Throwable) maybeDirectBufferConstructor); + directBufferConstructor = null; + } + } finally { + if (address != -1) { + UNSAFE.freeMemory(address); + } + } + DIRECT_BUFFER_CONSTRUCTOR = directBufferConstructor; + } catch (Throwable e) { + throw new RuntimeException("Failed to initialize MemoryUtil.", e); + } + } + + /** + * Given a {@link ByteBuffer}, gets the address the underlying memory space. + * + * @param buf the byte buffer. + * @return address of the underlying memory. + */ + public static long getByteBufferAddress(ByteBuffer buf) { + return UNSAFE.getLong(buf, BYTE_BUFFER_ADDRESS_OFFSET); + } + + private MemoryUtil() { + } + + /** + * Create nio byte buffer. + */ + public static ByteBuffer directBuffer(long address, int capacity) { + if (DIRECT_BUFFER_CONSTRUCTOR != null) { + if (capacity < 0) { + throw new IllegalArgumentException("Capacity is negative, has to be positive or 0"); + } + try { + return (ByteBuffer) DIRECT_BUFFER_CONSTRUCTOR.newInstance(address, capacity); + } catch (Throwable cause) { + throw new Error(cause); + } + } + throw new UnsupportedOperationException( + "sun.misc.Unsafe or java.nio.DirectByteBuffer.<init>(long, int) not available"); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/StackTrace.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/StackTrace.java new file mode 100644 index 000000000..d743af86c --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/StackTrace.java @@ -0,0 +1,70 @@ +/* + * 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.util; + +import java.util.Arrays; + +/** + * Convenient way of obtaining and manipulating stack traces for debugging. + */ +public class StackTrace { + + private final StackTraceElement[] stackTraceElements; + + /** + * Constructor. Captures the current stack trace. + */ + public StackTrace() { + // skip over the first element so that we don't include this constructor call + final StackTraceElement[] stack = Thread.currentThread().getStackTrace(); + stackTraceElements = Arrays.copyOfRange(stack, 1, stack.length - 1); + } + + /** + * Write the stack trace to a StringBuilder. + * + * @param sb where to write it + * @param indent how many double spaces to indent each line + */ + public void writeToBuilder(final StringBuilder sb, final int indent) { + // create the indentation string + final char[] indentation = new char[indent * 2]; + Arrays.fill(indentation, ' '); + + // write the stack trace in standard Java format + for (StackTraceElement ste : stackTraceElements) { + sb.append(indentation) + .append("at ") + .append(ste.getClassName()) + .append('.') + .append(ste.getMethodName()) + .append('(') + .append(ste.getFileName()) + .append(':') + .append(Integer.toString(ste.getLineNumber())) + .append(")\n"); + } + } + + @Override + public String toString() { + final StringBuilder sb = new StringBuilder(); + writeToBuilder(sb, 0); + return sb.toString(); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/ArrowBufHasher.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/ArrowBufHasher.java new file mode 100644 index 000000000..0de8e62a4 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/ArrowBufHasher.java @@ -0,0 +1,47 @@ +/* + * 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.util.hash; + +import org.apache.arrow.memory.ArrowBuf; + +/** + * Utility for calculating the hash code for a consecutive memory region. + * This class provides the basic framework for efficiently calculating the hash code. + * <p> + * A default light-weight implementation is given in {@link SimpleHasher}. + * </p> + */ +public interface ArrowBufHasher { + + /** + * Calculates the hash code for a memory region. + * @param address start address of the memory region. + * @param length length of the memory region. + * @return the hash code. + */ + int hashCode(long address, long length); + + /** + * Calculates the hash code for a memory region. + * @param buf the buffer for the memory region. + * @param offset offset within the buffer for the memory region. + * @param length length of the memory region. + * @return the hash code. + */ + int hashCode(ArrowBuf buf, long offset, long length); +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/MurmurHasher.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/MurmurHasher.java new file mode 100644 index 000000000..ea565dfca --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/MurmurHasher.java @@ -0,0 +1,175 @@ +/* + * 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.util.hash; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.util.MemoryUtil; + +/** + * Implementation of the Murmur hashing algorithm. + * Details of the algorithm can be found in + * https://en.wikipedia.org/wiki/MurmurHash + * <p> + * Murmur hashing is computationally expensive, as it involves several + * integer multiplications. However, the produced hash codes have + * good quality in the sense that they are uniformly distributed in the universe. + * </p> + * <p> + * Therefore, this algorithm is suitable for scenarios where uniform hashing + * is desired (e.g. in an open addressing hash table/hash set). + * </p> + */ +public class MurmurHasher implements ArrowBufHasher { + + private final int seed; + + /** + * Creates a default Murmur hasher, with seed 0. + */ + public MurmurHasher() { + this(0); + } + + /** + * Creates a Murmur hasher. + * @param seed the seed for the hasher. + */ + public MurmurHasher(int seed) { + this.seed = seed; + } + + @Override + public int hashCode(long address, long length) { + return hashCode(address, length, seed); + } + + @Override + public int hashCode(ArrowBuf buf, long offset, long length) { + buf.checkBytes(offset, offset + length); + return hashCode(buf.memoryAddress() + offset, length); + } + + /** + * Calculates the hash code for a memory region. + * @param buf the buffer for the memory region. + * @param offset offset within the buffer for the memory region. + * @param length length of the memory region. + * @param seed the seed. + * @return the hash code. + */ + public static int hashCode(ArrowBuf buf, long offset, long length, int seed) { + buf.checkBytes(offset, offset + length); + return hashCode(buf.memoryAddress() + offset, length, seed); + } + + /** + * Calculates the hash code for a memory region. + * @param address start address of the memory region. + * @param length length of the memory region. + * @param seed the seed. + * @return the hash code. + */ + public static int hashCode(long address, long length, int seed) { + int index = 0; + int hash = seed; + while (index + 4 <= length) { + int intValue = MemoryUtil.UNSAFE.getInt(address + index); + hash = combineHashCode(hash, intValue); + index += 4; + } + + if (index < length) { + // process remaining data as a integer in little endian + int intValue = 0; + for (int i = index - 1; i >= index; i--) { + intValue <<= 8; + intValue |= (MemoryUtil.UNSAFE.getByte(address + i) & 0x000000ff); + index += 1; + } + hash = combineHashCode(hash, intValue); + } + return finalizeHashCode(hash, length); + } + + /** + * Combine the current hash code and a new int value to calculate + * a new hash code. + * @param currentHashCode the current hash code. + * @param intValue the new int value. + * @return the new hah code. + */ + public static int combineHashCode(int currentHashCode, int intValue) { + int c1 = 0xcc9e2d51; + int c2 = 0x1b873593; + int r1 = 15; + int r2 = 13; + int m = 5; + int n = 0xe6546b64; + + int k = intValue; + k = k * c1; + k = rotateLeft(k, r1); + k = k * c2; + + int hash = currentHashCode; + hash = hash ^ k; + hash = rotateLeft(hash, r2); + hash = hash * m + n; + + return hash; + } + + /** + * Finalizing the hash code. + * @param hashCode the current hash code. + * @param length the length of the memory region. + * @return the finalized hash code. + */ + public static int finalizeHashCode(int hashCode, long length) { + hashCode = hashCode ^ (int) length; + + hashCode = hashCode ^ (hashCode >>> 16); + hashCode = hashCode * 0x85ebca6b; + hashCode = hashCode ^ (hashCode >>> 13); + hashCode = hashCode * 0xc2b2ae35; + hashCode = hashCode ^ (hashCode >>> 16); + + return hashCode; + } + + private static int rotateLeft(int value, int count) { + return (value << count) | (value >>> (32 - count)); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + MurmurHasher that = (MurmurHasher) o; + return seed == that.seed; + } + + @Override + public int hashCode() { + return seed; + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/SimpleHasher.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/SimpleHasher.java new file mode 100644 index 000000000..da0ee4829 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/util/hash/SimpleHasher.java @@ -0,0 +1,116 @@ +/* + * 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.util.hash; + + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.util.MemoryUtil; + +/** + * A simple hasher that calculates the hash code of integers as is, + * and does not perform any finalization. So the computation is extremely + * efficient. + * <p> + * This algorithm only provides the most basic semantics for the hash code. That is, + * if two objects are equal, they must have equal hash code. However, the quality of the + * produced hash code may not be good. In other words, the generated hash codes are + * far from being uniformly distributed in the universe. + * </p> + * <p> + * Therefore, this algorithm is suitable only for scenarios where the most basic semantics + * of the hash code is required (e.g. in scenarios that require fast and proactive data pruning) + * </p> + * <p> + * An object of this class is stateless, so it can be shared between threads. + * </p> + */ +public class SimpleHasher implements ArrowBufHasher { + + public static SimpleHasher INSTANCE = new SimpleHasher(); + + protected SimpleHasher() { + } + + /** + * Calculates the hash code for a memory region. + * @param address start address of the memory region. + * @param length length of the memory region. + * @return the hash code. + */ + public int hashCode(long address, long length) { + int hashValue = 0; + int index = 0; + while (index + 8 <= length) { + long longValue = MemoryUtil.UNSAFE.getLong(address + index); + int longHash = getLongHashCode(longValue); + hashValue = combineHashCode(hashValue, longHash); + index += 8; + } + + if (index + 4 <= length) { + int intValue = MemoryUtil.UNSAFE.getInt(address + index); + int intHash = intValue; + hashValue = combineHashCode(hashValue, intHash); + index += 4; + } + + while (index < length) { + byte byteValue = MemoryUtil.UNSAFE.getByte(address + index); + int byteHash = byteValue; + hashValue = combineHashCode(hashValue, byteHash); + index += 1; + } + + return finalizeHashCode(hashValue); + } + + /** + * Calculates the hash code for a memory region. + * @param buf the buffer for the memory region. + * @param offset offset within the buffer for the memory region. + * @param length length of the memory region. + * @return the hash code. + */ + @Override + public int hashCode(ArrowBuf buf, long offset, long length) { + buf.checkBytes(offset, offset + length); + return hashCode(buf.memoryAddress() + offset, length); + } + + protected int combineHashCode(int currentHashCode, int newHashCode) { + return currentHashCode * 37 + newHashCode; + } + + protected int getLongHashCode(long longValue) { + return Long.hashCode(longValue); + } + + protected int finalizeHashCode(int hashCode) { + return hashCode; + } + + @Override + public int hashCode() { + return 123; + } + + @Override + public boolean equals(Object obj) { + return obj != null && (obj instanceof SimpleHasher); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java new file mode 100644 index 000000000..d965f2aed --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/AutoCloseables.java @@ -0,0 +1,242 @@ +/** + * 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.util; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.stream.StreamSupport; + +/** + * Utilities for AutoCloseable classes. + */ +public final class AutoCloseables { + // Utility class. Should not be instantiated + private AutoCloseables() { + } + + /** + * Returns a new {@link AutoCloseable} that calls {@link #close(Iterable)} on <code>autoCloseables</code> + * when close is called. + */ + public static AutoCloseable all(final Collection<? extends AutoCloseable> autoCloseables) { + return new AutoCloseable() { + @Override + public void close() throws Exception { + AutoCloseables.close(autoCloseables); + } + }; + } + + /** + * Closes all autoCloseables if not null and suppresses exceptions by adding them to t. + * @param t the throwable to add suppressed exception to + * @param autoCloseables the closeables to close + */ + public static void close(Throwable t, AutoCloseable... autoCloseables) { + close(t, Arrays.asList(autoCloseables)); + } + + /** + * Closes all autoCloseables if not null and suppresses exceptions by adding them to t. + * @param t the throwable to add suppressed exception to + * @param autoCloseables the closeables to close + */ + public static void close(Throwable t, Iterable<? extends AutoCloseable> autoCloseables) { + try { + close(autoCloseables); + } catch (Exception e) { + t.addSuppressed(e); + } + } + + /** + * Closes all autoCloseables if not null and suppresses subsequent exceptions if more than one. + * @param autoCloseables the closeables to close + */ + public static void close(AutoCloseable... autoCloseables) throws Exception { + close(Arrays.asList(autoCloseables)); + } + + /** + * Closes all autoCloseables if not null and suppresses subsequent exceptions if more than one. + * @param ac the closeables to close + */ + public static void close(Iterable<? extends AutoCloseable> ac) throws Exception { + // this method can be called on a single object if it implements Iterable<AutoCloseable> + // like for example VectorContainer make sure we handle that properly + if (ac == null) { + return; + } else if (ac instanceof AutoCloseable) { + ((AutoCloseable) ac).close(); + return; + } + + Exception topLevelException = null; + for (AutoCloseable closeable : ac) { + try { + if (closeable != null) { + closeable.close(); + } + } catch (Exception e) { + if (topLevelException == null) { + topLevelException = e; + } else if (e != topLevelException) { + topLevelException.addSuppressed(e); + } + } + } + if (topLevelException != null) { + throw topLevelException; + } + } + + /** + * Calls {@link #close(Iterable)} on the flattened list of closeables. + */ + @SafeVarargs + public static void close(Iterable<? extends AutoCloseable>...closeables) throws Exception { + close(flatten(closeables)); + } + + @SafeVarargs + private static Iterable<AutoCloseable> flatten(Iterable<? extends AutoCloseable>... closeables) { + return new Iterable<AutoCloseable>() { + // Cast from Iterable<? extends AutoCloseable> to Iterable<AutoCloseable> is safe in this context + // since there's no modification of the original collection + @SuppressWarnings("unchecked") + @Override + public Iterator<AutoCloseable> iterator() { + return Arrays.stream(closeables) + .flatMap((Iterable<? extends AutoCloseable> i) + -> StreamSupport.stream(((Iterable<AutoCloseable>) i).spliterator(), /*parallel=*/false)) + .iterator(); + } + }; + } + + /** + * Converts <code>ac</code> to a {@link Iterable} filtering out any null values. + */ + public static Iterable<AutoCloseable> iter(AutoCloseable... ac) { + if (ac.length == 0) { + return Collections.emptyList(); + } else { + final List<AutoCloseable> nonNullAc = new ArrayList<>(); + for (AutoCloseable autoCloseable : ac) { + if (autoCloseable != null) { + nonNullAc.add(autoCloseable); + } + } + return nonNullAc; + } + } + + /** + * A closeable wrapper that will close the underlying closeables if a commit does not occur. + */ + public static class RollbackCloseable implements AutoCloseable { + + private boolean commit = false; + private List<AutoCloseable> closeables; + + public RollbackCloseable(AutoCloseable... closeables) { + this.closeables = new ArrayList<>(Arrays.asList(closeables)); + } + + public <T extends AutoCloseable> T add(T t) { + closeables.add(t); + return t; + } + + /** + * Add all of <code>list</code> to the rollback list. + */ + public void addAll(AutoCloseable... list) { + closeables.addAll(Arrays.asList(list)); + } + + /** + * Add all of <code>list</code> to the rollback list. + */ + public void addAll(Iterable<? extends AutoCloseable> list) { + for (AutoCloseable ac : list) { + closeables.add(ac); + } + } + + public void commit() { + commit = true; + } + + @Override + public void close() throws Exception { + if (!commit) { + AutoCloseables.close(closeables); + } + } + + } + + /** + * Creates an {@link RollbackCloseable} from the given closeables. + */ + public static RollbackCloseable rollbackable(AutoCloseable... closeables) { + return new RollbackCloseable(closeables); + } + + /** + * close() an {@link java.lang.AutoCloseable} without throwing a (checked) + * {@link java.lang.Exception}. This wraps the close() call with a + * try-catch that will rethrow an Exception wrapped with a + * {@link java.lang.RuntimeException}, providing a way to call close() + * without having to do the try-catch everywhere or propagate the Exception. + * + * @param autoCloseable the AutoCloseable to close; may be null + * @throws RuntimeException if an Exception occurs; the Exception is + * wrapped by the RuntimeException + */ + public static void closeNoChecked(final AutoCloseable autoCloseable) { + if (autoCloseable != null) { + try { + autoCloseable.close(); + } catch (final Exception e) { + throw new RuntimeException("Exception while closing: " + e.getMessage(), e); + } + } + } + + private static final AutoCloseable noOpAutocloseable = new AutoCloseable() { + @Override + public void close() { + } + }; + + /** + * Get an AutoCloseable that does nothing. + * + * @return A do-nothing autocloseable + */ + public static AutoCloseable noop() { + return noOpAutocloseable; + } +} + diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/Collections2.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/Collections2.java new file mode 100644 index 000000000..6b01a61eb --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/Collections2.java @@ -0,0 +1,92 @@ +/** + * 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.util; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Spliterator; +import java.util.Spliterators; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; + +/** + * Utility methods for manipulating {@link java.util.Collections} and their subclasses/implementations. + */ +public final class Collections2 { + private Collections2() {} + + /** + * Creates a {@link List} from the elements remaining in iterator. + */ + public static <T> List<T> toList(Iterator<T> iterator) { + List<T> target = new ArrayList<>(); + iterator.forEachRemaining(target::add); + return target; + } + + /** + * Converts the iterable into a new {@link List}. + */ + public static <T> List<T> toList(Iterable<T> iterable) { + if (iterable instanceof Collection<?>) { + // If iterable is a collection, take advantage of it for a more efficient copy + return new ArrayList<T>((Collection<T>) iterable); + } + return toList(iterable.iterator()); + } + + /** + * Converts the iterable into a new immutable {@link List}. + */ + public static <T> List<T> toImmutableList(Iterable<T> iterable) { + return Collections.unmodifiableList(toList(iterable)); + } + + + /** Copies the elements of <code>map</code> to a new unmodifiable map. */ + public static <K, V> Map<K, V> immutableMapCopy(Map<K, V> map) { + return Collections.unmodifiableMap(new HashMap<>(map)); + } + + /** Copies the elements of list to a new unmodifiable list. */ + public static <V> List<V> immutableListCopy(List<V> list) { + return Collections.unmodifiableList(new ArrayList<>(list)); + } + + /** Copies the values to a new unmodifiable list. */ + public static <V> List<V> asImmutableList(V...values) { + return Collections.unmodifiableList(Arrays.asList(values)); + } + + /** + * Creates a human readable string from the remaining elements in iterator. + * + * The output should be similar to {@code Arrays#toString(Object[])} + */ + public static String toString(Iterator<?> iterator) { + return StreamSupport.stream(Spliterators.spliteratorUnknownSize(iterator, Spliterator.ORDERED), false) + .map(String::valueOf) + .collect(Collectors.joining(", ", "[", "]")); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java new file mode 100644 index 000000000..0ffc9447e --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/Preconditions.java @@ -0,0 +1,1323 @@ +/* + * 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. + */ + +/* + * Lifted from Guava 20.0 to avoid dependency for core Arrow libraries. + * + * Copyright (C) 2007 The Guava Authors + * + * Licensed 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.util; + +/** + * Static convenience methods that help a method or constructor check whether it was invoked + * correctly (whether its <i>preconditions</i> have been met). These methods generally accept a + * {@code boolean} expression which is expected to be {@code true} (or in the case of {@code + * checkNotNull}, an object reference which is expected to be non-null). When {@code false} (or + * {@code null}) is passed instead, the {@code Preconditions} method throws an unchecked exception, + * which helps the calling method communicate to <i>its</i> caller that <i>that</i> caller has made + * a mistake. Example: <pre> {@code + * + * /** + * * Returns the positive square root of the given value. + * * + * * @throws IllegalArgumentException if the value is negative + * *}{@code / + * public static double sqrt(double value) { + * Preconditions.checkArgument(value >= 0.0, "negative value: %s", value); + * // calculate the square root + * } + * + * void exampleBadCaller() { + * double d = sqrt(-1.0); + * }}</pre> + * + * <p>In this example, {@code checkArgument} throws an {@code IllegalArgumentException} to indicate + * that {@code exampleBadCaller} made an error in <i>its</i> call to {@code sqrt}. + * + * <h3>Warning about performance</h3> + * + * <p>The goal of this class is to improve readability of code, but in some circumstances this may + * come at a significant performance cost. Remember that parameter values for message construction + * must all be computed eagerly, and autoboxing and varargs array creation may happen as well, even + * when the precondition check then succeeds (as it should almost always do in production). In some + * circumstances these wasted CPU cycles and allocations can add up to a real problem. + * Performance-sensitive precondition checks can always be converted to the customary form: + * <pre> {@code + * + * if (value < 0.0) { + * throw new IllegalArgumentException("negative value: " + value); + * }}</pre> + * + * <h3>Other types of preconditions</h3> + * + * <p>Not every type of precondition failure is supported by these methods. Continue to throw + * standard JDK exceptions such as {@link java.util.NoSuchElementException} or + * {@link UnsupportedOperationException} in the situations they are intended for. + * + * <h3>Non-preconditions</h3> + * + * <p>It is of course possible to use the methods of this class to check for invalid conditions + * which are <i>not the caller's fault</i>. Doing so is <b>not recommended</b> because it is + * misleading to future readers of the code and of stack traces. See + * <a href="https://github.com/google/guava/wiki/ConditionalFailuresExplained">Conditional failures + * explained</a> in the Guava User Guide for more advice. + * + * <h3>{@code java.util.Objects.requireNonNull()}</h3> + * + * <p>Projects which use {@code com.google.common} should generally avoid the use of + * {@link java.util.Objects#requireNonNull(Object)}. Instead, use whichever of + * {@link #checkNotNull(Object)} or {@link Verify#verifyNotNull(Object)} is appropriate to the + * situation. (The same goes for the message-accepting overloads.) + * + * <h3>Only {@code %s} is supported</h3> + * + * <p>In {@code Preconditions} error message template strings, only the {@code "%s"} specifier is + * supported, not the full range of {@link java.util.Formatter} specifiers. + * + * <h3>More information</h3> + * + * <p>See the Guava User Guide on + * <a href="https://github.com/google/guava/wiki/PreconditionsExplained">using {@code + * Preconditions}</a>. + * + * @author Kevin Bourrillion + * @since 2.0 + */ +public final class Preconditions { + private Preconditions() {} + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * @param expression a boolean expression + * @throws IllegalArgumentException if {@code expression} is false + */ + public static void checkArgument(boolean expression) { + if (!expression) { + throw new IllegalArgumentException(); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * @param expression a boolean expression + * @param errorMessage the exception message to use if the check fails; will be converted to a + * string using {@link String#valueOf(Object)} + * @throws IllegalArgumentException if {@code expression} is false + */ + public static void checkArgument(boolean expression, Object errorMessage) { + if (!expression) { + throw new IllegalArgumentException(String.valueOf(errorMessage)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * @param expression a boolean expression + * @param errorMessageTemplate a template for the exception message should the check fail. The + * message is formed by replacing each {@code %s} placeholder in the template with an + * argument. These are matched by position - the first {@code %s} gets {@code + * errorMessageArgs[0]}, etc. Unmatched arguments will be appended to the formatted message in + * square braces. Unmatched placeholders will be left as-is. + * @param errorMessageArgs the arguments to be substituted into the message template. Arguments + * are converted to strings using {@link String#valueOf(Object)}. + * @throws IllegalArgumentException if {@code expression} is false + * @throws NullPointerException if the check fails and either {@code errorMessageTemplate} or + * {@code errorMessageArgs} is null (don't let this happen) + */ + public static void checkArgument( + boolean expression, + String errorMessageTemplate, + Object... errorMessageArgs) { + if (!expression) { + throw new IllegalArgumentException(format(errorMessageTemplate, errorMessageArgs)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument(boolean b, String errorMessageTemplate, char p1) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument(boolean b, String errorMessageTemplate, int p1) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument(boolean b, String errorMessageTemplate, long p1) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, Object p1) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, char p1, char p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, char p1, int p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, char p1, long p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, char p1, Object p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, int p1, char p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, int p1, int p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, int p1, long p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, int p1, Object p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, long p1, char p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, long p1, int p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, long p1, long p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, long p1, Object p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, Object p1, char p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, Object p1, int p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, Object p1, long p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, String errorMessageTemplate, Object p1, Object p2) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, + String errorMessageTemplate, + Object p1, + Object p2, + Object p3) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2, p3)); + } + } + + /** + * Ensures the truth of an expression involving one or more parameters to the calling method. + * + * <p>See {@link #checkArgument(boolean, String, Object...)} for details. + */ + public static void checkArgument( + boolean b, + String errorMessageTemplate, + Object p1, + Object p2, + Object p3, + Object p4) { + if (!b) { + throw new IllegalArgumentException(format(errorMessageTemplate, p1, p2, p3, p4)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * @param expression a boolean expression + * @throws IllegalStateException if {@code expression} is false + */ + public static void checkState(boolean expression) { + if (!expression) { + throw new IllegalStateException(); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * @param expression a boolean expression + * @param errorMessage the exception message to use if the check fails; will be converted to a + * string using {@link String#valueOf(Object)} + * @throws IllegalStateException if {@code expression} is false + */ + public static void checkState(boolean expression, Object errorMessage) { + if (!expression) { + throw new IllegalStateException(String.valueOf(errorMessage)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * @param expression a boolean expression + * @param errorMessageTemplate a template for the exception message should the check fail. The + * message is formed by replacing each {@code %s} placeholder in the template with an + * argument. These are matched by position - the first {@code %s} gets {@code + * errorMessageArgs[0]}, etc. Unmatched arguments will be appended to the formatted message in + * square braces. Unmatched placeholders will be left as-is. + * @param errorMessageArgs the arguments to be substituted into the message template. Arguments + * are converted to strings using {@link String#valueOf(Object)}. + * @throws IllegalStateException if {@code expression} is false + * @throws NullPointerException if the check fails and either {@code errorMessageTemplate} or + * {@code errorMessageArgs} is null (don't let this happen) + */ + public static void checkState( + boolean expression, + String errorMessageTemplate, + Object... errorMessageArgs) { + if (!expression) { + throw new IllegalStateException(format(errorMessageTemplate, errorMessageArgs)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, char p1) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, int p1) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, long p1) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, Object p1) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, char p1, char p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, char p1, int p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, char p1, long p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, char p1, Object p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, int p1, char p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, int p1, int p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, int p1, long p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, int p1, Object p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, long p1, char p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState(boolean b, String errorMessageTemplate, long p1, int p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, long p1, long p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, long p1, Object p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, Object p1, char p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, Object p1, int p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, Object p1, long p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, String errorMessageTemplate, Object p1, Object p2) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, + String errorMessageTemplate, + Object p1, + Object p2, + Object p3) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2, p3)); + } + } + + /** + * Ensures the truth of an expression involving the state of the calling instance, but not + * involving any parameters to the calling method. + * + * <p>See {@link #checkState(boolean, String, Object...)} for details. + */ + public static void checkState( + boolean b, + String errorMessageTemplate, + Object p1, + Object p2, + Object p3, + Object p4) { + if (!b) { + throw new IllegalStateException(format(errorMessageTemplate, p1, p2, p3, p4)); + } + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * @param reference an object reference + * @return the non-null reference that was validated + * @throws NullPointerException if {@code reference} is null + */ + + public static <T> T checkNotNull(T reference) { + if (reference == null) { + throw new NullPointerException(); + } + return reference; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * @param reference an object reference + * @param errorMessage the exception message to use if the check fails; will be converted to a + * string using {@link String#valueOf(Object)} + * @return the non-null reference that was validated + * @throws NullPointerException if {@code reference} is null + */ + + public static <T> T checkNotNull(T reference, Object errorMessage) { + if (reference == null) { + throw new NullPointerException(String.valueOf(errorMessage)); + } + return reference; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * @param reference an object reference + * @param errorMessageTemplate a template for the exception message should the check fail. The + * message is formed by replacing each {@code %s} placeholder in the template with an + * argument. These are matched by position - the first {@code %s} gets {@code + * errorMessageArgs[0]}, etc. Unmatched arguments will be appended to the formatted message in + * square braces. Unmatched placeholders will be left as-is. + * @param errorMessageArgs the arguments to be substituted into the message template. Arguments + * are converted to strings using {@link String#valueOf(Object)}. + * @return the non-null reference that was validated + * @throws NullPointerException if {@code reference} is null + */ + + public static <T> T checkNotNull( + T reference, String errorMessageTemplate, Object... errorMessageArgs) { + if (reference == null) { + // If either of these parameters is null, the right thing happens anyway + throw new NullPointerException(format(errorMessageTemplate, errorMessageArgs)); + } + return reference; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, char p1) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, int p1) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, long p1) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, Object p1) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, char p1, char p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, char p1, int p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, char p1, long p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, char p1, Object p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, int p1, char p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, int p1, int p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, int p1, long p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, int p1, Object p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, long p1, char p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, long p1, int p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull(T obj, String errorMessageTemplate, long p1, long p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, long p1, Object p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, Object p1, char p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, Object p1, int p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, Object p1, long p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, String errorMessageTemplate, Object p1, Object p2) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, + String errorMessageTemplate, + Object p1, + Object p2, + Object p3) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2, p3)); + } + return obj; + } + + /** + * Ensures that an object reference passed as a parameter to the calling method is not null. + * + * <p>See {@link #checkNotNull(Object, String, Object...)} for details. + */ + + public static <T> T checkNotNull( + T obj, + String errorMessageTemplate, + Object p1, + Object p2, + Object p3, + Object p4) { + if (obj == null) { + throw new NullPointerException(format(errorMessageTemplate, p1, p2, p3, p4)); + } + return obj; + } + + /* + * All recent hotspots (as of 2009) *really* like to have the natural code + * + * if (guardExpression) { + * throw new BadException(messageExpression); + * } + * + * refactored so that messageExpression is moved to a separate String-returning method. + * + * if (guardExpression) { + * throw new BadException(badMsg(...)); + * } + * + * The alternative natural refactorings into void or Exception-returning methods are much slower. + * This is a big deal - we're talking factors of 2-8 in microbenchmarks, not just 10-20%. (This is + * a hotspot optimizer bug, which should be fixed, but that's a separate, big project). + * + * The coding pattern above is heavily used in java.util, e.g. in ArrayList. There is a + * RangeCheckMicroBenchmark in the JDK that was used to test this. + * + * But the methods in this class want to throw different exceptions, depending on the args, so it + * appears that this pattern is not directly applicable. But we can use the ridiculous, devious + * trick of throwing an exception in the middle of the construction of another exception. Hotspot + * is fine with that. + */ + + /** + * Ensures that {@code index} specifies a valid <i>element</i> in an array, list or string of size + * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive. + * + * @param index a user-supplied index identifying an element of an array, list or string + * @param size the size of that array, list or string + * @return the value of {@code index} + * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size} + * @throws IllegalArgumentException if {@code size} is negative + */ + + public static int checkElementIndex(int index, int size) { + return checkElementIndex(index, size, "index"); + } + + /** + * Ensures that {@code index} specifies a valid <i>element</i> in an array, list or string of size + * {@code size}. An element index may range from zero, inclusive, to {@code size}, exclusive. + * + * @param index a user-supplied index identifying an element of an array, list or string + * @param size the size of that array, list or string + * @param desc the text to use to describe this index in an error message + * @return the value of {@code index} + * @throws IndexOutOfBoundsException if {@code index} is negative or is not less than {@code size} + * @throws IllegalArgumentException if {@code size} is negative + */ + + public static int checkElementIndex(int index, int size, String desc) { + // Carefully optimized for execution by hotspot (explanatory comment above) + if (index < 0 || index >= size) { + throw new IndexOutOfBoundsException(badElementIndex(index, size, desc)); + } + return index; + } + + private static String badElementIndex(int index, int size, String desc) { + if (index < 0) { + return format("%s (%s) must not be negative", desc, index); + } else if (size < 0) { + throw new IllegalArgumentException("negative size: " + size); + } else { // index >= size + return format("%s (%s) must be less than size (%s)", desc, index, size); + } + } + + /** + * Ensures that {@code index} specifies a valid <i>position</i> in an array, list or string of + * size {@code size}. A position index may range from zero to {@code size}, inclusive. + * + * @param index a user-supplied index identifying a position in an array, list or string + * @param size the size of that array, list or string + * @return the value of {@code index} + * @throws IndexOutOfBoundsException if {@code index} is negative or is greater than {@code size} + * @throws IllegalArgumentException if {@code size} is negative + */ + + public static long checkPositionIndex(long index, long size) { + return checkPositionIndex(index, size, "index"); + } + + /** + * Ensures that {@code index} specifies a valid <i>position</i> in an array, list or string of + * size {@code size}. A position index may range from zero to {@code size}, inclusive. + * + * @param index a user-supplied index identifying a position in an array, list or string + * @param size the size of that array, list or string + * @param desc the text to use to describe this index in an error message + * @return the value of {@code index} + * @throws IndexOutOfBoundsException if {@code index} is negative or is greater than {@code size} + * @throws IllegalArgumentException if {@code size} is negative + */ + + public static long checkPositionIndex(long index, long size, String desc) { + // Carefully optimized for execution by hotspot (explanatory comment above) + if (index < 0 || index > size) { + throw new IndexOutOfBoundsException(badPositionIndex(index, size, desc)); + } + return index; + } + + private static String badPositionIndex(long index, long size, String desc) { + if (index < 0) { + return format("%s (%s) must not be negative", desc, index); + } else if (size < 0) { + throw new IllegalArgumentException("negative size: " + size); + } else { // index > size + return format("%s (%s) must not be greater than size (%s)", desc, index, size); + } + } + + /** + * Ensures that {@code start} and {@code end} specify a valid <i>positions</i> in an array, list + * or string of size {@code size}, and are in order. A position index may range from zero to + * {@code size}, inclusive. + * + * @param start a user-supplied index identifying a starting position in an array, list or string + * @param end a user-supplied index identifying a ending position in an array, list or string + * @param size the size of that array, list or string + * @throws IndexOutOfBoundsException if either index is negative or is greater than {@code size}, + * or if {@code end} is less than {@code start} + * @throws IllegalArgumentException if {@code size} is negative + */ + public static void checkPositionIndexes(int start, int end, int size) { + // Carefully optimized for execution by hotspot (explanatory comment above) + if (start < 0 || end < start || end > size) { + throw new IndexOutOfBoundsException(badPositionIndexes(start, end, size)); + } + } + + private static String badPositionIndexes(int start, int end, int size) { + if (start < 0 || start > size) { + return badPositionIndex(start, size, "start index"); + } + if (end < 0 || end > size) { + return badPositionIndex(end, size, "end index"); + } + // end < start + return format("end index (%s) must not be less than start index (%s)", end, start); + } + + /** + * Substitutes each {@code %s} in {@code template} with an argument. These are matched by + * position: the first {@code %s} gets {@code args[0]}, etc. If there are more arguments than + * placeholders, the unmatched arguments will be appended to the end of the formatted message in + * square braces. + * + * @param template a non-null string containing 0 or more {@code %s} placeholders. + * @param args the arguments to be substituted into the message template. Arguments are converted + * to strings using {@link String#valueOf(Object)}. Arguments can be null. + */ + // Note that this is somewhat-improperly used from Verify.java as well. + static String format(String template, Object... args) { + template = String.valueOf(template); // null -> "null" + + // start substituting the arguments into the '%s' placeholders + StringBuilder builder = new StringBuilder(template.length() + 16 * args.length); + int templateStart = 0; + int i = 0; + while (i < args.length) { + int placeholderStart = template.indexOf("%s", templateStart); + if (placeholderStart == -1) { + break; + } + builder.append(template, templateStart, placeholderStart); + builder.append(args[i++]); + templateStart = placeholderStart + 2; + } + builder.append(template, templateStart, template.length()); + + // if we run out of placeholders, append the extra args in square braces + if (i < args.length) { + builder.append(" ["); + builder.append(args[i++]); + while (i < args.length) { + builder.append(", "); + builder.append(args[i++]); + } + builder.append(']'); + } + + return builder.toString(); + } +} diff --git a/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/VisibleForTesting.java b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/VisibleForTesting.java new file mode 100644 index 000000000..b6ed378bb --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/util/VisibleForTesting.java @@ -0,0 +1,26 @@ +/* + * 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.util; + +/** + * Annotation to indicate a class member or class is visible + * only for the purposes of testing and otherwise should not + * be referenced by other classes. + */ +public @interface VisibleForTesting { +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java new file mode 100644 index 000000000..bfe496532 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java @@ -0,0 +1,63 @@ +/* + * 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 org.apache.arrow.memory.util.MemoryUtil; + +/** + * The default Allocation Manager Factory for a module. + * + * This is only used by tests and contains only a simplistic allocator method. + * + */ +public class DefaultAllocationManagerFactory implements AllocationManager.Factory { + + public static final AllocationManager.Factory FACTORY = new DefaultAllocationManagerFactory(); + private static final ArrowBuf EMPTY = new ArrowBuf(ReferenceManager.NO_OP, + null, + 0, + MemoryUtil.UNSAFE.allocateMemory(0)); + + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long size) { + return new AllocationManager(accountingAllocator) { + private final long allocatedSize = size; + private final long address = MemoryUtil.UNSAFE.allocateMemory(size); + + @Override + public long getSize() { + return allocatedSize; + } + + @Override + protected long memoryAddress() { + return address; + } + + @Override + protected void release0() { + MemoryUtil.UNSAFE.freeMemory(address); + } + }; + } + + @Override + public ArrowBuf empty() { + return EMPTY; + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestAccountant.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestAccountant.java new file mode 100644 index 000000000..13fa4a64d --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestAccountant.java @@ -0,0 +1,172 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.Assert; +import org.junit.Test; + +public class TestAccountant { + + @Test + public void basic() { + ensureAccurateReservations(null); + } + + @Test + public void nested() { + final Accountant parent = new Accountant(null, "test", 0, Long.MAX_VALUE); + ensureAccurateReservations(parent); + assertEquals(0, parent.getAllocatedMemory()); + assertEquals(parent.getLimit() - parent.getAllocatedMemory(), parent.getHeadroom()); + } + + @Test + public void multiThread() throws InterruptedException { + final Accountant parent = new Accountant(null, "test", 0, Long.MAX_VALUE); + + final int numberOfThreads = 32; + final int loops = 100; + Thread[] threads = new Thread[numberOfThreads]; + + for (int i = 0; i < numberOfThreads; i++) { + Thread t = new Thread() { + + @Override + public void run() { + try { + for (int i = 0; i < loops; i++) { + ensureAccurateReservations(parent); + } + } catch (Exception ex) { + ex.printStackTrace(); + Assert.fail(ex.getMessage()); + } + } + + }; + threads[i] = t; + t.start(); + } + + for (Thread thread : threads) { + thread.join(); + } + + assertEquals(0, parent.getAllocatedMemory()); + assertEquals(parent.getLimit() - parent.getAllocatedMemory(), parent.getHeadroom()); + } + + private void ensureAccurateReservations(Accountant outsideParent) { + final Accountant parent = new Accountant(outsideParent, "test", 0, 10); + assertEquals(0, parent.getAllocatedMemory()); + + final Accountant child = new Accountant(parent, "test", 2, Long.MAX_VALUE); + assertEquals(2, parent.getAllocatedMemory()); + assertEquals(10, child.getHeadroom()); + { + AllocationOutcome first = child.allocateBytes(1); + assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); + } + + // child will have new allocation + assertEquals(1, child.getAllocatedMemory()); + + // root has no change since within reservation + assertEquals(2, parent.getAllocatedMemory()); + + { + AllocationOutcome first = child.allocateBytes(1); + assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); + } + + // child will have new allocation + assertEquals(2, child.getAllocatedMemory()); + + // root has no change since within reservation + assertEquals(2, parent.getAllocatedMemory()); + + child.releaseBytes(1); + + // child will have new allocation + assertEquals(1, child.getAllocatedMemory()); + + // root has no change since within reservation + assertEquals(2, parent.getAllocatedMemory()); + + { + AllocationOutcome first = child.allocateBytes(2); + assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); + } + + // child will have new allocation + assertEquals(3, child.getAllocatedMemory()); + + // went beyond reservation, now in parent accountant + assertEquals(3, parent.getAllocatedMemory()); + + assertEquals(7, child.getHeadroom()); + assertEquals(7, parent.getHeadroom()); + + { + AllocationOutcome first = child.allocateBytes(7); + assertEquals(AllocationOutcome.Status.SUCCESS, first.getStatus()); + } + + // child will have new allocation + assertEquals(10, child.getAllocatedMemory()); + + // went beyond reservation, now in parent accountant + assertEquals(10, parent.getAllocatedMemory()); + + child.releaseBytes(9); + + assertEquals(1, child.getAllocatedMemory()); + assertEquals(9, child.getHeadroom()); + + // back to reservation size + assertEquals(2, parent.getAllocatedMemory()); + assertEquals(8, parent.getHeadroom()); + + AllocationOutcome first = child.allocateBytes(10); + assertEquals(AllocationOutcome.Status.FAILED_PARENT, first.getStatus()); + + // unchanged + assertEquals(1, child.getAllocatedMemory()); + assertEquals(2, parent.getAllocatedMemory()); + + boolean withinLimit = child.forceAllocate(10); + assertEquals(false, withinLimit); + + // at new limit + assertEquals(child.getAllocatedMemory(), 11); + assertEquals(parent.getAllocatedMemory(), 11); + assertEquals(-1, child.getHeadroom()); + assertEquals(-1, parent.getHeadroom()); + + child.releaseBytes(11); + assertEquals(child.getAllocatedMemory(), 0); + assertEquals(parent.getAllocatedMemory(), 2); + assertEquals(10, child.getHeadroom()); + assertEquals(8, parent.getHeadroom()); + + child.close(); + parent.close(); + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestAllocationManager.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestAllocationManager.java new file mode 100644 index 000000000..df28424b3 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestAllocationManager.java @@ -0,0 +1,39 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.Test; + +/** + * Test cases for {@link AllocationManager}. + */ +public class TestAllocationManager { + + @Test + public void testAllocationManagerType() { + + // test unknown allocation manager type + System.clearProperty(DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME); + DefaultAllocationManagerOption.AllocationManagerType mgrType = + DefaultAllocationManagerOption.getDefaultAllocationManagerType(); + + assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Unknown, mgrType); + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java new file mode 100644 index 000000000..ea74f21d8 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestArrowBuf.java @@ -0,0 +1,149 @@ +/* + * 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.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; + +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.Arrays; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +public class TestArrowBuf { + + private static final int MAX_ALLOCATION = 8 * 1024; + private static RootAllocator allocator; + + @BeforeClass + public static void beforeClass() { + allocator = new RootAllocator(MAX_ALLOCATION); + } + + /** Ensure the allocator is closed. */ + @AfterClass + public static void afterClass() { + if (allocator != null) { + allocator.close(); + } + } + + @Test(expected = IndexOutOfBoundsException.class) + public void testSliceOutOfBoundsLength_RaisesIndexOutOfBoundsException() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(2) + ) { + assertEquals(2, buf.capacity()); + buf.slice(0, 3); + } + } + + @Test(expected = IndexOutOfBoundsException.class) + public void testSliceOutOfBoundsIndexPlusLength_RaisesIndexOutOfBoundsException() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(2) + ) { + assertEquals(2, buf.capacity()); + buf.slice(1, 2); + } + } + + @Test(expected = IndexOutOfBoundsException.class) + public void testSliceOutOfBoundsIndex_RaisesIndexOutOfBoundsException() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(2) + ) { + assertEquals(2, buf.capacity()); + buf.slice(3, 0); + } + } + + @Test + public void testSliceWithinBoundsLength_ReturnsSlice() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(2) + ) { + assertEquals(2, buf.capacity()); + assertEquals(1, buf.slice(1, 1).capacity()); + assertEquals(2, buf.slice(0, 2).capacity()); + } + } + + @Test + public void testSetBytesSliced() { + int arrLength = 64; + byte[] expected = new byte[arrLength]; + for (int i = 0; i < expected.length; i++) { + expected[i] = (byte) i; + } + ByteBuffer data = ByteBuffer.wrap(expected); + try (ArrowBuf buf = allocator.buffer(expected.length)) { + buf.setBytes(0, data, 0, data.capacity()); + + byte[] actual = new byte[expected.length]; + buf.getBytes(0, actual); + assertArrayEquals(expected, actual); + } + } + + @Test + public void testSetBytesUnsliced() { + int arrLength = 64; + byte[] arr = new byte[arrLength]; + for (int i = 0; i < arr.length; i++) { + arr[i] = (byte) i; + } + ByteBuffer data = ByteBuffer.wrap(arr); + + int from = 10; + int to = arrLength; + byte[] expected = Arrays.copyOfRange(arr, from, to); + try (ArrowBuf buf = allocator.buffer(expected.length)) { + buf.setBytes(0, data, from, to - from); + + byte[] actual = new byte[expected.length]; + buf.getBytes(0, actual); + assertArrayEquals(expected, actual); + } + } + + /** ARROW-9221: guard against big-endian byte buffers. */ + @Test + public void testSetBytesBigEndian() { + final byte[] expected = new byte[64]; + for (int i = 0; i < expected.length; i++) { + expected[i] = (byte) i; + } + // Only this code path is susceptible: others use unsafe or byte-by-byte copies, while this override copies longs. + final ByteBuffer data = ByteBuffer.wrap(expected).asReadOnlyBuffer(); + assertFalse(data.hasArray()); + assertFalse(data.isDirect()); + assertEquals(ByteOrder.BIG_ENDIAN, data.order()); + try (ArrowBuf buf = allocator.buffer(expected.length)) { + buf.setBytes(0, data); + byte[] actual = new byte[expected.length]; + buf.getBytes(0, actual); + assertArrayEquals(expected, actual); + } + } + +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBoundaryChecking.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBoundaryChecking.java new file mode 100644 index 000000000..5b86bed40 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestBoundaryChecking.java @@ -0,0 +1,150 @@ +/* + * 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 java.lang.reflect.Field; +import java.net.URLClassLoader; + +import org.junit.Assert; +import org.junit.Test; + +/** + * Test cases for evaluating the value of {@link BoundsChecking#BOUNDS_CHECKING_ENABLED}. + */ +public class TestBoundaryChecking { + + /** + * Get a copy of the current class loader. + * @return the newly created class loader. + */ + private ClassLoader copyClassLoader() { + ClassLoader curClassLoader = this.getClass().getClassLoader(); + if (curClassLoader instanceof URLClassLoader) { + // for Java 1.8 + return new URLClassLoader(((URLClassLoader) curClassLoader).getURLs(), null); + } + + // for Java 1.9 and Java 11. + return null; + } + + /** + * Get the value of flag {@link BoundsChecking#BOUNDS_CHECKING_ENABLED}. + * @param classLoader the class loader from which to get the flag value. + * @return value of the flag. + */ + private boolean getFlagValue(ClassLoader classLoader) throws Exception { + Class<?> clazz = classLoader.loadClass("org.apache.arrow.memory.BoundsChecking"); + Field field = clazz.getField("BOUNDS_CHECKING_ENABLED"); + return (Boolean) field.get(null); + } + + /** + * Ensure the flag for bounds checking is enabled by default. + * This will protect users from JVM crashes. + */ + @Test + public void testDefaultValue() throws Exception { + ClassLoader classLoader = copyClassLoader(); + if (classLoader != null) { + boolean boundsCheckingEnabled = getFlagValue(classLoader); + Assert.assertTrue(boundsCheckingEnabled); + } + } + + /** + * Test setting the bounds checking flag by the old property. + * @throws Exception if loading class {@link BoundsChecking#BOUNDS_CHECKING_ENABLED} fails. + */ + @Test + public void testEnableOldProperty() throws Exception { + String savedOldProperty = System.getProperty("drill.enable_unsafe_memory_access"); + System.setProperty("drill.enable_unsafe_memory_access", "true"); + + ClassLoader classLoader = copyClassLoader(); + if (classLoader != null) { + boolean boundsCheckingEnabled = getFlagValue(classLoader); + Assert.assertFalse(boundsCheckingEnabled); + } + + // restore system property + if (savedOldProperty != null) { + System.setProperty("drill.enable_unsafe_memory_access", savedOldProperty); + } else { + System.clearProperty("drill.enable_unsafe_memory_access"); + } + } + + /** + * Test setting the bounds checking flag by the new property. + * @throws Exception if loading class {@link BoundsChecking#BOUNDS_CHECKING_ENABLED} fails. + */ + @Test + public void testEnableNewProperty() throws Exception { + String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access"); + + System.setProperty("arrow.enable_unsafe_memory_access", "true"); + + ClassLoader classLoader = copyClassLoader(); + if (classLoader != null) { + boolean boundsCheckingEnabled = getFlagValue(classLoader); + Assert.assertFalse(boundsCheckingEnabled); + } + + // restore system property + if (savedNewProperty != null) { + System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty); + } else { + System.clearProperty("arrow.enable_unsafe_memory_access"); + } + } + + /** + * Test setting the bounds checking flag by both old and new properties. + * In this case, the new property should take precedence. + * @throws Exception if loading class {@link BoundsChecking#BOUNDS_CHECKING_ENABLED} fails. + */ + @Test + public void testEnableBothProperties() throws Exception { + String savedOldProperty = System.getProperty("drill.enable_unsafe_memory_access"); + String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access"); + + System.setProperty("drill.enable_unsafe_memory_access", "false"); + System.setProperty("arrow.enable_unsafe_memory_access", "true"); + + // new property takes precedence. + ClassLoader classLoader = copyClassLoader(); + if (classLoader != null) { + boolean boundsCheckingEnabled = getFlagValue(classLoader); + Assert.assertFalse(boundsCheckingEnabled); + } + + // restore system property + if (savedOldProperty != null) { + System.setProperty("drill.enable_unsafe_memory_access", savedOldProperty); + } else { + System.clearProperty("drill.enable_unsafe_memory_access"); + } + + if (savedNewProperty != null) { + System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty); + } else { + System.clearProperty("arrow.enable_unsafe_memory_access"); + } + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestLowCostIdentityHashMap.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestLowCostIdentityHashMap.java new file mode 100644 index 000000000..0cabc4a05 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/TestLowCostIdentityHashMap.java @@ -0,0 +1,169 @@ +/* + * 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 junit.framework.TestCase.assertNotNull; +import static junit.framework.TestCase.assertTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; + +import org.junit.Test; + +/** + * To test simplified implementation of IdentityHashMap. + */ +public class TestLowCostIdentityHashMap { + + @Test + public void testIdentityHashMap() throws Exception { + LowCostIdentityHashMap<String, StringWithKey> hashMap = new LowCostIdentityHashMap<>(); + + StringWithKey obj1 = new StringWithKey("s1key", "s1value"); + StringWithKey obj2 = new StringWithKey("s2key", "s2value"); + StringWithKey obj3 = new StringWithKey("s3key", "s3value"); + StringWithKey obj4 = new StringWithKey("s1key", "s4value"); + StringWithKey obj5 = new StringWithKey("s5key", "s5value"); + + assertNull(hashMap.put(obj1)); + assertNull(hashMap.put(obj2)); + assertNull(hashMap.put(obj3)); + assertEquals(obj1, hashMap.put(obj4)); + assertNull(hashMap.put(obj5)); + + assertEquals(4, hashMap.size()); + + assertEquals(obj4, hashMap.get("s1key")); + + assertNull(hashMap.remove("abc")); + + assertEquals(obj3, hashMap.remove("s3key")); + + assertEquals(3, hashMap.size()); + + assertTrue(!hashMap.isEmpty()); + + StringWithKey nextValue = hashMap.getNextValue(); + + assertNotNull(nextValue); + + assertTrue((hashMap.get("s1key") == nextValue || hashMap.get("s2key") == nextValue || + hashMap.get("s5key") == nextValue)); + + assertTrue(hashMap.containsValue(obj4)); + assertTrue(hashMap.containsValue(obj2)); + assertTrue(hashMap.containsValue(obj5)); + + assertEquals(obj4, hashMap.remove("s1key")); + + nextValue = hashMap.getNextValue(); + + assertNotNull(nextValue); + + assertTrue(hashMap.get("s2key") == nextValue || hashMap.get("s5key") == nextValue); + + assertEquals(2, hashMap.size()); + + assertEquals(obj2, hashMap.remove("s2key")); + assertEquals(obj5, hashMap.remove("s5key")); + + assertEquals(0, hashMap.size()); + + assertTrue(hashMap.isEmpty()); + } + + @Test + public void testLargeMap() throws Exception { + LowCostIdentityHashMap<String, StringWithKey> hashMap = new LowCostIdentityHashMap<>(); + + String [] keys = new String[200]; + for (int i = 0; i < 200; i++) { + keys[i] = "s" + i + "key"; + } + + for (int i = 0; i < 100; i++) { + if (i % 5 == 0 && i != 0) { + StringWithKey obj = new StringWithKey(keys[i - 5], "s" + i + "value"); + StringWithKey retObj = hashMap.put(obj); + assertNotNull(retObj); + StringWithKey obj1 = new StringWithKey(keys[i], "s" + 2 * i + "value"); + StringWithKey retObj1 = hashMap.put(obj1); + assertNull(retObj1); + } else { + StringWithKey obj = new StringWithKey(keys[i], "s" + i + "value"); + StringWithKey retObj = hashMap.put(obj); + assertNull(retObj); + } + } + assertEquals(100, hashMap.size()); + for (int i = 0; i < 100; i++) { + StringWithKey returnObj = hashMap.get(keys[i]); + assertNotNull(returnObj); + if (i == 95) { + assertEquals("s190value", returnObj.getValue()); + continue; + } + if (i % 5 == 0) { + assertEquals("s" + (i + 5) + "value", returnObj.getValue()); + } else { + assertEquals("s" + i + "value", returnObj.getValue()); + } + } + + for (int i = 0; i < 100; i++) { + if (i % 4 == 0) { + StringWithKey returnObj = hashMap.remove(keys[i]); + assertNotNull(returnObj); + assertTrue(!hashMap.containsKey(keys[i])); + } + StringWithKey obj = new StringWithKey(keys[100 + i], "s" + (100 + i) + "value"); + StringWithKey retObj = hashMap.put(obj); + assertNull(retObj); + assertTrue(hashMap.containsKey(keys[100 + i])); + } + assertEquals(175, hashMap.size()); + for (int i = 0; i < 100; i++) { + StringWithKey retObj = hashMap.getNextValue(); + assertNotNull(retObj); + hashMap.remove(retObj.getKey()); + } + assertTrue(!hashMap.isEmpty()); + assertEquals(75, hashMap.size()); + hashMap.clear(); + assertTrue(hashMap.isEmpty()); + } + + private class StringWithKey implements ValueWithKeyIncluded<String> { + + private String myValue; + private String myKey; + + StringWithKey(String myKey, String myValue) { + this.myKey = myKey; + this.myValue = myValue; + } + + @Override + public String getKey() { + return myKey; + } + + String getValue() { + return myValue; + } + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestArrowBufPointer.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestArrowBufPointer.java new file mode 100644 index 000000000..c3ed0d057 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestArrowBufPointer.java @@ -0,0 +1,216 @@ +/* + * 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.util; + +import static junit.framework.TestCase.assertEquals; +import static junit.framework.TestCase.assertTrue; +import static org.junit.jupiter.api.Assertions.assertFalse; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.memory.util.hash.ArrowBufHasher; +import org.apache.arrow.memory.util.hash.SimpleHasher; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Test cases for {@link ArrowBufPointer}. + */ +public class TestArrowBufPointer { + + private final int BUFFER_LENGTH = 1024; + + private BufferAllocator allocator; + + @Before + public void prepare() { + allocator = new RootAllocator(1024 * 1024); + } + + @After + public void shutdown() { + allocator.close(); + } + + @Test + public void testArrowBufPointersEqual() { + try (ArrowBuf buf1 = allocator.buffer(BUFFER_LENGTH); + ArrowBuf buf2 = allocator.buffer(BUFFER_LENGTH)) { + for (int i = 0; i < BUFFER_LENGTH / 4; i++) { + buf1.setInt(i * 4, i * 1234); + buf2.setInt(i * 4, i * 1234); + } + + ArrowBufPointer ptr1 = new ArrowBufPointer(null, 0, 100); + ArrowBufPointer ptr2 = new ArrowBufPointer(null, 100, 5032); + assertTrue(ptr1.equals(ptr2)); + for (int i = 0; i < BUFFER_LENGTH / 4; i++) { + ptr1.set(buf1, i * 4, 4); + ptr2.set(buf2, i * 4, 4); + assertTrue(ptr1.equals(ptr2)); + } + } + } + + @Test + public void testArrowBufPointersHashCode() { + final int vectorLength = 100; + try (ArrowBuf buf1 = allocator.buffer(vectorLength * 4); + ArrowBuf buf2 = allocator.buffer(vectorLength * 4)) { + for (int i = 0; i < vectorLength; i++) { + buf1.setInt(i * 4, i); + buf2.setInt(i * 4, i); + } + + CounterHasher hasher1 = new CounterHasher(); + CounterHasher hasher2 = new CounterHasher(); + + ArrowBufPointer pointer1 = new ArrowBufPointer(hasher1); + assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer1.hashCode()); + + ArrowBufPointer pointer2 = new ArrowBufPointer(hasher2); + assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer2.hashCode()); + + for (int i = 0; i < vectorLength; i++) { + pointer1.set(buf1, i * 4, 4); + pointer2.set(buf2, i * 4, 4); + + assertEquals(pointer1.hashCode(), pointer2.hashCode()); + + // verify that the hash codes have been re-computed + assertEquals(hasher1.counter, i + 1); + assertEquals(hasher2.counter, i + 1); + } + } + } + + @Test + public void testNullPointersHashCode() { + ArrowBufPointer pointer = new ArrowBufPointer(); + assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer.hashCode()); + + pointer.set(null, 0, 0); + assertEquals(ArrowBufPointer.NULL_HASH_CODE, pointer.hashCode()); + } + + @Test + public void testReuseHashCode() { + try (ArrowBuf buf = allocator.buffer(10)) { + buf.setInt(0, 10); + buf.setInt(4, 20); + + CounterHasher hasher = new CounterHasher(); + ArrowBufPointer pointer = new ArrowBufPointer(hasher); + + pointer.set(buf, 0, 4); + pointer.hashCode(); + + // hash code computed + assertEquals(1, hasher.counter); + + // no hash code re-compute + pointer.hashCode(); + assertEquals(1, hasher.counter); + + // hash code re-computed + pointer.set(buf, 4, 4); + pointer.hashCode(); + assertEquals(2, hasher.counter); + } + } + + @Test + public void testHashersForEquality() { + try (ArrowBuf buf = allocator.buffer(10)) { + // pointer 1 uses the default hasher + ArrowBufPointer pointer1 = new ArrowBufPointer(buf, 0, 10); + + // pointer 2 uses the counter hasher + ArrowBufPointer pointer2 = new ArrowBufPointer(buf, 0, 10, new CounterHasher()); + + // the two pointers cannot be equal, since they have different hashers + assertFalse(pointer1.equals(pointer2)); + } + } + + @Test + public void testArrowBufPointersComparison() { + final int vectorLength = 100; + try (ArrowBuf buf1 = allocator.buffer(vectorLength); + ArrowBuf buf2 = allocator.buffer(vectorLength)) { + for (int i = 0; i < vectorLength; i++) { + buf1.setByte(i, i); + buf2.setByte(i, i); + } + + ArrowBufPointer pointer1 = new ArrowBufPointer(); + ArrowBufPointer pointer2 = new ArrowBufPointer(); + + pointer1.set(buf1, 0, 10); + pointer2.set(buf2, 0, 10); + assertEquals(0, pointer1.compareTo(pointer2)); + + pointer1.set(null, 0, 0); + pointer2.set(null, 0, 0); + assertEquals(0, pointer1.compareTo(pointer2)); + + pointer2.set(buf2, 0, 5); + assertTrue(pointer1.compareTo(pointer2) < 0); + + pointer1.set(buf1, 0, 10); + assertTrue(pointer1.compareTo(pointer2) > 0); + + pointer1.set(buf1, 1, 5); + pointer2.set(buf2, 3, 8); + assertTrue(pointer1.compareTo(pointer2) < 0); + } + } + + /** + * Hasher with a counter that increments each time a hash code is calculated. + * This is to validate that the hash code in {@link ArrowBufPointer} is reused. + */ + class CounterHasher implements ArrowBufHasher { + + protected int counter = 0; + + @Override + public int hashCode(long address, long length) { + counter += 1; + return SimpleHasher.INSTANCE.hashCode(address, length); + } + + @Override + public int hashCode(ArrowBuf buf, long offset, long length) { + counter += 1; + return SimpleHasher.INSTANCE.hashCode(buf, offset, length); + } + + @Override + public int hashCode() { + return super.hashCode(); + } + + @Override + public boolean equals(Object o) { + return o != null && this.getClass() == o.getClass(); + } + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestByteFunctionHelpers.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestByteFunctionHelpers.java new file mode 100644 index 000000000..04a715962 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestByteFunctionHelpers.java @@ -0,0 +1,167 @@ +/* + * 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.util; + +import static org.junit.Assert.assertEquals; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +public class TestByteFunctionHelpers { + + private BufferAllocator allocator; + + private static final int SIZE = 100; + + @Before + public void init() { + allocator = new RootAllocator(Long.MAX_VALUE); + + } + + @After + public void terminate() throws Exception { + allocator.close(); + } + + @Test + public void testEquals() { + ArrowBuf buffer1 = allocator.buffer(SIZE); + ArrowBuf buffer2 = allocator.buffer(SIZE); + + for (int i = 0; i < SIZE; i++) { + buffer1.setByte(i, i); + buffer2.setByte(i, i); + } + + //test three cases, length>8, length>3, length<3 + + assertEquals(1, ByteFunctionHelpers.equal(buffer1, 0, SIZE - 1, + buffer2, 0, SIZE - 1)); + assertEquals(1, ByteFunctionHelpers.equal(buffer1, 0, 6, + buffer2, 0, 6)); + assertEquals(1, ByteFunctionHelpers.equal(buffer1, 0, 2, + buffer2, 0, 2)); + + //change value at index1 + buffer1.setByte(1, 10); + + assertEquals(0, ByteFunctionHelpers.equal(buffer1, 0, SIZE - 1, + buffer2, 0, SIZE - 1)); + assertEquals(0, ByteFunctionHelpers.equal(buffer1, 0, 6, + buffer2, 0, 6)); + assertEquals(0, ByteFunctionHelpers.equal(buffer1, 0, 2, + buffer2, 0, 2)); + + buffer1.close(); + buffer2.close(); + + } + + @Test + public void testCompare() { + ArrowBuf buffer1 = allocator.buffer(SIZE); + ArrowBuf buffer2 = allocator.buffer(SIZE); + + for (int i = 0; i < SIZE; i++) { + buffer1.setByte(i, i); + buffer2.setByte(i, i); + } + + //test three cases, length>8, length>3, length<3 + + assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, + buffer2, 0, SIZE - 1)); + assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 6, + buffer2, 0, 6)); + assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 2, + buffer2, 0, 2)); + + //change value at index 1 + buffer1.setByte(1, 0); + + assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, + buffer2, 0, SIZE - 1)); + assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 6, + buffer2, 0, 6)); + assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 2, + buffer2, 0, 2)); + + buffer1.close(); + buffer2.close(); + + } + + @Test + public void testStringCompare() { + String[] leftStrings = {"cat", "cats", "catworld", "dogs", "bags"}; + String[] rightStrings = {"dog", "dogs", "dogworld", "dog", "sgab"}; + + for (int i = 0; i < leftStrings.length; ++i) { + String leftStr = leftStrings[i]; + String rightStr = rightStrings[i]; + + ArrowBuf left = allocator.buffer(SIZE); + left.setBytes(0, leftStr.getBytes()); + ArrowBuf right = allocator.buffer(SIZE); + right.setBytes(0, rightStr.getBytes()); + + assertEquals(leftStr.compareTo(rightStr) < 0 ? -1 : 1, + ByteFunctionHelpers.compare(left, 0, leftStr.length(), right, 0, rightStr.length())); + + left.close(); + right.close(); + } + } + + @Test + public void testCompareWithByteArray() { + ArrowBuf buffer1 = allocator.buffer(SIZE); + byte[] buffer2 = new byte[SIZE]; + + for (int i = 0; i < SIZE; i++) { + buffer1.setByte(i, i); + buffer2[i] = (byte) i; + } + + //test three cases, length>8, length>3, length<3 + + assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, + buffer2, 0, SIZE - 1)); + assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 6, + buffer2, 0, 6)); + assertEquals(0, ByteFunctionHelpers.compare(buffer1, 0, 2, + buffer2, 0, 2)); + + //change value at index 1 + buffer1.setByte(1, 0); + + assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, SIZE - 1, + buffer2, 0, SIZE - 1)); + assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 6, + buffer2, 0, 6)); + assertEquals(-1, ByteFunctionHelpers.compare(buffer1, 0, 2, + buffer2, 0, 2)); + + buffer1.close(); + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestLargeMemoryUtil.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestLargeMemoryUtil.java new file mode 100755 index 000000000..952fcb5f0 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/TestLargeMemoryUtil.java @@ -0,0 +1,105 @@ +/*
+ * 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.util;
+
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.net.URLClassLoader;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
+
+public class TestLargeMemoryUtil {
+
+ /**
+ * Get a copy of the current class loader.
+ * @return the newly created class loader.
+ */
+ private ClassLoader copyClassLoader() {
+ ClassLoader curClassLoader = this.getClass().getClassLoader();
+ if (curClassLoader instanceof URLClassLoader) {
+ // for Java 1.8
+ return new URLClassLoader(((URLClassLoader) curClassLoader).getURLs(), null);
+ }
+
+ // for Java 1.9 and Java 11.
+ return null;
+ }
+
+ /**
+ * Use the checkedCastToInt method from the current classloader.
+ * @param classLoader the class loader from which to call the method.
+ * @return the return value of the method.
+ */
+ private int checkedCastToInt(ClassLoader classLoader, long value) throws Exception {
+ Class<?> clazz = classLoader.loadClass("org.apache.arrow.memory.util.LargeMemoryUtil");
+ Method method = clazz.getMethod("checkedCastToInt", long.class);
+ return (int) method.invoke(null, value);
+ }
+
+ private void checkExpectedOverflow(ClassLoader classLoader, long value) {
+ InvocationTargetException ex = Assertions.assertThrows(InvocationTargetException.class, () -> {
+ checkedCastToInt(classLoader, value);
+ });
+ Assert.assertTrue(ex.getCause() instanceof ArithmeticException);
+ Assert.assertEquals("integer overflow", ex.getCause().getMessage());
+ }
+
+ @Test
+ public void testEnableLargeMemoryUtilCheck() throws Exception {
+ String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access");
+ System.setProperty("arrow.enable_unsafe_memory_access", "false");
+ try {
+ ClassLoader classLoader = copyClassLoader();
+ if (classLoader != null) {
+ Assert.assertEquals(Integer.MAX_VALUE, checkedCastToInt(classLoader, Integer.MAX_VALUE));
+ checkExpectedOverflow(classLoader, Integer.MAX_VALUE + 1L);
+ checkExpectedOverflow(classLoader, Integer.MIN_VALUE - 1L);
+ }
+ } finally {
+ // restore system property
+ if (savedNewProperty != null) {
+ System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty);
+ } else {
+ System.clearProperty("arrow.enable_unsafe_memory_access");
+ }
+ }
+ }
+
+ @Test
+ public void testDisabledLargeMemoryUtilCheck() throws Exception {
+ String savedNewProperty = System.getProperty("arrow.enable_unsafe_memory_access");
+ System.setProperty("arrow.enable_unsafe_memory_access", "true");
+ try {
+ ClassLoader classLoader = copyClassLoader();
+ if (classLoader != null) {
+ Assert.assertEquals(Integer.MAX_VALUE, checkedCastToInt(classLoader, Integer.MAX_VALUE));
+ Assert.assertEquals(Integer.MIN_VALUE, checkedCastToInt(classLoader, Integer.MAX_VALUE + 1L));
+ Assert.assertEquals(Integer.MAX_VALUE, checkedCastToInt(classLoader, Integer.MIN_VALUE - 1L));
+ }
+ } finally {
+ // restore system property
+ if (savedNewProperty != null) {
+ System.setProperty("arrow.enable_unsafe_memory_access", savedNewProperty);
+ } else {
+ System.clearProperty("arrow.enable_unsafe_memory_access");
+ }
+ }
+ }
+}
diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/hash/TestArrowBufHasher.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/hash/TestArrowBufHasher.java new file mode 100644 index 000000000..a8707e6ca --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/memory/util/hash/TestArrowBufHasher.java @@ -0,0 +1,123 @@ +/* + * 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.util.hash; + +import static org.junit.Assert.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Arrays; +import java.util.Collection; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; + +/** + * Test cases for {@link ArrowBufHasher} and its subclasses. + */ +@RunWith(Parameterized.class) +public class TestArrowBufHasher { + + private final int BUFFER_LENGTH = 1024; + + private BufferAllocator allocator; + + private ArrowBufHasher hasher; + + public TestArrowBufHasher(String name, ArrowBufHasher hasher) { + this.hasher = hasher; + } + + @Before + public void prepare() { + allocator = new RootAllocator(1024 * 1024); + } + + @After + public void shutdown() { + allocator.close(); + } + + @Test + public void testHasher() { + try (ArrowBuf buf1 = allocator.buffer(BUFFER_LENGTH); + ArrowBuf buf2 = allocator.buffer(BUFFER_LENGTH)) { + // prepare data + for (int i = 0; i < BUFFER_LENGTH / 4; i++) { + buf1.setFloat(i * 4, i / 10.0f); + buf2.setFloat(i * 4, i / 10.0f); + } + + verifyHashCodesEqual(buf1, 0, 100, buf2, 0, 100); + verifyHashCodesEqual(buf1, 1, 5, buf2, 1, 5); + verifyHashCodesEqual(buf1, 10, 17, buf2, 10, 17); + verifyHashCodesEqual(buf1, 33, 25, buf2, 33, 25); + verifyHashCodesEqual(buf1, 22, 22, buf2, 22, 22); + verifyHashCodesEqual(buf1, 123, 333, buf2, 123, 333); + verifyHashCodesEqual(buf1, 374, 1, buf2, 374, 1); + verifyHashCodesEqual(buf1, 11, 0, buf2, 11, 0); + verifyHashCodesEqual(buf1, 75, 25, buf2, 75, 25); + verifyHashCodesEqual(buf1, 0, 1024, buf2, 0, 1024); + } + } + + private void verifyHashCodesEqual(ArrowBuf buf1, int offset1, int length1, + ArrowBuf buf2, int offset2, int length2) { + int hashCode1 = hasher.hashCode(buf1, offset1, length1); + int hashCode2 = hasher.hashCode(buf2, offset2, length2); + assertEquals(hashCode1, hashCode2); + } + + @Test + public void testHasherNegative() { + try (ArrowBuf buf = allocator.buffer(BUFFER_LENGTH)) { + // prepare data + for (int i = 0; i < BUFFER_LENGTH / 4; i++) { + buf.setFloat(i * 4, i / 10.0f); + } + + assertThrows(IllegalArgumentException.class, () -> { + hasher.hashCode(buf, 0, -1); + }); + + assertThrows(IndexOutOfBoundsException.class, () -> { + hasher.hashCode(buf, 0, 1028); + }); + + assertThrows(IndexOutOfBoundsException.class, () -> { + hasher.hashCode(buf, 500, 1000); + }); + } + } + + @Parameterized.Parameters(name = "hasher = {0}") + public static Collection<Object[]> getHasher() { + return Arrays.asList( + new Object[] {SimpleHasher.class.getSimpleName(), + SimpleHasher.INSTANCE}, + new Object[] {MurmurHasher.class.getSimpleName(), + new MurmurHasher() + } + ); + } +} diff --git a/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/util/TestCollections2.java b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/util/TestCollections2.java new file mode 100644 index 000000000..c858ebe62 --- /dev/null +++ b/src/arrow/java/memory/memory-core/src/test/java/org/apache/arrow/util/TestCollections2.java @@ -0,0 +1,83 @@ +/* + * 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.util; + +import static org.junit.Assert.fail; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; + +import org.junit.Test; + +/** + * Tests for {@code Collections2} class. + */ +public class TestCollections2 { + + + @Test + public void testToImmutableListFromIterable() { + final List<String> source = new ArrayList<>(Arrays.asList("foo", "bar", "baz")); + + final List<String> copy = Collections2.toImmutableList(source); + assertEquals(source, copy); + + try { + copy.add("unexpected"); + fail("add operation should not be supported"); + } catch (UnsupportedOperationException ignored) { + } + + try { + copy.set(0, "unexpected"); + fail("set operation should not be supported"); + } catch (UnsupportedOperationException ignored) { + } + + try { + copy.remove(0); + fail("remove operation should not be supported"); + } catch (UnsupportedOperationException ignored) { + } + + source.set(1, "newvalue"); + source.add("anothervalue"); + + assertEquals("bar", copy.get(1)); + assertEquals(3, copy.size()); + } + + + @Test + public void testStringFromEmptyIterator() { + assertEquals("[]", Collections2.toString(Collections.emptyIterator())); + } + + @Test + public void testStringFromIterator() { + Iterator<String> iterator = Arrays.asList("foo", "bar", "baz").iterator(); + iterator.next(); + + assertEquals("[bar, baz]", Collections2.toString(iterator)); + assertEquals(false, iterator.hasNext()); + } +} diff --git a/src/arrow/java/memory/memory-netty/pom.xml b/src/arrow/java/memory/memory-netty/pom.xml new file mode 100644 index 000000000..dee06a321 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/pom.xml @@ -0,0 +1,72 @@ +<?xml version="1.0" encoding="UTF-8"?> +<!-- 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. --> +<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> + <parent> + <artifactId>arrow-memory</artifactId> + <groupId>org.apache.arrow</groupId> + <version>6.0.1</version> + </parent> + <modelVersion>4.0.0</modelVersion> + + <artifactId>arrow-memory-netty</artifactId> + <name>Arrow Memory - Netty</name> + <description>Netty allocator and utils for allocating memory in Arrow</description> + + <dependencies> + <dependency> + <groupId>org.apache.arrow</groupId> + <artifactId>arrow-memory-core</artifactId> + <version>${project.version}</version> + </dependency> + <dependency> + <groupId>io.netty</groupId> + <artifactId>netty-buffer</artifactId> + </dependency> + <dependency> + <groupId>io.netty</groupId> + <artifactId>netty-common</artifactId> + </dependency> + <dependency> + <groupId>org.slf4j</groupId> + <artifactId>slf4j-api</artifactId> + </dependency> + <dependency> + <groupId>org.immutables</groupId> + <artifactId>value</artifactId> + </dependency> + </dependencies> + + <profiles> + <profile> + <!-- This profile turns on integration testing. It activates the failsafe plugin and will run any tests + with the 'IT' prefix. This should be run in a separate CI build or on developers machines as it potentially + uses quite a bit of memory. Activate the tests by adding -Pintegration-tests to your maven command line --> + <id>integration-tests</id> + <build> + <plugins> + <plugin> + <groupId>org.apache.maven.plugins</groupId> + <artifactId>maven-failsafe-plugin</artifactId> + <executions> + <execution> + <goals> + <goal>integration-test</goal> + <goal>verify</goal> + </goals> + </execution> + </executions> + </plugin> + </plugins> + </build> + </profile> + </profiles> +</project> diff --git a/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/ExpandableByteBuf.java b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/ExpandableByteBuf.java new file mode 100644 index 000000000..09b730044 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/ExpandableByteBuf.java @@ -0,0 +1,56 @@ +/* + * 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 io.netty.buffer; + +import org.apache.arrow.memory.BufferAllocator; + +/** + * Allows us to decorate ArrowBuf to make it expandable so that we can use them in the context of + * the Netty framework + * (thus supporting RPC level memory accounting). + */ +public class ExpandableByteBuf extends MutableWrappedByteBuf { + + private final BufferAllocator allocator; + + public ExpandableByteBuf(ByteBuf buffer, BufferAllocator allocator) { + super(buffer); + this.allocator = allocator; + } + + @Override + public ByteBuf copy(int index, int length) { + return new ExpandableByteBuf(buffer.copy(index, length), allocator); + } + + @Override + public ByteBuf capacity(int newCapacity) { + if (newCapacity > capacity()) { + ByteBuf newBuf = NettyArrowBuf.unwrapBuffer(allocator.buffer(newCapacity)); + newBuf.writeBytes(buffer, 0, buffer.capacity()); + newBuf.readerIndex(buffer.readerIndex()); + newBuf.writerIndex(buffer.writerIndex()); + buffer.release(); + buffer = newBuf; + return newBuf; + } else { + return super.capacity(newCapacity); + } + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/LargeBuffer.java b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/LargeBuffer.java new file mode 100644 index 000000000..792b3b814 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/LargeBuffer.java @@ -0,0 +1,34 @@ +/* + * 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 io.netty.buffer; + +/** + * A MutableWrappedByteBuf that also maintains a metric of the number of huge buffer bytes and + * counts. + */ +public class LargeBuffer extends MutableWrappedByteBuf { + + public LargeBuffer(ByteBuf buffer) { + super(buffer); + } + + @Override + public ByteBuf copy(int index, int length) { + return new LargeBuffer(buffer.copy(index, length)); + } +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/MutableWrappedByteBuf.java b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/MutableWrappedByteBuf.java new file mode 100644 index 000000000..5221dd3c1 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/MutableWrappedByteBuf.java @@ -0,0 +1,448 @@ +/* + * 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 io.netty.buffer; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.channels.GatheringByteChannel; +import java.nio.channels.ScatteringByteChannel; + +import io.netty.util.ByteProcessor; + +/** + * This is basically a complete copy of netty's DuplicatedByteBuf. We copy because we want to override + * some behaviors and make buffer mutable. + */ +abstract class MutableWrappedByteBuf extends AbstractByteBuf { + + ByteBuf buffer; + + public MutableWrappedByteBuf(ByteBuf buffer) { + super(buffer.maxCapacity()); + + if (buffer instanceof MutableWrappedByteBuf) { + this.buffer = ((MutableWrappedByteBuf) buffer).buffer; + } else { + this.buffer = buffer; + } + + setIndex(buffer.readerIndex(), buffer.writerIndex()); + } + + @Override + public ByteBuffer nioBuffer(int index, int length) { + return unwrap().nioBuffer(index, length); + } + + @Override + public ByteBuf unwrap() { + return buffer; + } + + @Override + public ByteBufAllocator alloc() { + return buffer.alloc(); + } + + @Override + public ByteOrder order() { + return buffer.order(); + } + + @Override + public boolean isDirect() { + return buffer.isDirect(); + } + + @Override + public int capacity() { + return buffer.capacity(); + } + + @Override + public ByteBuf capacity(int newCapacity) { + buffer.capacity(newCapacity); + return this; + } + + @Override + public boolean hasArray() { + return buffer.hasArray(); + } + + @Override + public byte[] array() { + return buffer.array(); + } + + @Override + public int arrayOffset() { + return buffer.arrayOffset(); + } + + @Override + public boolean hasMemoryAddress() { + return buffer.hasMemoryAddress(); + } + + @Override + public long memoryAddress() { + return buffer.memoryAddress(); + } + + @Override + public byte getByte(int index) { + return _getByte(index); + } + + @Override + protected byte _getByte(int index) { + return buffer.getByte(index); + } + + @Override + public short getShort(int index) { + return _getShort(index); + } + + @Override + protected short _getShort(int index) { + return buffer.getShort(index); + } + + @Override + public short getShortLE(int index) { + return buffer.getShortLE(index); + } + + @Override + protected short _getShortLE(int index) { + return buffer.getShortLE(index); + } + + @Override + public int getUnsignedMedium(int index) { + return _getUnsignedMedium(index); + } + + @Override + protected int _getUnsignedMedium(int index) { + return buffer.getUnsignedMedium(index); + } + + @Override + public int getUnsignedMediumLE(int index) { + return buffer.getUnsignedMediumLE(index); + } + + @Override + protected int _getUnsignedMediumLE(int index) { + return buffer.getUnsignedMediumLE(index); + } + + @Override + public int getInt(int index) { + return _getInt(index); + } + + @Override + protected int _getInt(int index) { + return buffer.getInt(index); + } + + @Override + public int getIntLE(int index) { + return buffer.getIntLE(index); + } + + @Override + protected int _getIntLE(int index) { + return buffer.getIntLE(index); + } + + @Override + public long getLong(int index) { + return _getLong(index); + } + + @Override + protected long _getLong(int index) { + return buffer.getLong(index); + } + + @Override + public long getLongLE(int index) { + return buffer.getLongLE(index); + } + + @Override + protected long _getLongLE(int index) { + return buffer.getLongLE(index); + } + + @Override + public abstract ByteBuf copy(int index, int length); + + @Override + public ByteBuf slice(int index, int length) { + return new SlicedByteBuf(this, index, length); + } + + @Override + public ByteBuf getBytes(int index, ByteBuf dst, int dstIndex, int length) { + buffer.getBytes(index, dst, dstIndex, length); + return this; + } + + @Override + public ByteBuf getBytes(int index, byte[] dst, int dstIndex, int length) { + buffer.getBytes(index, dst, dstIndex, length); + return this; + } + + @Override + public ByteBuf getBytes(int index, ByteBuffer dst) { + buffer.getBytes(index, dst); + return this; + } + + @Override + public ByteBuf setByte(int index, int value) { + _setByte(index, value); + return this; + } + + @Override + protected void _setByte(int index, int value) { + buffer.setByte(index, value); + } + + @Override + public ByteBuf setShort(int index, int value) { + _setShort(index, value); + return this; + } + + @Override + protected void _setShort(int index, int value) { + buffer.setShort(index, value); + } + + @Override + public ByteBuf setShortLE(int index, int value) { + buffer.setShortLE(index, value); + return this; + } + + @Override + protected void _setShortLE(int index, int value) { + buffer.setShortLE(index, value); + } + + @Override + public ByteBuf setMedium(int index, int value) { + _setMedium(index, value); + return this; + } + + @Override + protected void _setMedium(int index, int value) { + buffer.setMedium(index, value); + } + + @Override + public ByteBuf setMediumLE(int index, int value) { + buffer.setMediumLE(index, value); + return this; + } + + @Override + protected void _setMediumLE(int index, int value) { + buffer.setMediumLE(index, value); + } + + @Override + public ByteBuf setInt(int index, int value) { + _setInt(index, value); + return this; + } + + @Override + protected void _setInt(int index, int value) { + buffer.setInt(index, value); + } + + @Override + public ByteBuf setIntLE(int index, int value) { + buffer.setIntLE(index, value); + return this; + } + + @Override + protected void _setIntLE(int index, int value) { + buffer.setIntLE(index, value); + } + + @Override + public ByteBuf setLong(int index, long value) { + _setLong(index, value); + return this; + } + + @Override + protected void _setLong(int index, long value) { + buffer.setLong(index, value); + } + + @Override + public ByteBuf setLongLE(int index, long value) { + buffer.setLongLE(index, value); + return this; + } + + @Override + protected void _setLongLE(int index, long value) { + buffer.setLongLE(index, value); + } + + @Override + public ByteBuf setBytes(int index, byte[] src, int srcIndex, int length) { + buffer.setBytes(index, src, srcIndex, length); + return this; + } + + @Override + public ByteBuf setBytes(int index, ByteBuf src, int srcIndex, int length) { + buffer.setBytes(index, src, srcIndex, length); + return this; + } + + @Override + public ByteBuf setBytes(int index, ByteBuffer src) { + buffer.setBytes(index, src); + return this; + } + + @Override + public int setBytes(int index, FileChannel in, long position, int length) + throws IOException { + return buffer.setBytes(index, in, position, length); + } + + @Override + public ByteBuf getBytes(int index, OutputStream out, int length) + throws IOException { + buffer.getBytes(index, out, length); + return this; + } + + @Override + public int getBytes(int index, GatheringByteChannel out, int length) + throws IOException { + return buffer.getBytes(index, out, length); + } + + @Override + public int setBytes(int index, InputStream in, int length) + throws IOException { + return buffer.setBytes(index, in, length); + } + + @Override + public int setBytes(int index, ScatteringByteChannel in, int length) + throws IOException { + return buffer.setBytes(index, in, length); + } + + + @Override + public int getBytes(int index, FileChannel out, long position, int length) + throws IOException { + return buffer.getBytes(index, out, position, length); + } + + @Override + public int nioBufferCount() { + return buffer.nioBufferCount(); + } + + @Override + public ByteBuffer[] nioBuffers(int index, int length) { + return buffer.nioBuffers(index, length); + } + + @Override + public ByteBuffer internalNioBuffer(int index, int length) { + return nioBuffer(index, length); + } + + @Override + public int forEachByte(int index, int length, ByteProcessor processor) { + return buffer.forEachByte(index, length, processor); + } + + @Override + public int forEachByteDesc(int index, int length, ByteProcessor processor) { + return buffer.forEachByteDesc(index, length, processor); + } + + @Override + public final int refCnt() { + return unwrap().refCnt(); + } + + @Override + public final ByteBuf touch() { + unwrap().touch(); + return this; + } + + @Override + public final ByteBuf touch(Object hint) { + unwrap().touch(hint); + return this; + } + + @Override + public final ByteBuf retain() { + unwrap().retain(); + return this; + } + + @Override + public final ByteBuf retain(int increment) { + unwrap().retain(increment); + return this; + } + + @Override + public boolean release() { + return release(1); + } + + @Override + public boolean release(int decrement) { + boolean released = unwrap().release(decrement); + return released; + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/NettyArrowBuf.java b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/NettyArrowBuf.java new file mode 100644 index 000000000..8681b005f --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/NettyArrowBuf.java @@ -0,0 +1,622 @@ +/* + * 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 io.netty.buffer; + +import static org.apache.arrow.memory.util.LargeMemoryUtil.checkedCastToInt; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.nio.channels.FileChannel; +import java.nio.channels.GatheringByteChannel; +import java.nio.channels.ScatteringByteChannel; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.ArrowByteBufAllocator; +import org.apache.arrow.memory.BoundsChecking; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.util.Preconditions; + +import io.netty.util.internal.PlatformDependent; + +/** + * Netty specific wrapper over ArrowBuf for use in Netty framework. + */ +public class NettyArrowBuf extends AbstractByteBuf implements AutoCloseable { + + private final ArrowBuf arrowBuf; + private final ArrowByteBufAllocator arrowByteBufAllocator; + private int length; + private final long address; + + /** + * Constructs a new instance. + * + * @param arrowBuf The buffer to wrap. + * @param bufferAllocator The allocator for the buffer. + * @param length The length of this buffer. + */ + public NettyArrowBuf( + final ArrowBuf arrowBuf, + final BufferAllocator bufferAllocator, + final int length) { + super(length); + this.arrowBuf = arrowBuf; + this.arrowByteBufAllocator = new ArrowByteBufAllocator(bufferAllocator); + this.length = length; + this.address = arrowBuf.memoryAddress(); + } + + @Override + public ByteBuf copy() { + throw new UnsupportedOperationException(); + } + + @Override + public ByteBuf copy(int index, int length) { + throw new UnsupportedOperationException(); + } + + @Override + public ByteBuf retain() { + arrowBuf.getReferenceManager().retain(); + return this; + } + + public ArrowBuf arrowBuf() { + return arrowBuf; + } + + @Override + public ByteBuf retain(final int increment) { + arrowBuf.getReferenceManager().retain(increment); + return this; + } + + @Override + public boolean isDirect() { + return true; + } + + @Override + public synchronized ByteBuf capacity(int newCapacity) { + if (newCapacity == length) { + return this; + } + Preconditions.checkArgument(newCapacity >= 0); + if (newCapacity < length) { + length = newCapacity; + return this; + } + throw new UnsupportedOperationException("Buffers don't support resizing that increases the size."); + } + + @Override + public ByteBuf unwrap() { + throw new UnsupportedOperationException("Unwrap not supported."); + } + + @Override + public int refCnt() { + return arrowBuf.getReferenceManager().getRefCount(); + } + + @Override + public ArrowByteBufAllocator alloc() { + return arrowByteBufAllocator; + } + + @Override + public boolean hasArray() { + return false; + } + + @Override + public byte[] array() { + throw new UnsupportedOperationException("Operation not supported on direct buffer"); + } + + @Override + public int arrayOffset() { + throw new UnsupportedOperationException("Operation not supported on direct buffer"); + } + + @Override + public boolean hasMemoryAddress() { + return true; + } + + @Override + public long memoryAddress() { + return this.address; + } + + @Override + public ByteBuf touch() { + return this; + } + + @Override + public ByteBuf touch(Object hint) { + return this; + } + + @Override + public int capacity() { + return (int) Math.min(Integer.MAX_VALUE, arrowBuf.capacity()); + } + + @Override + public NettyArrowBuf slice() { + return unwrapBuffer(arrowBuf.slice(readerIndex, writerIndex - readerIndex)); + } + + @Override + public NettyArrowBuf slice(int index, int length) { + return unwrapBuffer(arrowBuf.slice(index, length)); + } + + @Override + public void close() { + arrowBuf.close(); + } + + @Override + public boolean release() { + return arrowBuf.getReferenceManager().release(); + } + + @Override + public boolean release(int decrement) { + return arrowBuf.getReferenceManager().release(decrement); + } + + @Override + public NettyArrowBuf readerIndex(int readerIndex) { + super.readerIndex(readerIndex); + return this; + } + + @Override + public NettyArrowBuf writerIndex(int writerIndex) { + super.writerIndex(writerIndex); + return this; + } + + @Override + public int nioBufferCount() { + return 1; + } + + @Override + public ByteBuffer internalNioBuffer(int index, int length) { + ByteBuffer nioBuf = getDirectBuffer(index); + // Follows convention from other ByteBuf implementations. + return (ByteBuffer) nioBuf.clear().limit(length); + } + + @Override + public ByteBuffer[] nioBuffers() { + return new ByteBuffer[] {nioBuffer()}; + } + + @Override + public ByteBuffer[] nioBuffers(int index, int length) { + return new ByteBuffer[] {nioBuffer(index, length)}; + } + + @Override + public ByteBuffer nioBuffer() { + return nioBuffer(readerIndex(), readableBytes()); + } + + + /** + * Returns a buffer that is zero positioned but points + * to a slice of the original buffer starting at given index. + */ + @Override + public ByteBuffer nioBuffer(int index, int length) { + chk(index, length); + final ByteBuffer buffer = getDirectBuffer(index); + buffer.limit(length); + return buffer; + } + + /** + * Returns a buffer that is zero positioned but points + * to a slice of the original buffer starting at given index. + */ + public ByteBuffer nioBuffer(long index, int length) { + chk(index, length); + final ByteBuffer buffer = getDirectBuffer(index); + buffer.limit(length); + return buffer; + } + + /** + * Get this ArrowBuf as a direct {@link ByteBuffer}. + * + * @return ByteBuffer + */ + private ByteBuffer getDirectBuffer(long index) { + return PlatformDependent.directBuffer(addr(index), checkedCastToInt(length - index)); + } + + @Override + public ByteBuf getBytes(int index, ByteBuffer dst) { + arrowBuf.getBytes(index, dst); + return this; + } + + @Override + public ByteBuf setBytes(int index, ByteBuffer src) { + arrowBuf.setBytes(index, src); + return this; + } + + @Override + public ByteBuf getBytes(int index, byte[] dst, int dstIndex, int length) { + arrowBuf.getBytes(index, dst, dstIndex, length); + return this; + } + + @Override + public ByteBuf setBytes(int index, byte[] src, int srcIndex, int length) { + arrowBuf.setBytes(index, src, srcIndex, length); + return this; + } + + /** + * Determine if the requested {@code index} and {@code length} will fit within {@code capacity}. + * + * @param index The starting index. + * @param length The length which will be utilized (starting from {@code index}). + * @param capacity The capacity that {@code index + length} is allowed to be within. + * @return {@code true} if the requested {@code index} and {@code length} will fit within {@code capacity}. + * {@code false} if this would result in an index out of bounds exception. + */ + private static boolean isOutOfBounds(int index, int length, int capacity) { + return (index | length | (index + length) | (capacity - (index + length))) < 0; + } + + @Override + public ByteBuf getBytes(int index, ByteBuf dst, int dstIndex, int length) { + chk(index, length); + Preconditions.checkArgument(dst != null, "Expecting valid dst ByteBuffer"); + if (isOutOfBounds(dstIndex, length, dst.capacity())) { + throw new IndexOutOfBoundsException("dstIndex: " + dstIndex + " length: " + length); + } else { + final long srcAddress = addr(index); + if (dst.hasMemoryAddress()) { + final long dstAddress = dst.memoryAddress() + (long) dstIndex; + PlatformDependent.copyMemory(srcAddress, dstAddress, (long) length); + } else if (dst.hasArray()) { + dstIndex += dst.arrayOffset(); + PlatformDependent.copyMemory(srcAddress, dst.array(), dstIndex, (long) length); + } else { + dst.setBytes(dstIndex, this, index, length); + } + } + return this; + } + + @Override + public ByteBuf setBytes(int index, ByteBuf src, int srcIndex, int length) { + chk(index, length); + Preconditions.checkArgument(src != null, "Expecting valid src ByteBuffer"); + if (isOutOfBounds(srcIndex, length, src.capacity())) { + throw new IndexOutOfBoundsException("srcIndex: " + srcIndex + " length: " + length); + } else { + if (length != 0) { + final long dstAddress = addr(index); + if (src.hasMemoryAddress()) { + final long srcAddress = src.memoryAddress() + (long) srcIndex; + PlatformDependent.copyMemory(srcAddress, dstAddress, (long) length); + } else if (src.hasArray()) { + srcIndex += src.arrayOffset(); + PlatformDependent.copyMemory(src.array(), srcIndex, dstAddress, (long) length); + } else { + src.getBytes(srcIndex, this, index, length); + } + } + } + return this; + } + + @Override + public ByteBuf getBytes(int index, OutputStream out, int length) throws IOException { + arrowBuf.getBytes(index, out, length); + return this; + } + + @Override + public int setBytes(int index, InputStream in, int length) throws IOException { + return arrowBuf.setBytes(index, in, length); + } + + @Override + public int getBytes(int index, GatheringByteChannel out, int length) throws IOException { + Preconditions.checkArgument(out != null, "expecting valid gathering byte channel"); + chk(index, length); + if (length == 0) { + return 0; + } else { + final ByteBuffer tmpBuf = getDirectBuffer(index); + tmpBuf.clear().limit(length); + return out.write(tmpBuf); + } + } + + @Override + public int getBytes(int index, FileChannel out, long position, int length) throws IOException { + chk(index, length); + if (length == 0) { + return 0; + } else { + final ByteBuffer tmpBuf = getDirectBuffer(index); + tmpBuf.clear().limit(length); + return out.write(tmpBuf, position); + } + } + + @Override + public int setBytes(int index, ScatteringByteChannel in, int length) throws IOException { + return (int) in.read(nioBuffers(index, length)); + } + + @Override + public int setBytes(int index, FileChannel in, long position, int length) throws IOException { + return (int) in.read(nioBuffers(index, length)); + } + + @Override + public ByteOrder order() { + return ByteOrder.LITTLE_ENDIAN; + } + + @Override + public ByteBuf order(ByteOrder endianness) { + return this; + } + + @Override + protected int _getUnsignedMedium(int index) { + return getUnsignedMedium(index); + } + + @Override + protected int _getUnsignedMediumLE(int index) { + this.chk(index, 3); + long addr = this.addr(index); + return PlatformDependent.getByte(addr) & 255 | + (Short.reverseBytes(PlatformDependent.getShort(addr + 1L)) & '\uffff') << 8; + } + + + /*-------------------------------------------------* + | | + | get() APIs | + | | + *-------------------------------------------------*/ + + + @Override + protected byte _getByte(int index) { + return getByte(index); + } + + @Override + public byte getByte(int index) { + return arrowBuf.getByte(index); + } + + @Override + protected short _getShortLE(int index) { + short s = getShort(index); + return Short.reverseBytes(s); + } + + @Override + protected short _getShort(int index) { + return getShort(index); + } + + @Override + public short getShort(int index) { + return arrowBuf.getShort(index); + } + + @Override + protected int _getIntLE(int index) { + int value = getInt(index); + return Integer.reverseBytes(value); + } + + @Override + protected int _getInt(int index) { + return getInt(index); + } + + @Override + public int getInt(int index) { + return arrowBuf.getInt(index); + } + + @Override + protected long _getLongLE(int index) { + long value = getLong(index); + return Long.reverseBytes(value); + } + + @Override + protected long _getLong(int index) { + return getLong(index); + } + + @Override + public long getLong(int index) { + return arrowBuf.getLong(index); + } + + + /*-------------------------------------------------* + | | + | set() APIs | + | | + *-------------------------------------------------*/ + + + @Override + protected void _setByte(int index, int value) { + setByte(index, value); + } + + @Override + public NettyArrowBuf setByte(int index, int value) { + arrowBuf.setByte(index, value); + return this; + } + + @Override + protected void _setShortLE(int index, int value) { + this.chk(index, 2); + PlatformDependent.putShort(this.addr(index), Short.reverseBytes((short) value)); + } + + @Override + protected void _setShort(int index, int value) { + setShort(index, value); + } + + @Override + public NettyArrowBuf setShort(int index, int value) { + arrowBuf.setShort(index, value); + return this; + } + + private long addr(long index) { + return address + index; + } + + /** + * Helper function to do bounds checking at a particular + * index for particular length of data. + * + * @param index index (0 based relative to this ArrowBuf) + * @param fieldLength provided length of data for get/set + */ + private void chk(long index, long fieldLength) { + if (BoundsChecking.BOUNDS_CHECKING_ENABLED) { + // check reference count + ensureAccessible(); + // check bounds + if (fieldLength < 0) { + throw new IllegalArgumentException("length: " + fieldLength + " (expected: >= 0)"); + } + if (index < 0 || index > capacity() - fieldLength) { + throw new IndexOutOfBoundsException(String.format( + "index: %d, length: %d (expected: range(0, %d))", index, fieldLength, capacity())); + } + } + } + + @Override + protected void _setMedium(int index, int value) { + setMedium(index, value); + } + + @Override + protected void _setMediumLE(int index, int value) { + this.chk(index, 3); + long addr = this.addr(index); + PlatformDependent.putByte(addr, (byte) value); + PlatformDependent.putShort(addr + 1L, Short.reverseBytes((short) (value >>> 8))); + } + + @Override + public NettyArrowBuf setMedium(int index, int value) { + chk(index, 3); + final long addr = addr(index); + // we need to store 3 bytes starting from least significant byte + // and ignoring the most significant byte + // since arrow memory format is little endian, we will + // first store the first 2 bytes followed by third byte + // example: if the 4 byte int value is ABCD where A is MSB + // D is LSB then we effectively want to store DCB in increasing + // address to get Little Endian byte order + // (short)value will give us CD and PlatformDependent.putShort() + // will store them in LE order as DC starting at address addr + // in order to get B, we do ABCD >>> 16 = 00AB => (byte)AB which + // gives B. We store this at address addr + 2. So finally we get + // DCB + PlatformDependent.putShort(addr, (short) value); + PlatformDependent.putByte(addr + 2, (byte) (value >>> 16)); + return this; + } + + @Override + protected void _setInt(int index, int value) { + setInt(index, value); + } + + @Override + protected void _setIntLE(int index, int value) { + this.chk(index, 4); + PlatformDependent.putInt(this.addr(index), Integer.reverseBytes(value)); + } + + @Override + public NettyArrowBuf setInt(int index, int value) { + arrowBuf.setInt(index, value); + return this; + } + + @Override + protected void _setLong(int index, long value) { + setLong(index, value); + } + + @Override + public void _setLongLE(int index, long value) { + this.chk(index, 8); + PlatformDependent.putLong(this.addr(index), Long.reverseBytes(value)); + } + + @Override + public NettyArrowBuf setLong(int index, long value) { + arrowBuf.setLong(index, value); + return this; + } + + /** + * unwrap arrow buffer into a netty buffer. + */ + public static NettyArrowBuf unwrapBuffer(ArrowBuf buf) { + final NettyArrowBuf nettyArrowBuf = new NettyArrowBuf( + buf, + buf.getReferenceManager().getAllocator(), + checkedCastToInt(buf.capacity())); + nettyArrowBuf.readerIndex(checkedCastToInt(buf.readerIndex())); + nettyArrowBuf.writerIndex(checkedCastToInt(buf.writerIndex())); + return nettyArrowBuf; + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/PooledByteBufAllocatorL.java b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/PooledByteBufAllocatorL.java new file mode 100644 index 000000000..d0a5a9945 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/PooledByteBufAllocatorL.java @@ -0,0 +1,280 @@ +/* + * 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 io.netty.buffer; + +import static org.apache.arrow.memory.util.AssertionUtil.ASSERT_ENABLED; + +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.util.concurrent.atomic.AtomicLong; + +import org.apache.arrow.memory.OutOfMemoryException; +import org.apache.arrow.memory.util.LargeMemoryUtil; + +import io.netty.util.internal.OutOfDirectMemoryError; +import io.netty.util.internal.StringUtil; + +/** + * The base allocator that we use for all of Arrow's memory management. Returns + * UnsafeDirectLittleEndian buffers. + */ +public class PooledByteBufAllocatorL { + + private static final org.slf4j.Logger memoryLogger = org.slf4j.LoggerFactory.getLogger("arrow.allocator"); + + private static final int MEMORY_LOGGER_FREQUENCY_SECONDS = 60; + public final UnsafeDirectLittleEndian empty; + private final AtomicLong hugeBufferSize = new AtomicLong(0); + private final AtomicLong hugeBufferCount = new AtomicLong(0); + private final AtomicLong normalBufferSize = new AtomicLong(0); + private final AtomicLong normalBufferCount = new AtomicLong(0); + private final InnerAllocator allocator; + + public PooledByteBufAllocatorL() { + allocator = new InnerAllocator(); + empty = new UnsafeDirectLittleEndian(new DuplicatedByteBuf(Unpooled.EMPTY_BUFFER)); + } + + /** + * Returns a {@linkplain io.netty.buffer.UnsafeDirectLittleEndian} of the given size. + */ + public UnsafeDirectLittleEndian allocate(long size) { + try { + return allocator.directBuffer(LargeMemoryUtil.checkedCastToInt(size), Integer.MAX_VALUE); + } catch (OutOfMemoryError e) { + /* + * OutOfDirectMemoryError is thrown by Netty when we exceed the direct memory limit defined by + * -XX:MaxDirectMemorySize. OutOfMemoryError with "Direct buffer memory" message is thrown by + * java.nio.Bits when we exceed the direct memory limit. This should never be hit in practice + * as Netty is expected to throw an OutOfDirectMemoryError first. + */ + if (e instanceof OutOfDirectMemoryError || "Direct buffer memory".equals(e.getMessage())) { + throw new OutOfMemoryException("Failure allocating buffer.", e); + } + throw e; + } + } + + public int getChunkSize() { + return allocator.chunkSize; + } + + public long getHugeBufferSize() { + return hugeBufferSize.get(); + } + + public long getHugeBufferCount() { + return hugeBufferCount.get(); + } + + public long getNormalBufferSize() { + return normalBufferSize.get(); + } + + public long getNormalBufferCount() { + return normalBufferSize.get(); + } + + private static class AccountedUnsafeDirectLittleEndian extends UnsafeDirectLittleEndian { + + private final long initialCapacity; + private final AtomicLong count; + private final AtomicLong size; + + private AccountedUnsafeDirectLittleEndian(LargeBuffer buf, AtomicLong count, AtomicLong size) { + super(buf); + this.initialCapacity = buf.capacity(); + this.count = count; + this.size = size; + } + + private AccountedUnsafeDirectLittleEndian(PooledUnsafeDirectByteBuf buf, AtomicLong count, + AtomicLong size) { + super(buf); + this.initialCapacity = buf.capacity(); + this.count = count; + this.size = size; + } + + @Override + public ByteBuf copy() { + throw new UnsupportedOperationException("copy method is not supported"); + } + + @Override + public ByteBuf copy(int index, int length) { + throw new UnsupportedOperationException("copy method is not supported"); + } + + @Override + public boolean release(int decrement) { + boolean released = super.release(decrement); + if (released) { + count.decrementAndGet(); + size.addAndGet(-initialCapacity); + } + return released; + } + + } + + private class InnerAllocator extends PooledByteBufAllocator { + + private final PoolArena<ByteBuffer>[] directArenas; + private final MemoryStatusThread statusThread; + private final int chunkSize; + + public InnerAllocator() { + super(true); + + try { + Field f = PooledByteBufAllocator.class.getDeclaredField("directArenas"); + f.setAccessible(true); + this.directArenas = (PoolArena<ByteBuffer>[]) f.get(this); + } catch (Exception e) { + throw new RuntimeException("Failure while initializing allocator. Unable to retrieve direct arenas field.", e); + } + + this.chunkSize = directArenas[0].chunkSize; + + if (memoryLogger.isTraceEnabled()) { + statusThread = new MemoryStatusThread(); + statusThread.start(); + } else { + statusThread = null; + } + } + + private UnsafeDirectLittleEndian newDirectBufferL(int initialCapacity, int maxCapacity) { + PoolThreadCache cache = threadCache(); + PoolArena<ByteBuffer> directArena = cache.directArena; + + if (directArena != null) { + + if (initialCapacity > directArena.chunkSize) { + // This is beyond chunk size so we'll allocate separately. + ByteBuf buf = UnpooledByteBufAllocator.DEFAULT.directBuffer(initialCapacity, maxCapacity); + + hugeBufferSize.addAndGet(buf.capacity()); + hugeBufferCount.incrementAndGet(); + + // logger.debug("Allocating huge buffer of size {}", initialCapacity, new Exception()); + return new AccountedUnsafeDirectLittleEndian(new LargeBuffer(buf), hugeBufferCount, + hugeBufferSize); + } else { + // within chunk, use arena. + ByteBuf buf = directArena.allocate(cache, initialCapacity, maxCapacity); + if (!(buf instanceof PooledUnsafeDirectByteBuf)) { + fail(); + } + + if (!ASSERT_ENABLED) { + return new UnsafeDirectLittleEndian((PooledUnsafeDirectByteBuf) buf); + } + + normalBufferSize.addAndGet(buf.capacity()); + normalBufferCount.incrementAndGet(); + + return new AccountedUnsafeDirectLittleEndian((PooledUnsafeDirectByteBuf) buf, + normalBufferCount, normalBufferSize); + } + + } else { + throw fail(); + } + } + + private UnsupportedOperationException fail() { + return new UnsupportedOperationException( + "Arrow requires that the JVM used supports access sun.misc.Unsafe. This platform " + + "didn't provide that functionality."); + } + + @Override + public UnsafeDirectLittleEndian directBuffer(int initialCapacity, int maxCapacity) { + if (initialCapacity == 0 && maxCapacity == 0) { + newDirectBuffer(initialCapacity, maxCapacity); + } + validate(initialCapacity, maxCapacity); + return newDirectBufferL(initialCapacity, maxCapacity); + } + + @Override + public ByteBuf heapBuffer(int initialCapacity, int maxCapacity) { + throw new UnsupportedOperationException("Arrow doesn't support using heap buffers."); + } + + + private void validate(int initialCapacity, int maxCapacity) { + if (initialCapacity < 0) { + throw new IllegalArgumentException("initialCapacity: " + initialCapacity + " (expected: 0+)"); + } + if (initialCapacity > maxCapacity) { + throw new IllegalArgumentException(String.format( + "initialCapacity: %d (expected: not greater than maxCapacity(%d)", + initialCapacity, maxCapacity)); + } + } + + @Override + public String toString() { + StringBuilder buf = new StringBuilder(); + buf.append(directArenas.length); + buf.append(" direct arena(s):"); + buf.append(StringUtil.NEWLINE); + for (PoolArena<ByteBuffer> a : directArenas) { + buf.append(a); + } + + buf.append("Large buffers outstanding: "); + buf.append(hugeBufferCount.get()); + buf.append(" totaling "); + buf.append(hugeBufferSize.get()); + buf.append(" bytes."); + buf.append('\n'); + buf.append("Normal buffers outstanding: "); + buf.append(normalBufferCount.get()); + buf.append(" totaling "); + buf.append(normalBufferSize.get()); + buf.append(" bytes."); + return buf.toString(); + } + + private class MemoryStatusThread extends Thread { + + public MemoryStatusThread() { + super("allocation.logger"); + this.setDaemon(true); + } + + @Override + public void run() { + while (true) { + memoryLogger.trace("Memory Usage: \n{}", PooledByteBufAllocatorL.this.toString()); + try { + Thread.sleep(MEMORY_LOGGER_FREQUENCY_SECONDS * 1000); + } catch (InterruptedException e) { + return; + } + } + } + } + + + } +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/UnsafeDirectLittleEndian.java b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/UnsafeDirectLittleEndian.java new file mode 100644 index 000000000..e900b1ca7 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/io/netty/buffer/UnsafeDirectLittleEndian.java @@ -0,0 +1,270 @@ +/* + * 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 io.netty.buffer; + +import java.io.IOException; +import java.io.InputStream; +import java.io.OutputStream; +import java.nio.ByteOrder; +import java.util.concurrent.atomic.AtomicLong; + +import io.netty.util.internal.PlatformDependent; + +/** + * The underlying class we use for little-endian access to memory. Is used underneath ArrowBufs + * to abstract away the + * Netty classes and underlying Netty memory management. + */ +public class UnsafeDirectLittleEndian extends WrappedByteBuf { + + public static final boolean ASSERT_ENABLED; + private static final AtomicLong ID_GENERATOR = new AtomicLong(0); + + static { + boolean isAssertEnabled = false; + assert isAssertEnabled = true; + ASSERT_ENABLED = isAssertEnabled; + } + + public final long id = ID_GENERATOR.incrementAndGet(); + private final AbstractByteBuf wrapped; + private final long memoryAddress; + + UnsafeDirectLittleEndian(DuplicatedByteBuf buf) { + this(buf, true); + } + + UnsafeDirectLittleEndian(LargeBuffer buf) { + this(buf, true); + } + + UnsafeDirectLittleEndian(PooledUnsafeDirectByteBuf buf) { + this(buf, true); + } + + private UnsafeDirectLittleEndian(AbstractByteBuf buf, boolean fake) { + super(buf); + + this.wrapped = buf; + this.memoryAddress = buf.memoryAddress(); + } + + private long addr(int index) { + return memoryAddress + index; + } + + @Override + public long getLong(int index) { + // wrapped.checkIndex(index, 8); + long v = PlatformDependent.getLong(addr(index)); + return v; + } + + @Override + public float getFloat(int index) { + return Float.intBitsToFloat(getInt(index)); + } + + @Override + public ByteBuf slice() { + return slice(this.readerIndex(), readableBytes()); + } + + @Override + public ByteBuf slice(int index, int length) { + return new SlicedByteBuf(this, index, length); + } + + @Override + public ByteBuf order(ByteOrder endianness) { + return this; + } + + @Override + public double getDouble(int index) { + return Double.longBitsToDouble(getLong(index)); + } + + @Override + public char getChar(int index) { + return (char) getShort(index); + } + + @Override + public long getUnsignedInt(int index) { + return getInt(index) & 0xFFFFFFFFL; + } + + @Override + public int getInt(int index) { + int v = PlatformDependent.getInt(addr(index)); + return v; + } + + @Override + public int getUnsignedShort(int index) { + return getShort(index) & 0xFFFF; + } + + @Override + public short getShort(int index) { + short v = PlatformDependent.getShort(addr(index)); + return v; + } + + @Override + public ByteBuf setShort(int index, int value) { + wrapped.checkIndex(index, 2); + setShort_(index, value); + return this; + } + + @Override + public ByteBuf setInt(int index, int value) { + wrapped.checkIndex(index, 4); + setInt_(index, value); + return this; + } + + @Override + public ByteBuf setLong(int index, long value) { + wrapped.checkIndex(index, 8); + setLong_(index, value); + return this; + } + + @Override + public ByteBuf setChar(int index, int value) { + setShort(index, value); + return this; + } + + @Override + public ByteBuf setFloat(int index, float value) { + setInt(index, Float.floatToRawIntBits(value)); + return this; + } + + @Override + public ByteBuf setDouble(int index, double value) { + setLong(index, Double.doubleToRawLongBits(value)); + return this; + } + + @Override + public ByteBuf writeShort(int value) { + wrapped.ensureWritable(2); + setShort_(wrapped.writerIndex, value); + wrapped.writerIndex += 2; + return this; + } + + @Override + public ByteBuf writeInt(int value) { + wrapped.ensureWritable(4); + setInt_(wrapped.writerIndex, value); + wrapped.writerIndex += 4; + return this; + } + + @Override + public ByteBuf writeLong(long value) { + wrapped.ensureWritable(8); + setLong_(wrapped.writerIndex, value); + wrapped.writerIndex += 8; + return this; + } + + @Override + public ByteBuf writeChar(int value) { + writeShort(value); + return this; + } + + @Override + public ByteBuf writeFloat(float value) { + writeInt(Float.floatToRawIntBits(value)); + return this; + } + + @Override + public ByteBuf writeDouble(double value) { + writeLong(Double.doubleToRawLongBits(value)); + return this; + } + + private void setShort_(int index, int value) { + PlatformDependent.putShort(addr(index), (short) value); + } + + private void setInt_(int index, int value) { + PlatformDependent.putInt(addr(index), value); + } + + private void setLong_(int index, long value) { + PlatformDependent.putLong(addr(index), value); + } + + @Override + public byte getByte(int index) { + return PlatformDependent.getByte(addr(index)); + } + + @Override + public ByteBuf setByte(int index, int value) { + PlatformDependent.putByte(addr(index), (byte) value); + return this; + } + + @Override + public boolean release() { + return release(1); + } + + @Override + public int setBytes(int index, InputStream in, int length) throws IOException { + wrapped.checkIndex(index, length); + byte[] tmp = new byte[length]; + int readBytes = in.read(tmp); + if (readBytes > 0) { + PlatformDependent.copyMemory(tmp, 0, addr(index), readBytes); + } + return readBytes; + } + + @Override + public ByteBuf getBytes(int index, OutputStream out, int length) throws IOException { + wrapped.checkIndex(index, length); + if (length != 0) { + byte[] tmp = new byte[length]; + PlatformDependent.copyMemory(addr(index), tmp, 0, length); + out.write(tmp); + } + return this; + } + + @Override + public int hashCode() { + return System.identityHashCode(this); + } + + @Override + public boolean equals(Object obj) { + return this == obj; + } +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/ArrowByteBufAllocator.java b/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/ArrowByteBufAllocator.java new file mode 100644 index 000000000..ff40b49ff --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/ArrowByteBufAllocator.java @@ -0,0 +1,161 @@ +/* + * 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 io.netty.buffer.AbstractByteBufAllocator; +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.buffer.ExpandableByteBuf; +import io.netty.buffer.NettyArrowBuf; + +/** + * An implementation of ByteBufAllocator that wraps a Arrow BufferAllocator. This allows the RPC + * layer to be accounted + * and managed using Arrow's BufferAllocator infrastructure. The only thin different from a + * typical BufferAllocator is + * the signature and the fact that this Allocator returns ExpandableByteBufs which enable + * otherwise non-expandable + * ArrowBufs to be expandable. + * + * @deprecated This class may be removed in a future release. + */ +@Deprecated +public class ArrowByteBufAllocator extends AbstractByteBufAllocator { + + private static final int DEFAULT_BUFFER_SIZE = 4096; + private static final int DEFAULT_MAX_COMPOSITE_COMPONENTS = 16; + + private final BufferAllocator allocator; + + public ArrowByteBufAllocator(BufferAllocator allocator) { + this.allocator = allocator; + } + + public BufferAllocator unwrap() { + return allocator; + } + + @Override + public ByteBuf buffer() { + return buffer(DEFAULT_BUFFER_SIZE); + } + + @Override + public ByteBuf buffer(int initialCapacity) { + return new ExpandableByteBuf(NettyArrowBuf.unwrapBuffer(allocator.buffer(initialCapacity)), allocator); + } + + @Override + public ByteBuf buffer(int initialCapacity, int maxCapacity) { + return buffer(initialCapacity); + } + + @Override + public ByteBuf ioBuffer() { + return buffer(); + } + + @Override + public ByteBuf ioBuffer(int initialCapacity) { + return buffer(initialCapacity); + } + + @Override + public ByteBuf ioBuffer(int initialCapacity, int maxCapacity) { + return buffer(initialCapacity); + } + + @Override + public ByteBuf directBuffer() { + return buffer(); + } + + @Override + public ByteBuf directBuffer(int initialCapacity) { + return NettyArrowBuf.unwrapBuffer(allocator.buffer(initialCapacity)); + } + + @Override + public ByteBuf directBuffer(int initialCapacity, int maxCapacity) { + return buffer(initialCapacity, maxCapacity); + } + + @Override + public CompositeByteBuf compositeBuffer() { + return compositeBuffer(DEFAULT_MAX_COMPOSITE_COMPONENTS); + } + + @Override + public CompositeByteBuf compositeBuffer(int maxNumComponents) { + return new CompositeByteBuf(this, true, maxNumComponents); + } + + @Override + public CompositeByteBuf compositeDirectBuffer() { + return compositeBuffer(); + } + + @Override + public CompositeByteBuf compositeDirectBuffer(int maxNumComponents) { + return compositeBuffer(maxNumComponents); + } + + @Override + public boolean isDirectBufferPooled() { + return false; + } + + @Override + public ByteBuf heapBuffer() { + throw fail(); + } + + @Override + public ByteBuf heapBuffer(int initialCapacity) { + throw fail(); + } + + @Override + public ByteBuf heapBuffer(int initialCapacity, int maxCapacity) { + throw fail(); + } + + @Override + public CompositeByteBuf compositeHeapBuffer() { + throw fail(); + } + + @Override + public CompositeByteBuf compositeHeapBuffer(int maxNumComponents) { + throw fail(); + } + + @Override + protected ByteBuf newHeapBuffer(int initialCapacity, int maxCapacity) { + throw fail(); + } + + @Override + protected ByteBuf newDirectBuffer(int initialCapacity, int maxCapacity) { + return buffer(initialCapacity, maxCapacity); + } + + private RuntimeException fail() { + throw new UnsupportedOperationException("Allocator doesn't support heap-based memory."); + } +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java b/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java new file mode 100644 index 000000000..10cfb5c16 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java @@ -0,0 +1,38 @@ +/* + * 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; + +/** + * The default Allocation Manager Factory for a module. + * + */ +public class DefaultAllocationManagerFactory implements AllocationManager.Factory { + + public static final AllocationManager.Factory FACTORY = NettyAllocationManager.FACTORY; + + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long size) { + return FACTORY.create(accountingAllocator, size); + } + + @Override + public ArrowBuf empty() { + return FACTORY.empty(); + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java b/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java new file mode 100644 index 000000000..200047783 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/main/java/org/apache/arrow/memory/NettyAllocationManager.java @@ -0,0 +1,123 @@ +/* + * 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 io.netty.buffer.PooledByteBufAllocatorL; +import io.netty.buffer.UnsafeDirectLittleEndian; +import io.netty.util.internal.PlatformDependent; + +/** + * The default implementation of {@link AllocationManager}. The implementation is responsible for managing when memory + * is allocated and returned to the Netty-based PooledByteBufAllocatorL. + */ +public class NettyAllocationManager extends AllocationManager { + + public static final AllocationManager.Factory FACTORY = new AllocationManager.Factory() { + + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long size) { + return new NettyAllocationManager(accountingAllocator, size); + } + + @Override + public ArrowBuf empty() { + return EMPTY_BUFFER; + } + }; + + /** + * The default cut-off value for switching allocation strategies. + * If the request size is not greater than the cut-off value, we will allocate memory by + * {@link PooledByteBufAllocatorL} APIs, + * otherwise, we will use {@link PlatformDependent} APIs. + */ + public static final int DEFAULT_ALLOCATION_CUTOFF_VALUE = Integer.MAX_VALUE; + + private static final PooledByteBufAllocatorL INNER_ALLOCATOR = new PooledByteBufAllocatorL(); + static final UnsafeDirectLittleEndian EMPTY = INNER_ALLOCATOR.empty; + static final ArrowBuf EMPTY_BUFFER = new ArrowBuf(ReferenceManager.NO_OP, + null, + 0, + NettyAllocationManager.EMPTY.memoryAddress()); + static final long CHUNK_SIZE = INNER_ALLOCATOR.getChunkSize(); + + private final long allocatedSize; + private final UnsafeDirectLittleEndian memoryChunk; + private final long allocatedAddress; + + /** + * The cut-off value for switching allocation strategies. + */ + private final int allocationCutOffValue; + + NettyAllocationManager(BufferAllocator accountingAllocator, long requestedSize, int allocationCutOffValue) { + super(accountingAllocator); + this.allocationCutOffValue = allocationCutOffValue; + + if (requestedSize > allocationCutOffValue) { + this.memoryChunk = null; + this.allocatedAddress = PlatformDependent.allocateMemory(requestedSize); + this.allocatedSize = requestedSize; + } else { + this.memoryChunk = INNER_ALLOCATOR.allocate(requestedSize); + this.allocatedAddress = memoryChunk.memoryAddress(); + this.allocatedSize = memoryChunk.capacity(); + } + } + + NettyAllocationManager(BufferAllocator accountingAllocator, long requestedSize) { + this(accountingAllocator, requestedSize, DEFAULT_ALLOCATION_CUTOFF_VALUE); + } + + /** + * Get the underlying memory chunk managed by this AllocationManager. + * @return the underlying memory chunk if the request size is not greater than the + * {@link NettyAllocationManager#allocationCutOffValue}, or null otherwise. + * + * @deprecated this method will be removed in a future release. + */ + @Deprecated + UnsafeDirectLittleEndian getMemoryChunk() { + return memoryChunk; + } + + @Override + protected long memoryAddress() { + return allocatedAddress; + } + + @Override + protected void release0() { + if (memoryChunk == null) { + PlatformDependent.freeMemory(allocatedAddress); + } else { + memoryChunk.release(); + } + } + + /** + * Returns the underlying memory chunk size managed. + * + * <p>NettyAllocationManager rounds requested size up to the next power of two. + */ + @Override + public long getSize() { + return allocatedSize; + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/io/netty/buffer/TestNettyArrowBuf.java b/src/arrow/java/memory/memory-netty/src/test/java/io/netty/buffer/TestNettyArrowBuf.java new file mode 100644 index 000000000..916cf82e7 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/io/netty/buffer/TestNettyArrowBuf.java @@ -0,0 +1,141 @@ +/* + * 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 io.netty.buffer; + +import java.nio.ByteBuffer; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.ArrowByteBufAllocator; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.Assert; +import org.junit.Test; + +public class TestNettyArrowBuf { + + @Test + public void testSliceWithoutArgs() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); + nettyBuf.writerIndex(20); + nettyBuf.readerIndex(10); + NettyArrowBuf slicedBuffer = nettyBuf.slice(); + int readableBytes = slicedBuffer.readableBytes(); + Assert.assertEquals(10, readableBytes); + } + } + + @Test + public void testNioBuffer() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); + ByteBuffer byteBuffer = nettyBuf.nioBuffer(4, 6); + // Nio Buffers should always be 0 indexed + Assert.assertEquals(0, byteBuffer.position()); + Assert.assertEquals(6, byteBuffer.limit()); + // Underlying buffer has size 32 excluding 4 should have capacity of 28. + Assert.assertEquals(28, byteBuffer.capacity()); + + } + } + + @Test + public void testInternalNioBuffer() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); + ByteBuffer byteBuffer = nettyBuf.internalNioBuffer(4, 6); + Assert.assertEquals(0, byteBuffer.position()); + Assert.assertEquals(6, byteBuffer.limit()); + // Underlying buffer has size 32 excluding 4 should have capacity of 28. + Assert.assertEquals(28, byteBuffer.capacity()); + + } + } + + @Test + public void testSetLEValues() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf nettyBuf = NettyArrowBuf.unwrapBuffer(buf); + int [] intVals = new int[] {Integer.MIN_VALUE, Short.MIN_VALUE - 1, Short.MIN_VALUE, 0 , + Short.MAX_VALUE , Short.MAX_VALUE + 1, Integer.MAX_VALUE}; + for (int intValue :intVals ) { + nettyBuf._setInt(0, intValue); + Assert.assertEquals(nettyBuf._getIntLE(0), Integer.reverseBytes(intValue)); + } + + long [] longVals = new long[] {Long.MIN_VALUE, 0 , Long.MAX_VALUE}; + for (long longValue :longVals ) { + nettyBuf._setLong(0, longValue); + Assert.assertEquals(nettyBuf._getLongLE(0), Long.reverseBytes(longValue)); + } + + short [] shortVals = new short[] {Short.MIN_VALUE, 0 , Short.MAX_VALUE}; + for (short shortValue :shortVals ) { + nettyBuf._setShort(0, shortValue); + Assert.assertEquals(nettyBuf._getShortLE(0), Short.reverseBytes(shortValue)); + } + } + } + + @Test + public void testSetCompositeBuffer() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + NettyArrowBuf buf2 = NettyArrowBuf.unwrapBuffer(allocator.buffer(20)); + ) { + CompositeByteBuf byteBufs = new CompositeByteBuf(new ArrowByteBufAllocator(allocator), + true, 1); + int expected = 4; + buf2.setInt(0, expected); + buf2.writerIndex(4); + byteBufs.addComponent(true, buf2); + NettyArrowBuf.unwrapBuffer(buf).setBytes(0, byteBufs, 4); + int actual = buf.getInt(0); + Assert.assertEquals(expected, actual); + } + } + + @Test + public void testGetCompositeBuffer() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + CompositeByteBuf byteBufs = new CompositeByteBuf(new ArrowByteBufAllocator(allocator), + true, 1); + int expected = 4; + buf.setInt(0, expected); + NettyArrowBuf buf2 = NettyArrowBuf.unwrapBuffer(allocator.buffer(20)); + // composite buffers are a bit weird, need to jump hoops + // to set capacity. + byteBufs.addComponent(true, buf2); + byteBufs.capacity(20); + NettyArrowBuf.unwrapBuffer(buf).getBytes(0, byteBufs, 4); + int actual = byteBufs.getInt(0); + Assert.assertEquals(expected, actual); + byteBufs.component(0).release(); + } + } +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/io/netty/buffer/TestUnsafeDirectLittleEndian.java b/src/arrow/java/memory/memory-netty/src/test/java/io/netty/buffer/TestUnsafeDirectLittleEndian.java new file mode 100644 index 000000000..c2bd95bb3 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/io/netty/buffer/TestUnsafeDirectLittleEndian.java @@ -0,0 +1,77 @@ +/* + * 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 io.netty.buffer; + + +import static org.junit.Assert.assertEquals; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.ByteOrder; +import java.nio.charset.StandardCharsets; + +import org.junit.Test; + +public class TestUnsafeDirectLittleEndian { + private static final boolean LITTLE_ENDIAN = ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN; + + @Test + public void testPrimitiveGetSet() { + ByteBuf byteBuf = Unpooled.directBuffer(64); + UnsafeDirectLittleEndian unsafeDirect = new UnsafeDirectLittleEndian(new LargeBuffer(byteBuf)); + + unsafeDirect.setByte(0, Byte.MAX_VALUE); + unsafeDirect.setByte(1, -1); // 0xFF + unsafeDirect.setShort(2, Short.MAX_VALUE); + unsafeDirect.setShort(4, -2); // 0xFFFE + unsafeDirect.setInt(8, Integer.MAX_VALUE); + unsafeDirect.setInt(12, -66052); // 0xFFFE FDFC + unsafeDirect.setLong(16, Long.MAX_VALUE); + unsafeDirect.setLong(24, -4295098372L); // 0xFFFF FFFE FFFD FFFC + unsafeDirect.setFloat(32, 1.23F); + unsafeDirect.setFloat(36, -1.23F); + unsafeDirect.setDouble(40, 1.234567D); + unsafeDirect.setDouble(48, -1.234567D); + + assertEquals(Byte.MAX_VALUE, unsafeDirect.getByte(0)); + assertEquals(-1, unsafeDirect.getByte(1)); + assertEquals(Short.MAX_VALUE, unsafeDirect.getShort(2)); + assertEquals(-2, unsafeDirect.getShort(4)); + assertEquals((char) 65534, unsafeDirect.getChar(4)); + assertEquals(Integer.MAX_VALUE, unsafeDirect.getInt(8)); + assertEquals(-66052, unsafeDirect.getInt(12)); + assertEquals(4294901244L, unsafeDirect.getUnsignedInt(12)); + assertEquals(Long.MAX_VALUE, unsafeDirect.getLong(16)); + assertEquals(-4295098372L, unsafeDirect.getLong(24)); + assertEquals(1.23F, unsafeDirect.getFloat(32), 0.0); + assertEquals(-1.23F, unsafeDirect.getFloat(36), 0.0); + assertEquals(1.234567D, unsafeDirect.getDouble(40), 0.0); + assertEquals(-1.234567D, unsafeDirect.getDouble(48), 0.0); + + byte[] inBytes = "1234567".getBytes(StandardCharsets.UTF_8); + try (ByteArrayInputStream bais = new ByteArrayInputStream(inBytes); + ByteArrayOutputStream baos = new ByteArrayOutputStream()) { + assertEquals(5, unsafeDirect.setBytes(56, bais, 5)); + unsafeDirect.getBytes(56, baos, 5); + assertEquals("12345", new String(baos.toByteArray(), StandardCharsets.UTF_8)); + } catch (IOException e) { + e.printStackTrace(); + } + } +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/ITTestLargeArrowBuf.java b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/ITTestLargeArrowBuf.java new file mode 100644 index 000000000..fa8d510e3 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/ITTestLargeArrowBuf.java @@ -0,0 +1,72 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Integration test for large (more than 2GB) {@link org.apache.arrow.memory.ArrowBuf}. + * To run this test, please make sure there is at least 4GB memory in the system. + */ +public class ITTestLargeArrowBuf { + private static final Logger logger = LoggerFactory.getLogger(ITTestLargeArrowBuf.class); + + private void run(long bufSize) { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ArrowBuf largeBuf = allocator.buffer(bufSize)) { + assertEquals(bufSize, largeBuf.capacity()); + logger.trace("Successfully allocated a buffer with capacity {}", largeBuf.capacity()); + + for (long i = 0; i < bufSize / 8; i++) { + largeBuf.setLong(i * 8, i); + + if ((i + 1) % 10000 == 0) { + logger.trace("Successfully written {} long words", i + 1); + } + } + logger.trace("Successfully written {} long words", bufSize / 8); + + for (long i = 0; i < bufSize / 8; i++) { + long val = largeBuf.getLong(i * 8); + assertEquals(i, val); + + if ((i + 1) % 10000 == 0) { + logger.trace("Successfully read {} long words", i + 1); + } + } + logger.trace("Successfully read {} long words", bufSize / 8); + } + logger.trace("Successfully released the large buffer."); + } + + @Test + public void testLargeArrowBuf() { + run(4 * 1024 * 1024 * 1024L); + } + + @Test + public void testMaxIntArrowBuf() { + run(Integer.MAX_VALUE); + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestAllocationManagerNetty.java b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestAllocationManagerNetty.java new file mode 100644 index 000000000..2dbd56480 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestAllocationManagerNetty.java @@ -0,0 +1,39 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.Test; + +/** + * Test cases for {@link AllocationManager}. + */ +public class TestAllocationManagerNetty { + + @Test + public void testAllocationManagerType() { + // test netty allocation manager type + System.setProperty( + DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME, "Netty"); + DefaultAllocationManagerOption.AllocationManagerType mgrType = + DefaultAllocationManagerOption.getDefaultAllocationManagerType(); + + assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Netty, mgrType); + } +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java new file mode 100644 index 000000000..ef49e4178 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestBaseAllocator.java @@ -0,0 +1,1183 @@ +/* + * 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.Assert.assertEquals; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.lang.reflect.Field; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; + +import org.apache.arrow.memory.AllocationOutcomeDetails.Entry; +import org.apache.arrow.memory.rounding.RoundingPolicy; +import org.apache.arrow.memory.rounding.SegmentRoundingPolicy; +import org.apache.arrow.memory.util.AssertionUtil; +import org.junit.Ignore; +import org.junit.Test; +import org.junit.jupiter.api.Assertions; + +import sun.misc.Unsafe; + +public class TestBaseAllocator { + + private static final int MAX_ALLOCATION = 8 * 1024; + + /* + // ---------------------------------------- DEBUG ----------------------------------- + + @After + public void checkBuffers() { + final int bufferCount = UnsafeDirectLittleEndian.getBufferCount(); + if (bufferCount != 0) { + UnsafeDirectLittleEndian.logBuffers(logger); + UnsafeDirectLittleEndian.releaseBuffers(); + } + + assertEquals(0, bufferCount); + } + + // @AfterClass + // public static void dumpBuffers() { + // UnsafeDirectLittleEndian.logBuffers(logger); + // } + + // ---------------------------------------- DEBUG ------------------------------------ + */ + + + @Test + public void test_privateMax() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + final ArrowBuf arrowBuf1 = rootAllocator.buffer(MAX_ALLOCATION / 2); + assertNotNull("allocation failed", arrowBuf1); + + try (final BufferAllocator childAllocator = + rootAllocator.newChildAllocator("noLimits", 0, MAX_ALLOCATION)) { + final ArrowBuf arrowBuf2 = childAllocator.buffer(MAX_ALLOCATION / 2); + assertNotNull("allocation failed", arrowBuf2); + arrowBuf2.getReferenceManager().release(); + } + + arrowBuf1.getReferenceManager().release(); + } + } + + @Test(expected = IllegalStateException.class) + public void testRootAllocator_closeWithOutstanding() throws Exception { + try { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + final ArrowBuf arrowBuf = rootAllocator.buffer(512); + assertNotNull("allocation failed", arrowBuf); + } + } finally { + /* + * We expect there to be one unreleased underlying buffer because we're closing + * without releasing it. + */ + /* + // ------------------------------- DEBUG --------------------------------- + final int bufferCount = UnsafeDirectLittleEndian.getBufferCount(); + UnsafeDirectLittleEndian.releaseBuffers(); + assertEquals(1, bufferCount); + // ------------------------------- DEBUG --------------------------------- + */ + } + } + + @Test + public void testRootAllocator_getEmpty() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + final ArrowBuf arrowBuf = rootAllocator.buffer(0); + assertNotNull("allocation failed", arrowBuf); + assertEquals("capacity was non-zero", 0, arrowBuf.capacity()); + assertTrue("address should be valid", arrowBuf.memoryAddress() != 0); + arrowBuf.getReferenceManager().release(); + } + } + + @Ignore // TODO(DRILL-2740) + @Test(expected = IllegalStateException.class) + public void testAllocator_unreleasedEmpty() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + @SuppressWarnings("unused") + final ArrowBuf arrowBuf = rootAllocator.buffer(0); + } + } + + @Test + public void testAllocator_transferOwnership() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + final BufferAllocator childAllocator1 = + rootAllocator.newChildAllocator("changeOwnership1", 0, MAX_ALLOCATION); + final BufferAllocator childAllocator2 = + rootAllocator.newChildAllocator("changeOwnership2", 0, MAX_ALLOCATION); + + final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 4); + rootAllocator.verify(); + final ReferenceManager referenceManager = arrowBuf1.getReferenceManager(); + OwnershipTransferResult transferOwnership = referenceManager.transferOwnership(arrowBuf1, childAllocator2); + assertEquiv(arrowBuf1, transferOwnership.getTransferredBuffer()); + final boolean allocationFit = transferOwnership.getAllocationFit(); + rootAllocator.verify(); + assertTrue(allocationFit); + + arrowBuf1.getReferenceManager().release(); + childAllocator1.close(); + rootAllocator.verify(); + + transferOwnership.getTransferredBuffer().getReferenceManager().release(); + childAllocator2.close(); + } + } + + static <T> boolean equalsIgnoreOrder(Collection<T> c1, Collection<T> c2) { + return (c1.size() == c2.size() && c1.containsAll(c2)); + } + + @Test + public void testAllocator_getParentAndChild() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + assertEquals(rootAllocator.getParentAllocator(), null); + + try (final BufferAllocator childAllocator1 = + rootAllocator.newChildAllocator("child1", 0, MAX_ALLOCATION)) { + assertEquals(childAllocator1.getParentAllocator(), rootAllocator); + assertTrue( + equalsIgnoreOrder(Arrays.asList(childAllocator1), rootAllocator.getChildAllocators())); + + try (final BufferAllocator childAllocator2 = + rootAllocator.newChildAllocator("child2", 0, MAX_ALLOCATION)) { + assertEquals(childAllocator2.getParentAllocator(), rootAllocator); + assertTrue(equalsIgnoreOrder(Arrays.asList(childAllocator1, childAllocator2), + rootAllocator.getChildAllocators())); + + try (final BufferAllocator grandChildAllocator = + childAllocator1.newChildAllocator("grand-child", 0, MAX_ALLOCATION)) { + assertEquals(grandChildAllocator.getParentAllocator(), childAllocator1); + assertTrue(equalsIgnoreOrder(Arrays.asList(grandChildAllocator), + childAllocator1.getChildAllocators())); + } + } + } + } + } + + @Test + public void testAllocator_childRemovedOnClose() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator1 = + rootAllocator.newChildAllocator("child1", 0, MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator2 = + rootAllocator.newChildAllocator("child2", 0, MAX_ALLOCATION)) { + + // root has two child allocators + assertTrue(equalsIgnoreOrder(Arrays.asList(childAllocator1, childAllocator2), + rootAllocator.getChildAllocators())); + + try (final BufferAllocator grandChildAllocator = + childAllocator1.newChildAllocator("grand-child", 0, MAX_ALLOCATION)) { + + // child1 has one allocator i.e grand-child + assertTrue(equalsIgnoreOrder(Arrays.asList(grandChildAllocator), + childAllocator1.getChildAllocators())); + } + + // grand-child closed + assertTrue( + equalsIgnoreOrder(Collections.EMPTY_SET, childAllocator1.getChildAllocators())); + } + // root has only one child left + assertTrue( + equalsIgnoreOrder(Arrays.asList(childAllocator1), rootAllocator.getChildAllocators())); + } + // all child allocators closed. + assertTrue(equalsIgnoreOrder(Collections.EMPTY_SET, rootAllocator.getChildAllocators())); + } + } + + @Test + public void testAllocator_shareOwnership() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("shareOwnership1", 0, + MAX_ALLOCATION); + final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("shareOwnership2", 0, + MAX_ALLOCATION); + final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 4); + rootAllocator.verify(); + + // share ownership of buffer. + final ArrowBuf arrowBuf2 = arrowBuf1.getReferenceManager().retain(arrowBuf1, childAllocator2); + rootAllocator.verify(); + assertNotNull(arrowBuf2); + assertNotEquals(arrowBuf2, arrowBuf1); + assertEquiv(arrowBuf1, arrowBuf2); + + // release original buffer (thus transferring ownership to allocator 2. (should leave + // allocator 1 in empty state) + arrowBuf1.getReferenceManager().release(); + rootAllocator.verify(); + childAllocator1.close(); + rootAllocator.verify(); + + final BufferAllocator childAllocator3 = rootAllocator.newChildAllocator("shareOwnership3", 0, + MAX_ALLOCATION); + final ArrowBuf arrowBuf3 = arrowBuf1.getReferenceManager().retain(arrowBuf1, childAllocator3); + assertNotNull(arrowBuf3); + assertNotEquals(arrowBuf3, arrowBuf1); + assertNotEquals(arrowBuf3, arrowBuf2); + assertEquiv(arrowBuf1, arrowBuf3); + rootAllocator.verify(); + + arrowBuf2.getReferenceManager().release(); + rootAllocator.verify(); + childAllocator2.close(); + rootAllocator.verify(); + + arrowBuf3.getReferenceManager().release(); + rootAllocator.verify(); + childAllocator3.close(); + } + } + + @Test + public void testRootAllocator_createChildAndUse() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator( + "createChildAndUse", 0, MAX_ALLOCATION)) { + final ArrowBuf arrowBuf = childAllocator.buffer(512); + assertNotNull("allocation failed", arrowBuf); + arrowBuf.getReferenceManager().release(); + } + } + } + + @Test(expected = IllegalStateException.class) + public void testRootAllocator_createChildDontClose() throws Exception { + try { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + final BufferAllocator childAllocator = rootAllocator.newChildAllocator( + "createChildDontClose", 0, MAX_ALLOCATION); + final ArrowBuf arrowBuf = childAllocator.buffer(512); + assertNotNull("allocation failed", arrowBuf); + } + } finally { + /* + * We expect one underlying buffer because we closed a child allocator without + * releasing the buffer allocated from it. + */ + /* + // ------------------------------- DEBUG --------------------------------- + final int bufferCount = UnsafeDirectLittleEndian.getBufferCount(); + UnsafeDirectLittleEndian.releaseBuffers(); + assertEquals(1, bufferCount); + // ------------------------------- DEBUG --------------------------------- + */ + } + } + + @Test + public void testSegmentAllocator() { + RoundingPolicy policy = new SegmentRoundingPolicy(1024); + try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, 1024 * 1024, policy)) { + ArrowBuf buf = allocator.buffer(798); + assertEquals(1024, buf.capacity()); + buf.setInt(333, 959); + assertEquals(959, buf.getInt(333)); + buf.close(); + + buf = allocator.buffer(1025); + assertEquals(2048, buf.capacity()); + buf.setInt(193, 939); + assertEquals(939, buf.getInt(193)); + buf.close(); + } + } + + @Test + public void testSegmentAllocator_childAllocator() { + RoundingPolicy policy = new SegmentRoundingPolicy(1024); + try (RootAllocator allocator = new RootAllocator(AllocationListener.NOOP, 1024 * 1024, policy); + BufferAllocator childAllocator = allocator.newChildAllocator("child", 0, 512 * 1024)) { + + assertEquals("child", childAllocator.getName()); + + ArrowBuf buf = childAllocator.buffer(798); + assertEquals(1024, buf.capacity()); + buf.setInt(333, 959); + assertEquals(959, buf.getInt(333)); + buf.close(); + + buf = childAllocator.buffer(1025); + assertEquals(2048, buf.capacity()); + buf.setInt(193, 939); + assertEquals(939, buf.getInt(193)); + buf.close(); + } + } + + @Test + public void testSegmentAllocator_smallSegment() { + IllegalArgumentException e = Assertions.assertThrows( + IllegalArgumentException.class, + () -> new SegmentRoundingPolicy(128) + ); + assertEquals("The segment size cannot be smaller than 1024", e.getMessage()); + } + + @Test + public void testSegmentAllocator_segmentSizeNotPowerOf2() { + IllegalArgumentException e = Assertions.assertThrows( + IllegalArgumentException.class, + () -> new SegmentRoundingPolicy(4097) + ); + assertEquals("The segment size must be a power of 2", e.getMessage()); + } + + @Test + public void testCustomizedAllocationManager() { + try (BaseAllocator allocator = createAllocatorWithCustomizedAllocationManager()) { + final ArrowBuf arrowBuf1 = allocator.buffer(MAX_ALLOCATION); + assertNotNull("allocation failed", arrowBuf1); + + arrowBuf1.setInt(0, 1); + assertEquals(1, arrowBuf1.getInt(0)); + + try { + final ArrowBuf arrowBuf2 = allocator.buffer(1); + fail("allocated memory beyond max allowed"); + } catch (OutOfMemoryException e) { + // expected + } + arrowBuf1.getReferenceManager().release(); + + try { + arrowBuf1.getInt(0); + fail("data read from released buffer"); + } catch (RuntimeException e) { + // expected + } + } + } + + private BaseAllocator createAllocatorWithCustomizedAllocationManager() { + return new RootAllocator(BaseAllocator.configBuilder() + .maxAllocation(MAX_ALLOCATION) + .allocationManagerFactory(new AllocationManager.Factory() { + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long requestedSize) { + return new AllocationManager(accountingAllocator) { + private final Unsafe unsafe = getUnsafe(); + private final long address = unsafe.allocateMemory(requestedSize); + + @Override + protected long memoryAddress() { + return address; + } + + @Override + protected void release0() { + unsafe.setMemory(address, requestedSize, (byte) 0); + unsafe.freeMemory(address); + } + + @Override + public long getSize() { + return requestedSize; + } + + private Unsafe getUnsafe() { + Field f = null; + try { + f = Unsafe.class.getDeclaredField("theUnsafe"); + f.setAccessible(true); + return (Unsafe) f.get(null); + } catch (NoSuchFieldException | IllegalAccessException e) { + throw new RuntimeException(e); + } finally { + if (f != null) { + f.setAccessible(false); + } + } + } + }; + } + + @Override + public ArrowBuf empty() { + return null; + } + }).build()); + } + + // Allocation listener + // It counts the number of times it has been invoked, and how much memory allocation it has seen + // When set to 'expand on fail', it attempts to expand the associated allocator's limit + private static final class TestAllocationListener implements AllocationListener { + private int numPreCalls; + private int numCalls; + private int numReleaseCalls; + private int numChildren; + private long totalMem; + private boolean expandOnFail; + BufferAllocator expandAlloc; + long expandLimit; + + TestAllocationListener() { + this.numCalls = 0; + this.numChildren = 0; + this.totalMem = 0; + this.expandOnFail = false; + this.expandAlloc = null; + this.expandLimit = 0; + } + + @Override + public void onPreAllocation(long size) { + numPreCalls++; + } + + @Override + public void onAllocation(long size) { + numCalls++; + totalMem += size; + } + + @Override + public boolean onFailedAllocation(long size, AllocationOutcome outcome) { + if (expandOnFail) { + expandAlloc.setLimit(expandLimit); + return true; + } + return false; + } + + + @Override + public void onRelease(long size) { + numReleaseCalls++; + } + + @Override + public void onChildAdded(BufferAllocator parentAllocator, BufferAllocator childAllocator) { + ++numChildren; + } + + @Override + public void onChildRemoved(BufferAllocator parentAllocator, BufferAllocator childAllocator) { + --numChildren; + } + + void setExpandOnFail(BufferAllocator expandAlloc, long expandLimit) { + this.expandOnFail = true; + this.expandAlloc = expandAlloc; + this.expandLimit = expandLimit; + } + + int getNumPreCalls() { + return numPreCalls; + } + + int getNumReleaseCalls() { + return numReleaseCalls; + } + + int getNumCalls() { + return numCalls; + } + + int getNumChildren() { + return numChildren; + } + + long getTotalMem() { + return totalMem; + } + } + + @Test + public void testRootAllocator_listeners() throws Exception { + TestAllocationListener l1 = new TestAllocationListener(); + assertEquals(0, l1.getNumPreCalls()); + assertEquals(0, l1.getNumCalls()); + assertEquals(0, l1.getNumReleaseCalls()); + assertEquals(0, l1.getNumChildren()); + assertEquals(0, l1.getTotalMem()); + TestAllocationListener l2 = new TestAllocationListener(); + assertEquals(0, l2.getNumPreCalls()); + assertEquals(0, l2.getNumCalls()); + assertEquals(0, l2.getNumReleaseCalls()); + assertEquals(0, l2.getNumChildren()); + assertEquals(0, l2.getTotalMem()); + // root and first-level child share the first listener + // second-level and third-level child share the second listener + try (final RootAllocator rootAllocator = new RootAllocator(l1, MAX_ALLOCATION)) { + try (final BufferAllocator c1 = rootAllocator.newChildAllocator("c1", 0, MAX_ALLOCATION)) { + assertEquals(1, l1.getNumChildren()); + final ArrowBuf buf1 = c1.buffer(16); + assertNotNull("allocation failed", buf1); + assertEquals(1, l1.getNumPreCalls()); + assertEquals(1, l1.getNumCalls()); + assertEquals(0, l1.getNumReleaseCalls()); + assertEquals(16, l1.getTotalMem()); + buf1.getReferenceManager().release(); + try (final BufferAllocator c2 = c1.newChildAllocator("c2", l2, 0, MAX_ALLOCATION)) { + assertEquals(2, l1.getNumChildren()); // c1 got a new child, so c1's listener (l1) is notified + assertEquals(0, l2.getNumChildren()); + final ArrowBuf buf2 = c2.buffer(32); + assertNotNull("allocation failed", buf2); + assertEquals(1, l1.getNumCalls()); + assertEquals(16, l1.getTotalMem()); + assertEquals(1, l2.getNumPreCalls()); + assertEquals(1, l2.getNumCalls()); + assertEquals(0, l2.getNumReleaseCalls()); + assertEquals(32, l2.getTotalMem()); + buf2.getReferenceManager().release(); + try (final BufferAllocator c3 = c2.newChildAllocator("c3", 0, MAX_ALLOCATION)) { + assertEquals(2, l1.getNumChildren()); + assertEquals(1, l2.getNumChildren()); + final ArrowBuf buf3 = c3.buffer(64); + assertNotNull("allocation failed", buf3); + assertEquals(1, l1.getNumPreCalls()); + assertEquals(1, l1.getNumCalls()); + assertEquals(1, l1.getNumReleaseCalls()); + assertEquals(16, l1.getTotalMem()); + assertEquals(2, l2.getNumPreCalls()); + assertEquals(2, l2.getNumCalls()); + assertEquals(1, l2.getNumReleaseCalls()); + assertEquals(32 + 64, l2.getTotalMem()); + buf3.getReferenceManager().release(); + } + assertEquals(2, l1.getNumChildren()); + assertEquals(0, l2.getNumChildren()); // third-level child removed + } + assertEquals(1, l1.getNumChildren()); // second-level child removed + assertEquals(0, l2.getNumChildren()); + } + assertEquals(0, l1.getNumChildren()); // first-level child removed + + assertEquals(2, l2.getNumReleaseCalls()); + } + } + + @Test + public void testRootAllocator_listenerAllocationFail() throws Exception { + TestAllocationListener l1 = new TestAllocationListener(); + assertEquals(0, l1.getNumCalls()); + assertEquals(0, l1.getTotalMem()); + // Test attempts to allocate too much from a child whose limit is set to half of the max + // allocation. The listener's callback triggers, expanding the child allocator's limit, so then + // the allocation succeeds. + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator c1 = rootAllocator.newChildAllocator("c1", l1, 0, + MAX_ALLOCATION / 2)) { + try { + c1.buffer(MAX_ALLOCATION); + fail("allocated memory beyond max allowed"); + } catch (OutOfMemoryException e) { + // expected + } + assertEquals(0, l1.getNumCalls()); + assertEquals(0, l1.getTotalMem()); + + l1.setExpandOnFail(c1, MAX_ALLOCATION); + ArrowBuf arrowBuf = c1.buffer(MAX_ALLOCATION); + assertNotNull("allocation failed", arrowBuf); + assertEquals(1, l1.getNumCalls()); + assertEquals(MAX_ALLOCATION, l1.getTotalMem()); + arrowBuf.getReferenceManager().release(); + } + } + } + + private static void allocateAndFree(final BufferAllocator allocator) { + final ArrowBuf arrowBuf = allocator.buffer(512); + assertNotNull("allocation failed", arrowBuf); + arrowBuf.getReferenceManager().release(); + + final ArrowBuf arrowBuf2 = allocator.buffer(MAX_ALLOCATION); + assertNotNull("allocation failed", arrowBuf2); + arrowBuf2.getReferenceManager().release(); + + final int nBufs = 8; + final ArrowBuf[] arrowBufs = new ArrowBuf[nBufs]; + for (int i = 0; i < arrowBufs.length; ++i) { + ArrowBuf arrowBufi = allocator.buffer(MAX_ALLOCATION / nBufs); + assertNotNull("allocation failed", arrowBufi); + arrowBufs[i] = arrowBufi; + } + for (ArrowBuf arrowBufi : arrowBufs) { + arrowBufi.getReferenceManager().release(); + } + } + + @Test + public void testAllocator_manyAllocations() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = + rootAllocator.newChildAllocator("manyAllocations", 0, MAX_ALLOCATION)) { + allocateAndFree(childAllocator); + } + } + } + + @Test + public void testAllocator_overAllocate() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = + rootAllocator.newChildAllocator("overAllocate", 0, MAX_ALLOCATION)) { + allocateAndFree(childAllocator); + + try { + childAllocator.buffer(MAX_ALLOCATION + 1); + fail("allocated memory beyond max allowed"); + } catch (OutOfMemoryException e) { + // expected + } + } + } + } + + @Test + public void testAllocator_overAllocateParent() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = + rootAllocator.newChildAllocator("overAllocateParent", 0, MAX_ALLOCATION)) { + final ArrowBuf arrowBuf1 = rootAllocator.buffer(MAX_ALLOCATION / 2); + assertNotNull("allocation failed", arrowBuf1); + final ArrowBuf arrowBuf2 = childAllocator.buffer(MAX_ALLOCATION / 2); + assertNotNull("allocation failed", arrowBuf2); + + try { + childAllocator.buffer(MAX_ALLOCATION / 4); + fail("allocated memory beyond max allowed"); + } catch (OutOfMemoryException e) { + // expected + } + + arrowBuf1.getReferenceManager().release(); + arrowBuf2.getReferenceManager().release(); + } + } + } + + @Test + public void testAllocator_failureAtParentLimitOutcomeDetails() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = + rootAllocator.newChildAllocator("child", 0, MAX_ALLOCATION / 2)) { + try (final BufferAllocator grandChildAllocator = + childAllocator.newChildAllocator("grandchild", MAX_ALLOCATION / 4, MAX_ALLOCATION)) { + OutOfMemoryException e = assertThrows(OutOfMemoryException.class, + () -> grandChildAllocator.buffer(MAX_ALLOCATION)); + // expected + assertTrue(e.getMessage().contains("Unable to allocate buffer")); + + assertTrue("missing outcome details", e.getOutcomeDetails().isPresent()); + AllocationOutcomeDetails outcomeDetails = e.getOutcomeDetails().get(); + + assertEquals(outcomeDetails.getFailedAllocator(), childAllocator); + + // The order of allocators should be child to root (request propagates to parent if + // child cannot satisfy the request). + Iterator<Entry> iterator = outcomeDetails.allocEntries.iterator(); + AllocationOutcomeDetails.Entry first = iterator.next(); + assertEquals(MAX_ALLOCATION / 4, first.getAllocatedSize()); + assertEquals(MAX_ALLOCATION, first.getRequestedSize()); + assertEquals(false, first.isAllocationFailed()); + + AllocationOutcomeDetails.Entry second = iterator.next(); + assertEquals(MAX_ALLOCATION - MAX_ALLOCATION / 4, second.getRequestedSize()); + assertEquals(0, second.getAllocatedSize()); + assertEquals(true, second.isAllocationFailed()); + + assertFalse(iterator.hasNext()); + } + } + } + } + + @Test + public void testAllocator_failureAtRootLimitOutcomeDetails() throws Exception { + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = + rootAllocator.newChildAllocator("child", MAX_ALLOCATION / 2, Long.MAX_VALUE)) { + try (final BufferAllocator grandChildAllocator = + childAllocator.newChildAllocator("grandchild", MAX_ALLOCATION / 4, Long.MAX_VALUE)) { + OutOfMemoryException e = assertThrows(OutOfMemoryException.class, + () -> grandChildAllocator.buffer(MAX_ALLOCATION * 2)); + + assertTrue(e.getMessage().contains("Unable to allocate buffer")); + assertTrue("missing outcome details", e.getOutcomeDetails().isPresent()); + AllocationOutcomeDetails outcomeDetails = e.getOutcomeDetails().get(); + + assertEquals(outcomeDetails.getFailedAllocator(), rootAllocator); + + // The order of allocators should be child to root (request propagates to parent if + // child cannot satisfy the request). + Iterator<Entry> iterator = outcomeDetails.allocEntries.iterator(); + AllocationOutcomeDetails.Entry first = iterator.next(); + assertEquals(MAX_ALLOCATION / 4, first.getAllocatedSize()); + assertEquals(2 * MAX_ALLOCATION, first.getRequestedSize()); + assertEquals(false, first.isAllocationFailed()); + + AllocationOutcomeDetails.Entry second = iterator.next(); + assertEquals(MAX_ALLOCATION / 4, second.getAllocatedSize()); + assertEquals(2 * MAX_ALLOCATION - MAX_ALLOCATION / 4, second.getRequestedSize()); + assertEquals(false, second.isAllocationFailed()); + + AllocationOutcomeDetails.Entry third = iterator.next(); + assertEquals(0, third.getAllocatedSize()); + assertEquals(true, third.isAllocationFailed()); + + assertFalse(iterator.hasNext()); + } + } + } + } + + private static void testAllocator_sliceUpBufferAndRelease( + final RootAllocator rootAllocator, final BufferAllocator bufferAllocator) { + final ArrowBuf arrowBuf1 = bufferAllocator.buffer(MAX_ALLOCATION / 2); + rootAllocator.verify(); + + final ArrowBuf arrowBuf2 = arrowBuf1.slice(16, arrowBuf1.capacity() - 32); + rootAllocator.verify(); + final ArrowBuf arrowBuf3 = arrowBuf2.slice(16, arrowBuf2.capacity() - 32); + rootAllocator.verify(); + @SuppressWarnings("unused") + final ArrowBuf arrowBuf4 = arrowBuf3.slice(16, arrowBuf3.capacity() - 32); + rootAllocator.verify(); + + arrowBuf3.getReferenceManager().release(); // since they share refcounts, one is enough to release them all + rootAllocator.verify(); + } + + @Test + public void testAllocator_createSlices() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + testAllocator_sliceUpBufferAndRelease(rootAllocator, rootAllocator); + + try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("createSlices", 0, + MAX_ALLOCATION)) { + testAllocator_sliceUpBufferAndRelease(rootAllocator, childAllocator); + } + rootAllocator.verify(); + + testAllocator_sliceUpBufferAndRelease(rootAllocator, rootAllocator); + + try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator("createSlices", 0, + MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator2 = + childAllocator.newChildAllocator("createSlices", 0, MAX_ALLOCATION)) { + final ArrowBuf arrowBuf1 = childAllocator2.buffer(MAX_ALLOCATION / 8); + @SuppressWarnings("unused") + final ArrowBuf arrowBuf2 = arrowBuf1.slice(MAX_ALLOCATION / 16, MAX_ALLOCATION / 16); + testAllocator_sliceUpBufferAndRelease(rootAllocator, childAllocator); + arrowBuf1.getReferenceManager().release(); + rootAllocator.verify(); + } + rootAllocator.verify(); + + testAllocator_sliceUpBufferAndRelease(rootAllocator, childAllocator); + } + rootAllocator.verify(); + } + } + + @Test + public void testAllocator_sliceRanges() throws Exception { + // final AllocatorOwner allocatorOwner = new NamedOwner("sliceRanges"); + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + // Populate a buffer with byte values corresponding to their indices. + final ArrowBuf arrowBuf = rootAllocator.buffer(256); + assertEquals(256, arrowBuf.capacity()); + assertEquals(0, arrowBuf.readerIndex()); + assertEquals(0, arrowBuf.readableBytes()); + assertEquals(0, arrowBuf.writerIndex()); + assertEquals(256, arrowBuf.writableBytes()); + + final ArrowBuf slice3 = arrowBuf.slice(); + assertEquals(0, slice3.readerIndex()); + assertEquals(0, slice3.readableBytes()); + assertEquals(0, slice3.writerIndex()); + // assertEquals(256, slice3.capacity()); + // assertEquals(256, slice3.writableBytes()); + + for (int i = 0; i < 256; ++i) { + arrowBuf.writeByte(i); + } + assertEquals(0, arrowBuf.readerIndex()); + assertEquals(256, arrowBuf.readableBytes()); + assertEquals(256, arrowBuf.writerIndex()); + assertEquals(0, arrowBuf.writableBytes()); + + final ArrowBuf slice1 = arrowBuf.slice(); + assertEquals(0, slice1.readerIndex()); + assertEquals(256, slice1.readableBytes()); + for (int i = 0; i < 10; ++i) { + assertEquals(i, slice1.readByte()); + } + assertEquals(256 - 10, slice1.readableBytes()); + for (int i = 0; i < 256; ++i) { + assertEquals((byte) i, slice1.getByte(i)); + } + + final ArrowBuf slice2 = arrowBuf.slice(25, 25); + assertEquals(0, slice2.readerIndex()); + assertEquals(25, slice2.readableBytes()); + for (int i = 25; i < 50; ++i) { + assertEquals(i, slice2.readByte()); + } + + /* + for(int i = 256; i > 0; --i) { + slice3.writeByte(i - 1); + } + for(int i = 0; i < 256; ++i) { + assertEquals(255 - i, slice1.getByte(i)); + } + */ + + arrowBuf.getReferenceManager().release(); // all the derived buffers share this fate + } + } + + @Test + public void testAllocator_slicesOfSlices() throws Exception { + // final AllocatorOwner allocatorOwner = new NamedOwner("slicesOfSlices"); + try (final RootAllocator rootAllocator = + new RootAllocator(MAX_ALLOCATION)) { + // Populate a buffer with byte values corresponding to their indices. + final ArrowBuf arrowBuf = rootAllocator.buffer(256); + for (int i = 0; i < 256; ++i) { + arrowBuf.writeByte(i); + } + + // Slice it up. + final ArrowBuf slice0 = arrowBuf.slice(0, arrowBuf.capacity()); + for (int i = 0; i < 256; ++i) { + assertEquals((byte) i, arrowBuf.getByte(i)); + } + + final ArrowBuf slice10 = slice0.slice(10, arrowBuf.capacity() - 10); + for (int i = 10; i < 256; ++i) { + assertEquals((byte) i, slice10.getByte(i - 10)); + } + + final ArrowBuf slice20 = slice10.slice(10, arrowBuf.capacity() - 20); + for (int i = 20; i < 256; ++i) { + assertEquals((byte) i, slice20.getByte(i - 20)); + } + + final ArrowBuf slice30 = slice20.slice(10, arrowBuf.capacity() - 30); + for (int i = 30; i < 256; ++i) { + assertEquals((byte) i, slice30.getByte(i - 30)); + } + + arrowBuf.getReferenceManager().release(); + } + } + + @Test + public void testAllocator_transferSliced() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("transferSliced1", 0, MAX_ALLOCATION); + final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("transferSliced2", 0, MAX_ALLOCATION); + + final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 8); + final ArrowBuf arrowBuf2 = childAllocator2.buffer(MAX_ALLOCATION / 8); + + final ArrowBuf arrowBuf1s = arrowBuf1.slice(0, arrowBuf1.capacity() / 2); + final ArrowBuf arrowBuf2s = arrowBuf2.slice(0, arrowBuf2.capacity() / 2); + + rootAllocator.verify(); + + OwnershipTransferResult result1 = arrowBuf2s.getReferenceManager().transferOwnership(arrowBuf2s, childAllocator1); + assertEquiv(arrowBuf2s, result1.getTransferredBuffer()); + rootAllocator.verify(); + OwnershipTransferResult result2 = arrowBuf1s.getReferenceManager().transferOwnership(arrowBuf1s, childAllocator2); + assertEquiv(arrowBuf1s, result2.getTransferredBuffer()); + rootAllocator.verify(); + + result1.getTransferredBuffer().getReferenceManager().release(); + result2.getTransferredBuffer().getReferenceManager().release(); + + arrowBuf1s.getReferenceManager().release(); // releases arrowBuf1 + arrowBuf2s.getReferenceManager().release(); // releases arrowBuf2 + + childAllocator1.close(); + childAllocator2.close(); + } + } + + @Test + public void testAllocator_shareSliced() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("transferSliced", 0, MAX_ALLOCATION); + final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("transferSliced", 0, MAX_ALLOCATION); + + final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 8); + final ArrowBuf arrowBuf2 = childAllocator2.buffer(MAX_ALLOCATION / 8); + + final ArrowBuf arrowBuf1s = arrowBuf1.slice(0, arrowBuf1.capacity() / 2); + final ArrowBuf arrowBuf2s = arrowBuf2.slice(0, arrowBuf2.capacity() / 2); + + rootAllocator.verify(); + + final ArrowBuf arrowBuf2s1 = arrowBuf2s.getReferenceManager().retain(arrowBuf2s, childAllocator1); + assertEquiv(arrowBuf2s, arrowBuf2s1); + final ArrowBuf arrowBuf1s2 = arrowBuf1s.getReferenceManager().retain(arrowBuf1s, childAllocator2); + assertEquiv(arrowBuf1s, arrowBuf1s2); + rootAllocator.verify(); + + arrowBuf1s.getReferenceManager().release(); // releases arrowBuf1 + arrowBuf2s.getReferenceManager().release(); // releases arrowBuf2 + rootAllocator.verify(); + + arrowBuf2s1.getReferenceManager().release(); // releases the shared arrowBuf2 slice + arrowBuf1s2.getReferenceManager().release(); // releases the shared arrowBuf1 slice + + childAllocator1.close(); + childAllocator2.close(); + } + } + + @Test + public void testAllocator_transferShared() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator("transferShared1", 0, MAX_ALLOCATION); + final BufferAllocator childAllocator2 = rootAllocator.newChildAllocator("transferShared2", 0, MAX_ALLOCATION); + final BufferAllocator childAllocator3 = rootAllocator.newChildAllocator("transferShared3", 0, MAX_ALLOCATION); + + final ArrowBuf arrowBuf1 = childAllocator1.buffer(MAX_ALLOCATION / 8); + + boolean allocationFit; + + ArrowBuf arrowBuf2 = arrowBuf1.getReferenceManager().retain(arrowBuf1, childAllocator2); + rootAllocator.verify(); + assertNotNull(arrowBuf2); + assertNotEquals(arrowBuf2, arrowBuf1); + assertEquiv(arrowBuf1, arrowBuf2); + + final ReferenceManager refManager1 = arrowBuf1.getReferenceManager(); + final OwnershipTransferResult result1 = refManager1.transferOwnership(arrowBuf1, childAllocator3); + allocationFit = result1.getAllocationFit(); + final ArrowBuf arrowBuf3 = result1.getTransferredBuffer(); + assertTrue(allocationFit); + assertEquiv(arrowBuf1, arrowBuf3); + rootAllocator.verify(); + + // Since childAllocator3 now has childAllocator1's buffer, 1, can close + arrowBuf1.getReferenceManager().release(); + childAllocator1.close(); + rootAllocator.verify(); + + arrowBuf2.getReferenceManager().release(); + childAllocator2.close(); + rootAllocator.verify(); + + final BufferAllocator childAllocator4 = rootAllocator.newChildAllocator("transferShared4", 0, MAX_ALLOCATION); + final ReferenceManager refManager3 = arrowBuf3.getReferenceManager(); + final OwnershipTransferResult result3 = refManager3.transferOwnership(arrowBuf3, childAllocator4); + allocationFit = result3.getAllocationFit(); + final ArrowBuf arrowBuf4 = result3.getTransferredBuffer(); + assertTrue(allocationFit); + assertEquiv(arrowBuf3, arrowBuf4); + rootAllocator.verify(); + + arrowBuf3.getReferenceManager().release(); + childAllocator3.close(); + rootAllocator.verify(); + + arrowBuf4.getReferenceManager().release(); + childAllocator4.close(); + rootAllocator.verify(); + } + } + + @Test + public void testAllocator_unclaimedReservation() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator1 = + rootAllocator.newChildAllocator("unclaimedReservation", 0, MAX_ALLOCATION)) { + try (final AllocationReservation reservation = childAllocator1.newReservation()) { + assertTrue(reservation.add(64)); + } + rootAllocator.verify(); + } + } + } + + @Test + public void testAllocator_claimedReservation() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + + try (final BufferAllocator childAllocator1 = rootAllocator.newChildAllocator( + "claimedReservation", 0, MAX_ALLOCATION)) { + + try (final AllocationReservation reservation = childAllocator1.newReservation()) { + assertTrue(reservation.add(32)); + assertTrue(reservation.add(32)); + + final ArrowBuf arrowBuf = reservation.allocateBuffer(); + assertEquals(64, arrowBuf.capacity()); + rootAllocator.verify(); + + arrowBuf.getReferenceManager().release(); + rootAllocator.verify(); + } + rootAllocator.verify(); + } + } + } + + @Test + public void testInitReservationAndLimit() throws Exception { + try (final RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + try (final BufferAllocator childAllocator = rootAllocator.newChildAllocator( + "child", 2048, 4096)) { + assertEquals(2048, childAllocator.getInitReservation()); + assertEquals(4096, childAllocator.getLimit()); + } + } + } + + @Test + public void multiple() throws Exception { + final String owner = "test"; + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + + final int op = 100000; + + BufferAllocator frag1 = allocator.newChildAllocator(owner, 1500000, Long.MAX_VALUE); + BufferAllocator frag2 = allocator.newChildAllocator(owner, 500000, Long.MAX_VALUE); + + allocator.verify(); + + BufferAllocator allocator11 = frag1.newChildAllocator(owner, op, Long.MAX_VALUE); + ArrowBuf b11 = allocator11.buffer(1000000); + + allocator.verify(); + + BufferAllocator allocator12 = frag1.newChildAllocator(owner, op, Long.MAX_VALUE); + ArrowBuf b12 = allocator12.buffer(500000); + + allocator.verify(); + + BufferAllocator allocator21 = frag1.newChildAllocator(owner, op, Long.MAX_VALUE); + + allocator.verify(); + + BufferAllocator allocator22 = frag2.newChildAllocator(owner, op, Long.MAX_VALUE); + ArrowBuf b22 = allocator22.buffer(2000000); + + allocator.verify(); + + BufferAllocator frag3 = allocator.newChildAllocator(owner, 1000000, Long.MAX_VALUE); + + allocator.verify(); + + BufferAllocator allocator31 = frag3.newChildAllocator(owner, op, Long.MAX_VALUE); + ArrowBuf b31a = allocator31.buffer(200000); + + allocator.verify(); + + // Previously running operator completes + b22.getReferenceManager().release(); + + allocator.verify(); + + allocator22.close(); + + b31a.getReferenceManager().release(); + allocator31.close(); + + b12.getReferenceManager().release(); + allocator12.close(); + + allocator21.close(); + + b11.getReferenceManager().release(); + allocator11.close(); + + frag1.close(); + frag2.close(); + frag3.close(); + + } + } + + // This test needs to run in non-debug mode. So disabling the assertion status through class loader for this. + // The test passes if run individually with -Dtest=TestBaseAllocator#testMemoryLeakWithReservation + // but fails generally since the assertion status cannot be changed once the class is initialized. + // So setting the test to @ignore + @Test(expected = IllegalStateException.class) + @Ignore + public void testMemoryLeakWithReservation() throws Exception { + // disabling assertion status + AssertionUtil.class.getClassLoader().setClassAssertionStatus(AssertionUtil.class.getName(), false); + try (RootAllocator rootAllocator = new RootAllocator(MAX_ALLOCATION)) { + ChildAllocator childAllocator1 = (ChildAllocator) rootAllocator.newChildAllocator( + "child1", 1024, MAX_ALLOCATION); + rootAllocator.verify(); + + ChildAllocator childAllocator2 = (ChildAllocator) childAllocator1.newChildAllocator( + "child2", 1024, MAX_ALLOCATION); + rootAllocator.verify(); + + ArrowBuf buff = childAllocator2.buffer(256); + + Exception exception = assertThrows(IllegalStateException.class, () -> { + childAllocator2.close(); + }); + String exMessage = exception.getMessage(); + assertTrue(exMessage.contains("Memory leaked: (256)")); + + exception = assertThrows(IllegalStateException.class, () -> { + childAllocator1.close(); + }); + exMessage = exception.getMessage(); + assertTrue(exMessage.contains("Memory leaked: (256)")); + } + } + + public void assertEquiv(ArrowBuf origBuf, ArrowBuf newBuf) { + assertEquals(origBuf.readerIndex(), newBuf.readerIndex()); + assertEquals(origBuf.writerIndex(), newBuf.writerIndex()); + } +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestEmptyArrowBuf.java b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestEmptyArrowBuf.java new file mode 100644 index 000000000..3fd7ce74a --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestEmptyArrowBuf.java @@ -0,0 +1,88 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.Test; + +import io.netty.buffer.PooledByteBufAllocatorL; + +public class TestEmptyArrowBuf { + + private static final int MAX_ALLOCATION = 8 * 1024; + private static RootAllocator allocator; + + @BeforeClass + public static void beforeClass() { + allocator = new RootAllocator(MAX_ALLOCATION); + } + + /** Ensure the allocator is closed. */ + @AfterClass + public static void afterClass() { + if (allocator != null) { + allocator.close(); + } + } + + @Test + public void testZeroBuf() { + // Exercise the historical log inside the empty ArrowBuf. This is initialized statically, and there is a circular + // dependency between ArrowBuf and BaseAllocator, so if the initialization happens in the wrong order, the + // historical log will be null even though BaseAllocator.DEBUG is true. + allocator.getEmpty().print(new StringBuilder(), 0, BaseAllocator.Verbosity.LOG_WITH_STACKTRACE); + } + + @Test + public void testEmptyArrowBuf() { + ArrowBuf buf = new ArrowBuf(ReferenceManager.NO_OP, null, + 1024, new PooledByteBufAllocatorL().empty.memoryAddress()); + + buf.getReferenceManager().retain(); + buf.getReferenceManager().retain(8); + assertEquals(1024, buf.capacity()); + assertEquals(1, buf.getReferenceManager().getRefCount()); + assertEquals(0, buf.getActualMemoryConsumed()); + + for (int i = 0; i < 10; i++) { + buf.setByte(i, i); + } + assertEquals(0, buf.getActualMemoryConsumed()); + assertEquals(0, buf.getReferenceManager().getSize()); + assertEquals(0, buf.getReferenceManager().getAccountedSize()); + assertEquals(false, buf.getReferenceManager().release()); + assertEquals(false, buf.getReferenceManager().release(2)); + assertEquals(0, buf.getReferenceManager().getAllocator().getLimit()); + assertEquals(buf, buf.getReferenceManager().transferOwnership(buf, allocator).getTransferredBuffer()); + assertEquals(0, buf.readerIndex()); + assertEquals(0, buf.writerIndex()); + assertEquals(1, buf.refCnt()); + + ArrowBuf derive = buf.getReferenceManager().deriveBuffer(buf, 0, 100); + assertEquals(derive, buf); + assertEquals(1, buf.refCnt()); + assertEquals(1, derive.refCnt()); + + buf.close(); + + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestEndianness.java b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestEndianness.java new file mode 100644 index 000000000..dcaeb2488 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestEndianness.java @@ -0,0 +1,51 @@ +/* + * 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.Assert.assertEquals; + +import java.nio.ByteOrder; + +import org.junit.Test; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.NettyArrowBuf; + +public class TestEndianness { + + @Test + public void testNativeEndian() { + final BufferAllocator a = new RootAllocator(10000); + final ByteBuf b = NettyArrowBuf.unwrapBuffer(a.buffer(4)); + b.setInt(0, 35); + if (ByteOrder.nativeOrder() == ByteOrder.LITTLE_ENDIAN) { + assertEquals(b.getByte(0), 35); + assertEquals(b.getByte(1), 0); + assertEquals(b.getByte(2), 0); + assertEquals(b.getByte(3), 0); + } else { + assertEquals(b.getByte(0), 0); + assertEquals(b.getByte(1), 0); + assertEquals(b.getByte(2), 0); + assertEquals(b.getByte(3), 35); + } + b.release(); + a.close(); + } + +} diff --git a/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java new file mode 100644 index 000000000..1b64cd733 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/java/org/apache/arrow/memory/TestNettyAllocationManager.java @@ -0,0 +1,108 @@ +/* + * 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.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; + +/** + * Test cases for {@link NettyAllocationManager}. + */ +public class TestNettyAllocationManager { + + static int CUSTOMIZED_ALLOCATION_CUTOFF_VALUE = 1024; + + private BaseAllocator createCustomizedAllocator() { + return new RootAllocator(BaseAllocator.configBuilder() + .allocationManagerFactory(new AllocationManager.Factory() { + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long size) { + return new NettyAllocationManager(accountingAllocator, size, CUSTOMIZED_ALLOCATION_CUTOFF_VALUE); + } + + @Override + public ArrowBuf empty() { + return null; + } + }).build()); + } + + private void readWriteArrowBuf(ArrowBuf buffer) { + // write buffer + for (long i = 0; i < buffer.capacity() / 8; i++) { + buffer.setLong(i * 8, i); + } + + // read buffer + for (long i = 0; i < buffer.capacity() / 8; i++) { + long val = buffer.getLong(i * 8); + assertEquals(i, val); + } + } + + /** + * Test the allocation strategy for small buffers.. + */ + @Test + public void testSmallBufferAllocation() { + final long bufSize = CUSTOMIZED_ALLOCATION_CUTOFF_VALUE - 512L; + try (BaseAllocator allocator = createCustomizedAllocator(); + ArrowBuf buffer = allocator.buffer(bufSize)) { + + assertTrue(buffer.getReferenceManager() instanceof BufferLedger); + BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); + + // make sure we are using netty allocation manager + AllocationManager allocMgr = bufferLedger.getAllocationManager(); + assertTrue(allocMgr instanceof NettyAllocationManager); + NettyAllocationManager nettyMgr = (NettyAllocationManager) allocMgr; + + // for the small buffer allocation strategy, the chunk is not null + assertNotNull(nettyMgr.getMemoryChunk()); + + readWriteArrowBuf(buffer); + } + } + + /** + * Test the allocation strategy for large buffers.. + */ + @Test + public void testLargeBufferAllocation() { + final long bufSize = CUSTOMIZED_ALLOCATION_CUTOFF_VALUE + 1024L; + try (BaseAllocator allocator = createCustomizedAllocator(); + ArrowBuf buffer = allocator.buffer(bufSize)) { + assertTrue(buffer.getReferenceManager() instanceof BufferLedger); + BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); + + // make sure we are using netty allocation manager + AllocationManager allocMgr = bufferLedger.getAllocationManager(); + assertTrue(allocMgr instanceof NettyAllocationManager); + NettyAllocationManager nettyMgr = (NettyAllocationManager) allocMgr; + + // for the large buffer allocation strategy, the chunk is null + assertNull(nettyMgr.getMemoryChunk()); + + readWriteArrowBuf(buffer); + } + } +} diff --git a/src/arrow/java/memory/memory-netty/src/test/resources/logback.xml b/src/arrow/java/memory/memory-netty/src/test/resources/logback.xml new file mode 100644 index 000000000..4c54d18a2 --- /dev/null +++ b/src/arrow/java/memory/memory-netty/src/test/resources/logback.xml @@ -0,0 +1,28 @@ +<?xml version="1.0" encoding="UTF-8" ?> +<!-- 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. --> + +<configuration> + <statusListener class="ch.qos.logback.core.status.NopStatusListener"/> + <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender"> + <!-- encoders are assigned the type + ch.qos.logback.classic.encoder.PatternLayoutEncoder by default --> + <encoder> + <pattern>%d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} - %msg%n</pattern> + </encoder> + </appender> + + <logger name="org.apache.arrow" additivity="false"> + <level value="info" /> + <appender-ref ref="STDOUT" /> + </logger> + +</configuration> diff --git a/src/arrow/java/memory/memory-unsafe/pom.xml b/src/arrow/java/memory/memory-unsafe/pom.xml new file mode 100644 index 000000000..6358784d5 --- /dev/null +++ b/src/arrow/java/memory/memory-unsafe/pom.xml @@ -0,0 +1,33 @@ +<?xml version="1.0" encoding="UTF-8"?> +<!-- 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. --> +<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> + <parent> + <artifactId>arrow-memory</artifactId> + <groupId>org.apache.arrow</groupId> + <version>6.0.1</version> + </parent> + <modelVersion>4.0.0</modelVersion> + + <artifactId>arrow-memory-unsafe</artifactId> + <name>Arrow Memory - Unsafe</name> + <description>Allocator and utils for allocating memory in Arrow based on sun.misc.Unsafe</description> + + + <dependencies> + <dependency> + <groupId>org.apache.arrow</groupId> + <artifactId>arrow-memory-core</artifactId> + <version>${project.version}</version> + </dependency> + </dependencies> + +</project> diff --git a/src/arrow/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java b/src/arrow/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java new file mode 100644 index 000000000..720c3d02d --- /dev/null +++ b/src/arrow/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/DefaultAllocationManagerFactory.java @@ -0,0 +1,37 @@ +/* + * 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; + +/** + * The default Allocation Manager Factory for a module. + * + */ +public class DefaultAllocationManagerFactory implements AllocationManager.Factory { + + public static final AllocationManager.Factory FACTORY = UnsafeAllocationManager.FACTORY; + + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long size) { + return FACTORY.create(accountingAllocator, size); + } + + @Override + public ArrowBuf empty() { + return UnsafeAllocationManager.FACTORY.empty(); + } +} diff --git a/src/arrow/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java b/src/arrow/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java new file mode 100644 index 000000000..b10aba359 --- /dev/null +++ b/src/arrow/java/memory/memory-unsafe/src/main/java/org/apache/arrow/memory/UnsafeAllocationManager.java @@ -0,0 +1,70 @@ +/* + * 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 org.apache.arrow.memory.util.MemoryUtil; + +/** + * Allocation manager based on unsafe API. + */ +public final class UnsafeAllocationManager extends AllocationManager { + + private static final ArrowBuf EMPTY = new ArrowBuf(ReferenceManager.NO_OP, + null, + 0, + MemoryUtil.UNSAFE.allocateMemory(0) + ); + + public static final AllocationManager.Factory FACTORY = new Factory() { + @Override + public AllocationManager create(BufferAllocator accountingAllocator, long size) { + return new UnsafeAllocationManager(accountingAllocator, size); + } + + @Override + public ArrowBuf empty() { + return EMPTY; + } + }; + + private final long allocatedSize; + + private final long allocatedAddress; + + UnsafeAllocationManager(BufferAllocator accountingAllocator, long requestedSize) { + super(accountingAllocator); + allocatedAddress = MemoryUtil.UNSAFE.allocateMemory(requestedSize); + allocatedSize = requestedSize; + } + + @Override + public long getSize() { + return allocatedSize; + } + + @Override + protected long memoryAddress() { + return allocatedAddress; + } + + @Override + protected void release0() { + MemoryUtil.UNSAFE.freeMemory(allocatedAddress); + } + +} diff --git a/src/arrow/java/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/TestAllocationManagerUnsafe.java b/src/arrow/java/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/TestAllocationManagerUnsafe.java new file mode 100644 index 000000000..33abe92e5 --- /dev/null +++ b/src/arrow/java/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/TestAllocationManagerUnsafe.java @@ -0,0 +1,41 @@ +/* + * 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.Assert.assertEquals; + +import org.junit.Test; + +/** + * Test cases for {@link AllocationManager}. + */ +public class TestAllocationManagerUnsafe { + + @Test + public void testAllocationManagerType() { + + // test unsafe allocation manager type + System.setProperty( + DefaultAllocationManagerOption.ALLOCATION_MANAGER_TYPE_PROPERTY_NAME, "Unsafe"); + DefaultAllocationManagerOption.AllocationManagerType mgrType = + DefaultAllocationManagerOption.getDefaultAllocationManagerType(); + + assertEquals(DefaultAllocationManagerOption.AllocationManagerType.Unsafe, mgrType); + + } +} diff --git a/src/arrow/java/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/TestUnsafeAllocationManager.java b/src/arrow/java/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/TestUnsafeAllocationManager.java new file mode 100644 index 000000000..c15882a37 --- /dev/null +++ b/src/arrow/java/memory/memory-unsafe/src/test/java/org/apache/arrow/memory/TestUnsafeAllocationManager.java @@ -0,0 +1,68 @@ +/* + * 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.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +import org.junit.Test; + +/** + * Test cases for {@link UnsafeAllocationManager}. + */ +public class TestUnsafeAllocationManager { + + private BaseAllocator createUnsafeAllocator() { + return new RootAllocator(BaseAllocator.configBuilder().allocationManagerFactory(UnsafeAllocationManager.FACTORY) + .build()); + } + + private void readWriteArrowBuf(ArrowBuf buffer) { + // write buffer + for (long i = 0; i < buffer.capacity() / 8; i++) { + buffer.setLong(i * 8, i); + } + + // read buffer + for (long i = 0; i < buffer.capacity() / 8; i++) { + long val = buffer.getLong(i * 8); + assertEquals(i, val); + } + } + + /** + * Test the memory allocation for {@link UnsafeAllocationManager}. + */ + @Test + public void testBufferAllocation() { + final long bufSize = 4096L; + try (BaseAllocator allocator = createUnsafeAllocator(); + ArrowBuf buffer = allocator.buffer(bufSize)) { + assertTrue(buffer.getReferenceManager() instanceof BufferLedger); + BufferLedger bufferLedger = (BufferLedger) buffer.getReferenceManager(); + + // make sure we are using unsafe allocation manager + AllocationManager allocMgr = bufferLedger.getAllocationManager(); + assertTrue(allocMgr instanceof UnsafeAllocationManager); + UnsafeAllocationManager unsafeMgr = (UnsafeAllocationManager) allocMgr; + + assertEquals(bufSize, unsafeMgr.getSize()); + readWriteArrowBuf(buffer); + } + } +} diff --git a/src/arrow/java/memory/pom.xml b/src/arrow/java/memory/pom.xml new file mode 100644 index 000000000..7f2e642f0 --- /dev/null +++ b/src/arrow/java/memory/pom.xml @@ -0,0 +1,29 @@ +<?xml version="1.0"?> +<!-- 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. --> +<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd"> + <modelVersion>4.0.0</modelVersion> + <parent> + <groupId>org.apache.arrow</groupId> + <artifactId>arrow-java-root</artifactId> + <version>6.0.1</version> + </parent> + <artifactId>arrow-memory</artifactId> + <name>Arrow Memory</name> + <packaging>pom</packaging> + + <modules> + <module>memory-core</module> + <module>memory-unsafe</module> + <module>memory-netty</module> + </modules> + +</project> diff --git a/src/arrow/java/memory/src/test/java/io/netty/buffer/TestExpandableByteBuf.java b/src/arrow/java/memory/src/test/java/io/netty/buffer/TestExpandableByteBuf.java new file mode 100644 index 000000000..b39cca8e8 --- /dev/null +++ b/src/arrow/java/memory/src/test/java/io/netty/buffer/TestExpandableByteBuf.java @@ -0,0 +1,117 @@ +/* + * 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 io.netty.buffer; + +import org.apache.arrow.memory.ArrowBuf; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.junit.Assert; +import org.junit.Test; + +public class TestExpandableByteBuf { + + @Test + public void testCapacity() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); + ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); + ByteBuf newByteBuf = expandableByteBuf.capacity(31); + int capacity = newByteBuf.capacity(); + Assert.assertEquals(32, capacity); + } + } + + @Test + public void testCapacity1() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); + ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); + ByteBuf newByteBuf = expandableByteBuf.capacity(32); + int capacity = newByteBuf.capacity(); + Assert.assertEquals(32, capacity); + } + } + + @Test + public void testSetAndGetIntValues() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); + ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); + int [] intVals = new int[] {Integer.MIN_VALUE, Short.MIN_VALUE - 1, Short.MIN_VALUE, 0 , + Short.MAX_VALUE , Short.MAX_VALUE + 1, Integer.MAX_VALUE}; + for (int intValue :intVals) { + expandableByteBuf.setInt(0, intValue); + Assert.assertEquals(expandableByteBuf.getInt(0), intValue); + Assert.assertEquals(expandableByteBuf.getIntLE(0), Integer.reverseBytes(intValue)); + } + } + } + + @Test + public void testSetAndGetLongValues() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); + ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); + long [] longVals = new long[] {Long.MIN_VALUE, 0 , Long.MAX_VALUE}; + for (long longValue :longVals) { + expandableByteBuf.setLong(0, longValue); + Assert.assertEquals(expandableByteBuf.getLong(0), longValue); + Assert.assertEquals(expandableByteBuf.getLongLE(0), Long.reverseBytes(longValue)); + } + } + } + + @Test + public void testSetAndGetShortValues() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); + ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); + short [] shortVals = new short[] {Short.MIN_VALUE, 0 , Short.MAX_VALUE}; + for (short shortValue :shortVals) { + expandableByteBuf.setShort(0, shortValue); + Assert.assertEquals(expandableByteBuf.getShort(0), shortValue); + Assert.assertEquals(expandableByteBuf.getShortLE(0), Short.reverseBytes(shortValue)); + } + } + } + + @Test + public void testSetAndGetByteValues() { + try (BufferAllocator allocator = new RootAllocator(128); + ArrowBuf buf = allocator.buffer(20); + ) { + NettyArrowBuf srcByteBuf = NettyArrowBuf.unwrapBuffer(buf); + ExpandableByteBuf expandableByteBuf = new ExpandableByteBuf(srcByteBuf, allocator); + byte [] byteVals = new byte[] {Byte.MIN_VALUE, 0 , Byte.MAX_VALUE}; + for (short byteValue :byteVals) { + expandableByteBuf.setByte(0, byteValue); + Assert.assertEquals(expandableByteBuf.getByte(0), byteValue); + } + } + } +} |