mem-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_MEM_POOL_H
#define IMPALA_RUNTIME_MEM_POOL_H

#include <stdio.h>
#include <algorithm>
#include <vector>
#include <string>

#include "common/logging.h"
#include "util/runtime-profile.h"

namespace impala {

class MemTracker;

//
//
//
//

class MemPool {
 public:
  MemPool(MemTracker* mem_tracker, int chunk_size = 0);

  ~MemPool();

  uint8_t* Allocate(int size) {
    return Allocate<false>(size);
  }

  uint8_t* TryAllocate(int size) {
    return Allocate<true>(size);
  }

  void ReturnPartialAllocation(int byte_size) {
    DCHECK_GE(byte_size, 0);
    DCHECK(current_chunk_idx_ != -1);
    ChunkInfo& info = chunks_[current_chunk_idx_];
    DCHECK(info.owns_data);
    DCHECK_GE(info.allocated_bytes, byte_size);
    info.allocated_bytes -= byte_size;
    total_allocated_bytes_ -= byte_size;
  }

  void Clear() {
    current_chunk_idx_ = -1;
    for (std::vector<ChunkInfo>::iterator chunk = chunks_.begin();
         chunk != chunks_.end(); ++chunk) {
      chunk->cumulative_allocated_bytes = 0;
      chunk->allocated_bytes = 0;
    }
    total_allocated_bytes_ = 0;
    DCHECK(CheckIntegrity(false));
  }

  void FreeAll();

  void AcquireData(MemPool* src, bool keep_current);

  bool Contains(uint8_t* ptr, int size);

  std::string DebugString();

  int64_t total_allocated_bytes() const { return total_allocated_bytes_; }
  int64_t peak_allocated_bytes() const { return peak_allocated_bytes_; }
  int64_t total_reserved_bytes() const { return total_reserved_bytes_; }
  MemTracker* mem_tracker() { return mem_tracker_; }

  int64_t GetTotalChunkSizes() const;

  int GetOffset(uint8_t* data);

  int GetCurrentOffset() const { return total_allocated_bytes_; }

  uint8_t* GetDataPtr(int offset);

  void GetChunkInfo(std::vector<std::pair<uint8_t*, int> >* chunk_info);

  std::string DebugPrint();

  static const char* LLVM_CLASS_NAME;

 private:
  friend class MemPoolTest;
  static const int DEFAULT_INITIAL_CHUNK_SIZE = 4 * 1024;

  struct ChunkInfo {
    bool owns_data;  // true if we eventually need to dealloc data
    uint8_t* data;
    int size;  // in bytes

    int cumulative_allocated_bytes;

    int allocated_bytes;

    explicit ChunkInfo(int size);

    ChunkInfo()
      : owns_data(true),
        data(NULL),
        size(0),
        cumulative_allocated_bytes(0),
        allocated_bytes(0) {}
  };

  int current_chunk_idx_;

  int last_offset_conversion_chunk_idx_;

  int chunk_size_;  // if != 0, use this size for new chunks

  int64_t total_allocated_bytes_;

  int64_t peak_allocated_bytes_;

  int64_t total_reserved_bytes_;

  std::vector<ChunkInfo> chunks_;

  MemTracker* mem_tracker_;

  bool FindChunk(int min_size, bool check_limits);

  bool CheckIntegrity(bool current_chunk_empty);

  int GetOffsetHelper(uint8_t* data);
  uint8_t* GetDataPtrHelper(int offset);

  int GetFreeOffset() const {
    if (current_chunk_idx_ == -1) return 0;
    return chunks_[current_chunk_idx_].allocated_bytes;
  }

  template <bool CHECK_LIMIT_FIRST>
  uint8_t* Allocate(int size) {
    if (size == 0) return NULL;

    int num_bytes = ((size + 7) / 8) * 8;  // round up to nearest 8 bytes
    if (current_chunk_idx_ == -1
        || num_bytes + chunks_[current_chunk_idx_].allocated_bytes
          > chunks_[current_chunk_idx_].size) {
      if (CHECK_LIMIT_FIRST) {
        // If we couldn't allocate a new chunk, return NULL.
        if (!FindChunk(num_bytes, true)) return NULL;
      } else {
        FindChunk(num_bytes, false);
      }
    }
    ChunkInfo& info = chunks_[current_chunk_idx_];
    DCHECK(info.owns_data);
    uint8_t* result = info.data + info.allocated_bytes;
    DCHECK_LE(info.allocated_bytes + num_bytes, info.size);
    info.allocated_bytes += num_bytes;
    total_allocated_bytes_ += num_bytes;
    DCHECK_LE(current_chunk_idx_, chunks_.size() - 1);
    peak_allocated_bytes_ = std::max(total_allocated_bytes_, peak_allocated_bytes_);
    return result;
  }
};

inline
int MemPool::GetOffset(uint8_t* data) {
  if (last_offset_conversion_chunk_idx_ != -1) {
    const ChunkInfo& info = chunks_[last_offset_conversion_chunk_idx_];
    if (info.data <= data && info.data + info.allocated_bytes > data) {
      return info.cumulative_allocated_bytes + data - info.data;
    }
  }
  return GetOffsetHelper(data);
}

inline
uint8_t* MemPool::GetDataPtr(int offset) {
  if (last_offset_conversion_chunk_idx_ != -1) {
    const ChunkInfo& info = chunks_[last_offset_conversion_chunk_idx_];
    if (info.cumulative_allocated_bytes <= offset
        && info.cumulative_allocated_bytes + info.allocated_bytes > offset) {
      return info.data + offset - info.cumulative_allocated_bytes;
    }
  }
  return GetDataPtrHelper(offset);
}

}

#endif