unsafe/memory/UnsafeMemoryAllocator.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.unsafe.memory;
0019
0020 import org.apache.spark.unsafe.Platform;
0021
0022 /**
0023  * A simple {@link MemoryAllocator} that uses {@code Unsafe} to allocate off-heap memory.
0024  */
0025 public class UnsafeMemoryAllocator implements MemoryAllocator {
0026
0027   @Override
0028   public MemoryBlock allocate(long size) throws OutOfMemoryError {
0029     long address = Platform.allocateMemory(size);
0030     MemoryBlock memory = new MemoryBlock(null, address, size);
0031     if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
0032       memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
0033     }
0034     return memory;
0035   }
0036
0037   @Override
0038   public void free(MemoryBlock memory) {
0039     assert (memory.obj == null) :
0040       "baseObject not null; are you trying to use the off-heap allocator to free on-heap memory?";
0041     assert (memory.pageNumber != MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER) :
0042       "page has already been freed";
0043     assert ((memory.pageNumber == MemoryBlock.NO_PAGE_NUMBER)
0044             || (memory.pageNumber == MemoryBlock.FREED_IN_TMM_PAGE_NUMBER)) :
0045       "TMM-allocated pages must be freed via TMM.freePage(), not directly in allocator free()";
0046
0047     if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
0048       memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_FREED_VALUE);
0049     }
0050     Platform.freeMemory(memory.offset);
0051     // As an additional layer of defense against use-after-free bugs, we mutate the
0052     // MemoryBlock to reset its pointer.
0053     memory.offset = 0;
0054     // Mark the page as freed (so we can detect double-frees).
0055     memory.pageNumber = MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER;
0056   }
0057 }