/*
 * 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.checkerframework.checker.nullness.qual.Nullable;

/**
 * An AllocationManager is the implementation of a physical memory allocation.
 *
 * <p>Manages the relationship between the allocators and a particular memory allocation. Ensures
 * that one allocator owns the memory that multiple allocators may be referencing. Manages a
 * BufferLedger between each of its associated allocators. It does not track the reference count;
 * that is the role of {@link BufferLedger} (aka {@link ReferenceManager}).
 *
 * <p>This is a public interface implemented by concrete allocator implementations (e.g. Netty or
 * Unsafe).
 *
 * <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 physical 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 {
  // The RootAllocator we are associated with. An allocation can only ever be associated with a
  // single RootAllocator.
  private final BufferAllocator root;
  // An allocation can be tracked by multiple allocators. (This is because an allocator is more like
  // a ledger.)
  // All such allocators track reference counts individually, via BufferLedger instances. When an
  // individual
  // reference count reaches zero, the allocator will be dissociated from this allocation. If that
  // was via the
  // owningLedger, then no more allocators should be tracking this allocation, and the allocation
  // will be freed.
  // ARROW-1627: Trying to minimize memory overhead caused by previously used IdentityHashMap
  private final LowCostIdentityHashMap<BufferAllocator, BufferLedger> map =
      new LowCostIdentityHashMap<>();
  // The primary BufferLedger (i.e. reference count) tracking this allocation.
  // This is mostly a semantic constraint on the API user: if the reference count reaches 0 in the
  // owningLedger, then
  // there are not supposed to be any references through other allocators. In practice, this doesn't
  // do anything
  // as the implementation just forces ownership to be transferred to one of the other extant
  // references.
  private volatile @Nullable BufferLedger owningLedger;

  @SuppressWarnings(
      "nullness:method.invocation") // call to associate(a, b) not allowed on the given receiver
  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);
  }

  @Nullable
  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) {
        // We were already being tracked by the given allocator, just return it
        if (retain) {
          // bump the ref count for the ledger
          ledger.increment();
        }
        return ledger;
      }

      // We weren't previously being tracked by the given allocator; create a new 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);
    Preconditions.checkState(
        oldLedger != null, "Expecting a mapping for allocator and reference manager");
    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());
        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();
  }
}
