0
# Memory Management
1

2
Core memory allocation and management functionality providing both heap and off-heap memory strategies with automatic pooling, debug support, and memory block abstraction for Spark's internal operations.
3

4
## Capabilities
5

6
### MemoryAllocator Interface
7

8
The primary interface for memory allocation strategies, providing abstract methods for allocating and freeing memory blocks.
9

10
```java { .api }
11
/**
12
 * Interface for memory allocation strategies
13
 */
14
interface MemoryAllocator {
15
    /** 
16
     * Allocate a memory block of the specified size
17
     * @param size Size in bytes to allocate
18
     * @return MemoryBlock representing the allocated memory
19
     * @throws OutOfMemoryError if allocation fails
20
     */
21
    MemoryBlock allocate(long size) throws OutOfMemoryError;
22
    
23
    /**
24
     * Free a previously allocated memory block
25
     * @param memory MemoryBlock to free
26
     */
27
    void free(MemoryBlock memory);
28
    
29
    // Standard allocator instances
30
    static final MemoryAllocator HEAP = new HeapMemoryAllocator();
31
    static final MemoryAllocator UNSAFE = new UnsafeMemoryAllocator();
32
    
33
    // Debug configuration
34
    static final boolean MEMORY_DEBUG_FILL_ENABLED;
35
    static final byte MEMORY_DEBUG_FILL_CLEAN_VALUE = (byte)0xa5;
36
    static final byte MEMORY_DEBUG_FILL_FREED_VALUE = (byte)0x5a;
37
}
38
```
39

40
**Usage Examples:**
41

42
```java
43
import org.apache.spark.unsafe.memory.*;
44

45
// Using heap allocator
46
MemoryAllocator heapAllocator = MemoryAllocator.HEAP;
47
MemoryBlock heapBlock = heapAllocator.allocate(1024);
48
// ... use block ...
49
heapAllocator.free(heapBlock);
50

51
// Using off-heap allocator
52
MemoryAllocator unsafeAllocator = MemoryAllocator.UNSAFE;  
53
MemoryBlock offHeapBlock = unsafeAllocator.allocate(2048);
54
// ... use block ...
55
unsafeAllocator.free(offHeapBlock);
56
```
57

58
### HeapMemoryAllocator
59

60
JVM heap-based memory allocator with pooling for large allocations, designed to reduce GC pressure by reusing large byte arrays.
61

62
```java { .api }
63
/**
64
 * Memory allocator using JVM heap with pooling for large allocations
65
 */
66
class HeapMemoryAllocator implements MemoryAllocator {
67
    /**
68
     * Allocate heap memory, using pooling for large allocations
69
     * @param size Size in bytes to allocate
70
     * @return MemoryBlock backed by heap memory
71
     * @throws OutOfMemoryError if allocation fails
72
     */
73
    public MemoryBlock allocate(long size) throws OutOfMemoryError;
74
    
75
    /**
76
     * Free heap memory, returning large allocations to pool
77
     * @param memory MemoryBlock to free
78
     */
79
    public void free(MemoryBlock memory);
80
}
81
```
82

83
### UnsafeMemoryAllocator
84

85
Off-heap memory allocator using sun.misc.Unsafe for direct memory allocation, providing memory outside of JVM heap management.
86

87
```java { .api }
88
/**
89
 * Memory allocator using off-heap memory via Unsafe
90
 */
91
class UnsafeMemoryAllocator implements MemoryAllocator {
92
    /**
93
     * Allocate off-heap memory using Unsafe
94
     * @param size Size in bytes to allocate
95
     * @return MemoryBlock backed by off-heap memory
96
     * @throws OutOfMemoryError if allocation fails
97
     */
98
    public MemoryBlock allocate(long size) throws OutOfMemoryError;
99
    
100
    /**
101
     * Free off-heap memory using Unsafe
102
     * @param memory MemoryBlock to free
103
     */
104
    public void free(MemoryBlock memory);
105
}
106
```
107

108
### MemoryLocation
109

110
Base class for tracking memory locations using either object references (for heap) or direct addresses (for off-heap).
111

112
```java { .api }
113
/**
114
 * Memory location tracked by address or object reference
115
 */
116
class MemoryLocation {
117
    /**
118
     * Create memory location with object and offset
119
     * @param obj Base object (null for off-heap)
120
     * @param offset Offset within object or direct address
121
     */
122
    public MemoryLocation(Object obj, long offset);
123
    
124
    /**
125
     * Default constructor for uninitialized location
126
     */
127
    public MemoryLocation();
128
    
129
    /**
130
     * Update the memory location
131
     * @param newObj New base object
132
     * @param newOffset New offset
133
     */
134
    public void setObjAndOffset(Object newObj, long newOffset);
135
    
136
    /**
137
     * Get the base object (null for off-heap memory)
138
     * @return Base object or null
139
     */
140
    public final Object getBaseObject();
141
    
142
    /**
143
     * Get the offset within object or direct address
144
     * @return Offset or address
145
     */
146
    public final long getBaseOffset();
147
}
148
```
149

150
### MemoryBlock
151

152
Represents a consecutive block of memory with fixed size, extending MemoryLocation with size information and utility methods.
153

154
```java { .api }
155
/**
156
 * Consecutive block of memory with fixed size
157
 */
158
class MemoryBlock extends MemoryLocation {
159
    /**
160
     * Create memory block with object, offset, and length
161
     * @param obj Base object (null for off-heap)
162
     * @param offset Offset within object or direct address
163
     * @param length Size of the block in bytes
164
     */
165
    public MemoryBlock(Object obj, long offset, long length);
166
    
167
    /**
168
     * Get the size of this memory block
169
     * @return Size in bytes
170
     */
171
    public long size();
172
    
173
    /**
174
     * Fill the entire block with a specific byte value
175
     * @param value Byte value to fill with
176
     */
177
    public void fill(byte value);
178
    
179
    /**
180
     * Create a memory block backed by a long array
181
     * @param array Long array to wrap
182
     * @return MemoryBlock backed by the array
183
     */
184
    public static MemoryBlock fromLongArray(long[] array);
185
    
186
    // Page number for TaskMemoryManager integration
187
    public int pageNumber;
188
    
189
    // Special page number constants
190
    public static final int NO_PAGE_NUMBER = -1;
191
    public static final int FREED_IN_TMM_PAGE_NUMBER = -2;
192
    public static final int FREED_IN_ALLOCATOR_PAGE_NUMBER = -3;
193
}
194
```
195

196
**Usage Examples:**
197

198
```java
199
import org.apache.spark.unsafe.memory.*;
200

201
// Create and use memory blocks
202
MemoryAllocator allocator = MemoryAllocator.HEAP;
203
MemoryBlock block = allocator.allocate(1024);
204

205
// Fill block with zeros
206
block.fill((byte) 0);
207

208
// Get memory location info
209
Object baseObj = block.getBaseObject();
210
long baseOffset = block.getBaseOffset();
211
long size = block.size();
212

213
// Create block from long array
214
long[] data = {1L, 2L, 3L, 4L};
215
MemoryBlock arrayBlock = MemoryBlock.fromLongArray(data);
216

217
// Clean up
218
allocator.free(block);
219
// Note: arrayBlock doesn't need explicit freeing as it's backed by a Java array
220
```
221

222
## Memory Allocation Patterns
223

224
### Choosing Allocation Strategy
225

226
```java
227
// For temporary data that should be GC managed
228
MemoryAllocator heapAllocator = MemoryAllocator.HEAP;
229

230
// For large data that should not contribute to GC pressure  
231
MemoryAllocator offHeapAllocator = MemoryAllocator.UNSAFE;
232

233
// The choice depends on:
234
// - Data lifetime (short-lived: heap, long-lived: off-heap)
235
// - Size (small: heap, large: off-heap)
236
// - GC pressure considerations
237
```
238

239
### Safe Memory Management
240

241
```java
242
import org.apache.spark.unsafe.memory.*;
243

244
public void processData() {
245
    MemoryAllocator allocator = MemoryAllocator.HEAP;
246
    MemoryBlock block = null;
247
    
248
    try {
249
        block = allocator.allocate(1024);
250
        
251
        // Use the memory block
252
        block.fill((byte) 0);
253
        
254
        // ... processing logic ...
255
        
256
    } catch (OutOfMemoryError e) {
257
        // Handle allocation failure
258
        System.err.println("Failed to allocate memory: " + e.getMessage());
259
    } finally {
260
        // Always free allocated memory
261
        if (block != null) {
262
            allocator.free(block);
263
        }
264
    }
265
}
266
```
267

268
## Debug and Monitoring
269

270
When `MEMORY_DEBUG_FILL_ENABLED` is true, the allocators will:
271
- Fill newly allocated memory with `MEMORY_DEBUG_FILL_CLEAN_VALUE` (0xa5)
272
- Fill freed memory with `MEMORY_DEBUG_FILL_FREED_VALUE` (0x5a)
273

274
This helps detect use-after-free bugs and uninitialized memory access during development and testing.