free-pool.h

Go to the documentation of this file.

// Copyright 2012 Cloudera Inc.
//
// 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.


#ifndef IMPALA_RUNTIME_FREE_POOL_H
#define IMPALA_RUNTIME_FREE_POOL_H

#include <stdio.h>
#include <string.h>
#include <string>
#include "common/logging.h"
#include "runtime/mem-pool.h"
#include "util/bit-util.h"

DECLARE_bool(disable_mem_pools);

namespace impala {

class FreePool {
 public:
  FreePool(MemPool* mem_pool)
    : mem_pool_(mem_pool),
      net_allocations_(0) {
    memset(&lists_, 0, sizeof(lists_));
  }

  uint8_t* Allocate(int size) {
    ++net_allocations_;
    if (FLAGS_disable_mem_pools) return reinterpret_cast<uint8_t*>(malloc(size));

    if (size == 0) return reinterpret_cast<uint8_t*>(0x1);

    int free_list_idx = BitUtil::Log2(size);
    DCHECK_LT(free_list_idx, NUM_LISTS);

    FreeListNode* allocation = lists_[free_list_idx].next;
    if (allocation == NULL) {
      // There wasn't an existing allocation of the right size, allocate a new one.
      size = 1 << free_list_idx;
      allocation = reinterpret_cast<FreeListNode*>(
          mem_pool_->Allocate(size + sizeof(FreeListNode)));
    } else {
      // Remove this allocation from the list.
      lists_[free_list_idx].next = allocation->next;
    }
    DCHECK(allocation != NULL);
    // Set the back node to point back to the list it came from so know where
    // to add it on Free().
    allocation->list = &lists_[free_list_idx];
    return reinterpret_cast<uint8_t*>(allocation) + sizeof(FreeListNode);
  }

  void Free(uint8_t* ptr) {
    --net_allocations_;
    if (FLAGS_disable_mem_pools) {
      free(ptr);
      return;
    }
    if (ptr == NULL || reinterpret_cast<int64_t>(ptr) == 0x1) return;
    FreeListNode* node = reinterpret_cast<FreeListNode*>(ptr - sizeof(FreeListNode));
    FreeListNode* list = node->list;
#ifndef NDEBUG
    CheckValidAllocation(list, ptr);
#endif
    // Add node to front of list.
    node->next = list->next;
    list->next = node;
  }

  uint8_t* Reallocate(uint8_t* ptr, int size) {
    if (FLAGS_disable_mem_pools) {
      return reinterpret_cast<uint8_t*>(realloc(reinterpret_cast<void*>(ptr), size));
    }
    if (ptr == NULL || reinterpret_cast<int64_t>(ptr) == 0x1) return Allocate(size);
    FreeListNode* node = reinterpret_cast<FreeListNode*>(ptr - sizeof(FreeListNode));
    FreeListNode* list = node->list;
#ifndef NDEBUG
    CheckValidAllocation(list, ptr);
#endif
    int bucket_idx = (list - &lists_[0]);
    // This is the actual size of ptr.
    int allocation_size = 1 << bucket_idx;

    // If it's already big enough, just return the ptr.
    if (allocation_size >= size) return ptr;

    // Make a new one. Since Allocate() already rounds up to powers of 2, this effectively
    // doubles for the caller.
    uint8_t* new_ptr = Allocate(size);
    memcpy(new_ptr, ptr, allocation_size);
    Free(ptr);
    return new_ptr;
  }

  MemTracker* mem_tracker() { return mem_pool_->mem_tracker(); }
  int64_t net_allocations() const { return net_allocations_; }

 private:
  static const int NUM_LISTS = 64;

  struct FreeListNode {
    union {
      FreeListNode* next; // Used when it is in the free list
      FreeListNode* list; // Used when it is being used by the caller.
    };
  };

  void CheckValidAllocation(FreeListNode* computed_list_ptr, uint8_t* allocation) const {
    // On debug, check that list is valid.
    bool found = false;
    for (int i = 0; i < NUM_LISTS && !found; ++i) {
      if (computed_list_ptr == &lists_[i]) found = true;
    }
    DCHECK(found) << "Could not find list for ptr: "
                  << reinterpret_cast<void*>(allocation)
                  << ". Allocation could have already been freed." << std::endl
                  << DebugString();
  }

  std::string DebugString() const {
    std::stringstream ss;
    ss << "FreePool: " << this << std::endl;
    for (int i = 0; i < NUM_LISTS; ++i) {
      FreeListNode* n = lists_[i].next;
      if (n == NULL) continue;
      ss << i << ": ";
      while (n != NULL) {
        uint8_t* ptr = reinterpret_cast<uint8_t*>(n);
        ptr += sizeof(FreeListNode);
        ss << reinterpret_cast<void*>(ptr);
        n = n->next;
        if (n != NULL) ss << "->";
      }
      ss << std::endl;
    }
    return ss.str();
  }

  MemPool* mem_pool_;

  FreeListNode lists_[NUM_LISTS];

  int64_t net_allocations_;
};

}

#endif