ceph/src/arrow/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.memory;
  19
  20 import java.util.concurrent.atomic.AtomicLong;
  21
  22 import org.apache.arrow.util.Preconditions;
  23
  24 /**
  25  * The abstract base class of AllocationManager.
  26  *
  27  * <p>Manages the relationship between one or more allocators and a particular UDLE. Ensures that
  28  * one allocator owns the
  29  * memory that multiple allocators may be referencing. Manages a BufferLedger between each of its
  30  * associated allocators.
  31  *
  32  * <p>The only reason that this isn't package private is we're forced to put ArrowBuf in Netty's
  33  * package which need access
  34  * to these objects or methods.
  35  *
  36  * <p>Threading: AllocationManager manages thread-safety internally. Operations within the context
  37  * of a single BufferLedger
  38  * are lockless in nature and can be leveraged by multiple threads. Operations that cross the
  39  * context of two ledgers
  40  * will acquire a lock on the AllocationManager instance. Important note, there is one
  41  * AllocationManager per
  42  * UnsafeDirectLittleEndian buffer allocation. As such, there will be thousands of these in a
  43  * typical query. The
  44  * contention of acquiring a lock on AllocationManager should be very low.
  45  */
  46 public abstract class AllocationManager {
  47
  48   private static final AtomicLong MANAGER_ID_GENERATOR = new AtomicLong(0);
  49
  50   private final BufferAllocator root;
  51   private final long allocatorManagerId = MANAGER_ID_GENERATOR.incrementAndGet();
  52   // ARROW-1627 Trying to minimize memory overhead caused by previously used IdentityHashMap
  53   // see JIRA for details
  54   private final LowCostIdentityHashMap<BufferAllocator, BufferLedger> map = new LowCostIdentityHashMap<>();
  55   private final long amCreationTime = System.nanoTime();
  56
  57   // The ReferenceManager created at the time of creation of this AllocationManager
  58   // is treated as the owning reference manager for the underlying chunk of memory
  59   // managed by this allocation manager
  60   private volatile BufferLedger owningLedger;
  61   private volatile long amDestructionTime = 0;
  62
  63   protected AllocationManager(BufferAllocator accountingAllocator) {
  64     Preconditions.checkNotNull(accountingAllocator);
  65     accountingAllocator.assertOpen();
  66
  67     this.root = accountingAllocator.getRoot();
  68
  69     // we do a no retain association since our creator will want to retrieve the newly created
  70     // ledger and will create a reference count at that point
  71     this.owningLedger = associate(accountingAllocator, false);
  72   }
  73
  74   BufferLedger getOwningLedger() {
  75     return owningLedger;
  76   }
  77
  78   void setOwningLedger(final BufferLedger ledger) {
  79     this.owningLedger = ledger;
  80   }
  81
  82   /**
  83    * Associate the existing underlying buffer with a new allocator. This will increase the
  84    * reference count on the corresponding buffer ledger by 1
  85    *
  86    * @param allocator The target allocator to associate this buffer with.
  87    * @return The reference manager (new or existing) that associates the underlying
  88    *         buffer to this new ledger.
  89    */
  90   BufferLedger associate(final BufferAllocator allocator) {
  91     return associate(allocator, true);
  92   }
  93
  94   private BufferLedger associate(final BufferAllocator allocator, final boolean retain) {
  95     allocator.assertOpen();
  96     Preconditions.checkState(root == allocator.getRoot(),
  97           "A buffer can only be associated between two allocators that share the same root");
  98
  99     synchronized (this) {
 100       BufferLedger ledger = map.get(allocator);
 101       if (ledger != null) {
 102         if (retain) {
 103           // bump the ref count for the ledger
 104           ledger.increment();
 105         }
 106         return ledger;
 107       }
 108
 109       ledger = new BufferLedger(allocator, this);
 110
 111       if (retain) {
 112         // the new reference manager will have a ref count of 1
 113         ledger.increment();
 114       }
 115
 116       // store the mapping for <allocator, reference manager>
 117       BufferLedger oldLedger = map.put(ledger);
 118       Preconditions.checkState(oldLedger == null,
 119           "Detected inconsistent state: A reference manager already exists for this allocator");
 120
 121       if (allocator instanceof BaseAllocator) {
 122         // needed for debugging only: keep a pointer to reference manager inside allocator
 123         // to dump state, verify allocator state etc
 124         ((BaseAllocator) allocator).associateLedger(ledger);
 125       }
 126       return ledger;
 127     }
 128   }
 129
 130   /**
 131    * The way that a particular ReferenceManager (BufferLedger) communicates back to the
 132    * AllocationManager that it no longer needs to hold a reference to a particular
 133    * piece of memory. Reference manager needs to hold a lock to invoke this method
 134    * It is called when the shared refcount of all the ArrowBufs managed by the
 135    * calling ReferenceManager drops to 0.
 136    */
 137   void release(final BufferLedger ledger) {
 138     final BufferAllocator allocator = ledger.getAllocator();
 139     allocator.assertOpen();
 140
 141     // remove the <BaseAllocator, BufferLedger> mapping for the allocator
 142     // of calling BufferLedger
 143     Preconditions.checkState(map.containsKey(allocator),
 144         "Expecting a mapping for allocator and reference manager");
 145     final BufferLedger oldLedger = map.remove(allocator);
 146
 147     BufferAllocator oldAllocator = oldLedger.getAllocator();
 148     if (oldAllocator instanceof BaseAllocator) {
 149       // needed for debug only: tell the allocator that AllocationManager is removing a
 150       // reference manager associated with this particular allocator
 151       ((BaseAllocator) oldAllocator).dissociateLedger(oldLedger);
 152     }
 153
 154     if (oldLedger == owningLedger) {
 155       // the release call was made by the owning reference manager
 156       if (map.isEmpty()) {
 157         // the only <allocator, reference manager> mapping was for the owner
 158         // which now has been removed, it implies we can safely destroy the
 159         // underlying memory chunk as it is no longer being referenced
 160         oldAllocator.releaseBytes(getSize());
 161         // free the memory chunk associated with the allocation manager
 162         release0();
 163         oldAllocator.getListener().onRelease(getSize());
 164         amDestructionTime = System.nanoTime();
 165         owningLedger = null;
 166       } else {
 167         // since the refcount dropped to 0 for the owning reference manager and allocation
 168         // manager will no longer keep a mapping for it, we need to change the owning
 169         // reference manager to whatever the next available <allocator, reference manager>
 170         // mapping exists.
 171         BufferLedger newOwningLedger = map.getNextValue();
 172         // we'll forcefully transfer the ownership and not worry about whether we
 173         // exceeded the limit since this consumer can't do anything with this.
 174         oldLedger.transferBalance(newOwningLedger);
 175       }
 176     } else {
 177       // the release call was made by a non-owning reference manager, so after remove there have
 178       // to be 1 or more <allocator, reference manager> mappings
 179       Preconditions.checkState(map.size() > 0,
 180           "The final removal of reference manager should be connected to owning reference manager");
 181     }
 182   }
 183
 184   /**
 185    * Return the size of underlying chunk of memory managed by this Allocation Manager.
 186    *
 187    * <p>The underlying memory chunk managed can be different from the original requested size.
 188    *
 189    * @return size of underlying memory chunk
 190    */
 191   public abstract long getSize();
 192
 193   /**
 194    * Return the absolute memory address pointing to the fist byte of underlying memory chunk.
 195    */
 196   protected abstract long memoryAddress();
 197
 198   /**
 199    * Release the underlying memory chunk.
 200    */
 201   protected abstract void release0();
 202
 203   /**
 204    * A factory interface for creating {@link AllocationManager}.
 205    * One may extend this interface to use a user-defined AllocationManager implementation.
 206    */
 207   public interface Factory {
 208
 209     /**
 210      * Create an {@link AllocationManager}.
 211      *
 212      * @param accountingAllocator The allocator that are expected to be associated with newly created AllocationManager.
 213      *                            Currently it is always equivalent to "this"
 214      * @param size Size (in bytes) of memory managed by the AllocationManager
 215      * @return The created AllocationManager used by this allocator
 216      */
 217     AllocationManager create(BufferAllocator accountingAllocator, long size);
 218
 219     ArrowBuf empty();
 220   }
 221 }