[SPARK-16021] Fill freed memory in test to help catch correctness bugs

## What changes were proposed in this pull request?

This patches `MemoryAllocator` to fill clean and freed memory with known byte values, similar to https://github.com/jemalloc/jemalloc/wiki/Use-Case:-Find-a-memory-corruption-bug . Memory filling is flag-enabled in test only by default.

## How was this patch tested?

Unit test that it's on in test.

cc sameeragarwal

Author: Eric Liang <ekl@databricks.com>

Closes #13983 from ericl/spark-16021.

6 files changed, 56 insertions, 3 deletions

diff --git a/common/unsafe/src/main/java/org/apache/spark/unsafe/Platform.java b/common/unsafe/src/main/java/org/apache/spark/unsafe/Platform.java
index 77c8c398be..a2ee45c37e 100644
--- a/common/unsafe/src/main/java/org/apache/spark/unsafe/Platform.java
+++ b/common/unsafe/src/main/java/org/apache/spark/unsafe/Platform.java
@@ -176,6 +176,10 @@ public final class Platform {
     throw new IllegalStateException("unreachable");
   }
 
+  public static void setMemory(Object object, long offset, long size, byte value) {
+    _UNSAFE.setMemory(object, offset, size, value);
+  }
+
   public static void setMemory(long address, byte value, long size) {
     _UNSAFE.setMemory(address, size, value);
   }
diff --git a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java
index 09847cec9c..3cd4264680 100644
--- a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java
+++ b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/HeapMemoryAllocator.java
@@ -24,6 +24,7 @@ import java.util.LinkedList;
 import java.util.Map;
 
 import org.apache.spark.unsafe.Platform;
+import org.apache.spark.unsafe.memory.MemoryAllocator;
 
 /**
  * A simple {@link MemoryAllocator} that can allocate up to 16GB using a JVM long primitive array.
@@ -64,12 +65,19 @@ public class HeapMemoryAllocator implements MemoryAllocator {
       }
     }
     long[] array = new long[(int) ((size + 7) / 8)];
-    return new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, size);
+    MemoryBlock memory = new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, size);
+    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
+      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
+    }
+    return memory;
   }
 
   @Override
   public void free(MemoryBlock memory) {
     final long size = memory.size();
+    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
+      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_FREED_VALUE);
+    }
     if (shouldPool(size)) {
       synchronized (this) {
         LinkedList<WeakReference<MemoryBlock>> pool = bufferPoolsBySize.get(size);
diff --git a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryAllocator.java b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryAllocator.java
index 5192f68c86..8bd2b06db8 100644
--- a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryAllocator.java
+++ b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryAllocator.java
@@ -20,8 +20,19 @@ package org.apache.spark.unsafe.memory;
 public interface MemoryAllocator {
 
   /**
+   * Whether to fill newly allocated and deallocated memory with 0xa5 and 0x5a bytes respectively.
+   * This helps catch misuse of uninitialized or freed memory, but imposes some overhead.
+   */
+  public static final boolean MEMORY_DEBUG_FILL_ENABLED = Boolean.parseBoolean(
+    System.getProperty("spark.memory.debugFill", "false"));
+
+  // Same as jemalloc's debug fill values.
+  public static final byte MEMORY_DEBUG_FILL_CLEAN_VALUE = (byte)0xa5;
+  public static final byte MEMORY_DEBUG_FILL_FREED_VALUE = (byte)0x5a;
+
+  /**
    * Allocates a contiguous block of memory. Note that the allocated memory is not guaranteed
-   * to be zeroed out (call `zero()` on the result if this is necessary).
+   * to be zeroed out (call `fill(0)` on the result if this is necessary).
    */
   MemoryBlock allocate(long size) throws OutOfMemoryError;
 
diff --git a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryBlock.java b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryBlock.java
index 1bc924d424..cd1d378bc1 100644
--- a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryBlock.java
+++ b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/MemoryBlock.java
@@ -53,4 +53,11 @@ public class MemoryBlock extends MemoryLocation {
   public static MemoryBlock fromLongArray(final long[] array) {
     return new MemoryBlock(array, Platform.LONG_ARRAY_OFFSET, array.length * 8L);
   }
+
+  /**
+   * Fills the memory block with the specified byte value.
+   */
+  public void fill(byte value) {
+    Platform.setMemory(obj, offset, length, value);
+  }
 }
diff --git a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java
index 98ce711176..55bcdf1ed7 100644
--- a/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java
+++ b/common/unsafe/src/main/java/org/apache/spark/unsafe/memory/UnsafeMemoryAllocator.java
@@ -27,13 +27,20 @@ public class UnsafeMemoryAllocator implements MemoryAllocator {
   @Override
   public MemoryBlock allocate(long size) throws OutOfMemoryError {
     long address = Platform.allocateMemory(size);
-    return new MemoryBlock(null, address, size);
+    MemoryBlock memory = new MemoryBlock(null, address, size);
+    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
+      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
+    }
+    return memory;
   }
 
   @Override
   public void free(MemoryBlock memory) {
     assert (memory.obj == null) :
       "baseObject not null; are you trying to use the off-heap allocator to free on-heap memory?";
+    if (MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED) {
+      memory.fill(MemoryAllocator.MEMORY_DEBUG_FILL_FREED_VALUE);
+    }
     Platform.freeMemory(memory.offset);
   }
 }
diff --git a/common/unsafe/src/test/java/org/apache/spark/unsafe/PlatformUtilSuite.java b/common/unsafe/src/test/java/org/apache/spark/unsafe/PlatformUtilSuite.java
index 693ec6ec58..a77ba826fc 100644
--- a/common/unsafe/src/test/java/org/apache/spark/unsafe/PlatformUtilSuite.java
+++ b/common/unsafe/src/test/java/org/apache/spark/unsafe/PlatformUtilSuite.java
@@ -17,6 +17,9 @@
 
 package org.apache.spark.unsafe;
 
+import org.apache.spark.unsafe.memory.MemoryAllocator;
+import org.apache.spark.unsafe.memory.MemoryBlock;
+
 import org.junit.Assert;
 import org.junit.Test;
 
@@ -58,4 +61,17 @@ public class PlatformUtilSuite {
       Assert.assertEquals((byte)i, data[i + 1]);
     }
   }
+
+  @Test
+  public void memoryDebugFillEnabledInTest() {
+    Assert.assertTrue(MemoryAllocator.MEMORY_DEBUG_FILL_ENABLED);
+    MemoryBlock onheap = MemoryAllocator.HEAP.allocate(1);
+    MemoryBlock offheap = MemoryAllocator.UNSAFE.allocate(1);
+    Assert.assertEquals(
+      Platform.getByte(onheap.getBaseObject(), onheap.getBaseOffset()),
+      MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
+    Assert.assertEquals(
+      Platform.getByte(offheap.getBaseObject(), offheap.getBaseOffset()),
+      MemoryAllocator.MEMORY_DEBUG_FILL_CLEAN_VALUE);
+  }
 }