ARROW-4086: [Java] Add apis to debug memory alloc failures

- On failures, capture state for each allocator in chain and add the
  details to the exception.
- add APIs to get parent/child allocators
- Do not mask the original exception when allocating from vectors

Author: Pindikura Ravindra <[email protected]>

Closes #4369 from pravindra/memdebug and squashes the following commits:

92a269b <Pindikura Ravindra> ARROW-4086 :  Add apis to debug memory alloc failures

Author: Pindikura Ravindra

java/memory/src/main/java/org/apache/arrow/memory/Accountant.java

@@ -36,6 +36,8 @@ class Accountant implements AutoCloseable {
    */
   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
@@ -57,7 +59,8 @@ class Accountant implements AutoCloseable {
    */
   private final AtomicLong locallyHeldMemory = new AtomicLong();
 
-  public Accountant(Accountant parent, long reservation, long maxAllocation) {
+  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 " +
@@ -68,6 +71,7 @@ public Accountant(Accountant parent, long reservation, long maxAllocation) {
         "reserve memory.");
 
     this.parent = parent;
+    this.name = name;
     this.reservation = reservation;
     this.allocationLimit.set(maxAllocation);
 
@@ -77,28 +81,43 @@ public Accountant(Accountant parent, long reservation, long maxAllocation) {
       if (!outcome.isOk()) {
         throw new OutOfMemoryException(String.format(
             "Failure trying to allocate initial reservation for Allocator. " +
-                "Attempted to allocate %d bytes and received an outcome of %s.", reservation,
-            outcome.name()));
+                "Attempted to allocate %d bytes.", reservation,
+            outcome.getStatus().name()), outcome.getDetails());
       }
     }
   }
 
   /**
    * Attempt to allocate the requested amount of memory. Either completely succeeds or completely
-   * fails. Constructs a a
-   * log of delta
-   *
-   * <p>If it fails, no changes are made to accounting.
+   * fails. If it fails, no changes are made to accounting.
    *
    * @param size The amount of memory to reserve in bytes.
-   * @return True if the allocation was successful, false if the allocation failed.
+   * @return the status and details of allocation at each allocator in the chain.
    */
   AllocationOutcome allocateBytes(long size) {
-    final AllocationOutcome outcome = allocate(size, true, false);
-    if (!outcome.isOk()) {
+    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 outcome;
+    return status;
+  }
+
+  private AllocationOutcome.Status allocateBytesInternal(long size) {
+    return allocateBytesInternal(size, null /*details*/);
   }
 
   private void updatePeak() {
@@ -126,7 +145,7 @@ private void updatePeak() {
    * @return Whether the allocation fit within limits.
    */
   boolean forceAllocate(long size) {
-    final AllocationOutcome outcome = allocate(size, true, true);
+    final AllocationOutcome.Status outcome = allocate(size, true, true, null);
     return outcome.isOk();
   }
 
@@ -152,21 +171,38 @@ boolean forceAllocate(long size) {
    * @param forceAllocation    Whether we should force the allocation.
    * @return The outcome of the allocation.
    */
-  private AllocationOutcome allocate(final long size, final boolean incomingUpdatePeak, final boolean forceAllocation) {
+  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);
 
-    AllocationOutcome parentOutcome = AllocationOutcome.SUCCESS;
+    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);
+      parentOutcome = parent.allocate(parentRequest, updatePeak, forceAllocation, details);
     }
 
-    final AllocationOutcome finalOutcome = beyondLimit ? AllocationOutcome.FAILED_LOCAL :
-        parentOutcome.isOk() ? AllocationOutcome.SUCCESS : AllocationOutcome.FAILED_PARENT;
+    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();
@@ -206,6 +242,15 @@ public void close() {
     }
   }
 
+  /**
+   * Return the name of the accountant.
+   *
+   * @return name of accountant
+   */
+  public String getName() {
+    return name;
+  }
+
   /**
    * Return the current limit of this Accountant.
    *

java/memory/src/main/java/org/apache/arrow/memory/AllocationOutcome.java

@@ -17,38 +17,81 @@
 
 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 enum AllocationOutcome {
+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);
+  }
 
   /**
-   * Allocation succeeded.
+   * Get the status of the allocation.
+   * @return status code.
    */
-  SUCCESS(true),
+  public Status getStatus() {
+    return status;
+  }
 
   /**
-   * Allocation succeeded but only because the allocator was forced to move beyond a limit.
+   * Get additional details of the allocation (like the status at each allocator in the hierarchy).
+   * @return details of allocation
    */
-  FORCED_SUCCESS(true),
+  public Optional<AllocationOutcomeDetails> getDetails() {
+    return Optional.ofNullable(details);
+  }
 
   /**
-   * Allocation failed because the local allocator's limits were exceeded.
+   * Returns true if the allocation was a success.
+   * @return true if allocation was successful, false otherwise.
    */
-  FAILED_LOCAL(false),
+  public boolean isOk() {
+    return status.isOk();
+  }
 
   /**
-   * Allocation failed because a parent allocator's limits were exceeded.
+   * Allocation status code.
    */
-  FAILED_PARENT(false);
+  public enum Status {
+    /**
+     * Allocation succeeded.
+     */
+    SUCCESS(true),
 
-  private final boolean ok;
+    /**
+     * Allocation succeeded but only because the allocator was forced to move beyond a limit.
+     */
+    FORCED_SUCCESS(true),
 
-  AllocationOutcome(boolean ok) {
-    this.ok = ok;
-  }
+    /**
+     * Allocation failed because the local allocator's limits were exceeded.
+     */
+    FAILED_LOCAL(false),
 
-  public boolean isOk() {
-    return ok;
+    /**
+     * 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;
+    }
   }
 }

java/memory/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 ? "success" : "fail"))
+          .append("\n")
+          .toString();
+    }
+  }
+
+}