HIVE-29703: LLAP: Parquet footer larger than alloc.max throws RuntimeException during caching

llap-server/src/java/org/apache/hadoop/hive/llap/io/metadata/MetadataCache.java

@@ -257,41 +257,86 @@ public LlapBufferOrBuffers putFileMetadata(Object fileKey, int length, InputStre
   }
 
 
-  @SuppressWarnings({ "rawtypes", "unchecked" })
+  @SuppressWarnings("unchecked")
   private LlapBufferOrBuffers wrapBbForFile(LlapBufferOrBuffers result,
       Object fileKey, int length, InputStream stream, CacheTag tag, AtomicBoolean isStopped) throws IOException {
-    if (result != null) return result;
+    if (result != null) {
+      return result;
+    }
     int maxAlloc = allocator.getMaxAllocation();
-    LlapMetadataBuffer<Object>[] largeBuffers = null;
-    if (maxAlloc < length) {
-      largeBuffers = new LlapMetadataBuffer[length / maxAlloc];
-      for (int i = 0; i < largeBuffers.length; ++i) {
-        largeBuffers[i] = new LlapMetadataBuffer<>(fileKey, tag);
+    // Buffers are locked and handed to the cache only after this returns, so release whatever we
+    // allocated if a later read or allocation throws - otherwise it leaks (nothing else reclaims it).
+    if (length <= maxAlloc) {
+      // The whole footer fits in a single buffer - the overwhelmingly common case.
+      LlapMetadataBuffer<Object> buffer = new LlapMetadataBuffer<>(fileKey, tag);
+      allocator.allocateMultiple(new MemoryBuffer[] { buffer }, length, null, isStopped);
+      boolean done = false;
+      try {
+        readIntoCacheBuffer(stream, length, buffer);
+        done = true;
+        return buffer;
+      } finally {
+        if (!done) {
+          deallocateQuietly(buffer);
+        }
       }
-      allocator.allocateMultiple(largeBuffers, maxAlloc, null, isStopped);
+    }
+    // Larger footers are split across maxAlloc-sized chunks, the last one holding the remainder.
+    LlapMetadataBuffer<Object>[] largeBuffers = new LlapMetadataBuffer[length / maxAlloc];
+    for (int i = 0; i < largeBuffers.length; ++i) {
+      largeBuffers[i] = new LlapMetadataBuffer<>(fileKey, tag);
+    }
+    // allocateMultiple is all-or-nothing: on success every chunk is allocated; on failure it
+    // releases whatever it reserved.
+    allocator.allocateMultiple(largeBuffers, maxAlloc, null, isStopped);
+
+    LlapMetadataBuffer<Object> smallBuffer = null;
+    boolean done = false;
+    try {
       for (int i = 0; i < largeBuffers.length; ++i) {
         readIntoCacheBuffer(stream, maxAlloc, largeBuffers[i]);
       }
-    }
-    int smallSize = length % maxAlloc;
-    if (smallSize == 0) {
-      return new LlapMetadataBuffers(largeBuffers);
-    } else {
-      LlapMetadataBuffer<Object>[] smallBuffer = new LlapMetadataBuffer[1];
-      smallBuffer[0] = new LlapMetadataBuffer(fileKey, tag);
-      allocator.allocateMultiple(smallBuffer, length, null, isStopped);
-      readIntoCacheBuffer(stream, smallSize, smallBuffer[0]);
-      if (largeBuffers == null) {
-        return smallBuffer[0]; // This is the overwhelmingly common case.
-      } else {
-        LlapMetadataBuffer<Object>[] cacheData = new LlapMetadataBuffer[largeBuffers.length + 1];
-        System.arraycopy(largeBuffers, 0, cacheData, 0, largeBuffers.length);
-        cacheData[largeBuffers.length] = smallBuffer[0];
-        return new LlapMetadataBuffers<>(cacheData);
+      int smallSize = length % maxAlloc;
+      if (smallSize == 0) {
+        done = true;
+        return new LlapMetadataBuffers<>(largeBuffers);
+      }
+      // Allocate the remainder only; the last chunk is smaller than maxAlloc.
+      LlapMetadataBuffer<Object> remainder = new LlapMetadataBuffer<>(fileKey, tag);
+      allocator.allocateMultiple(new MemoryBuffer[] { remainder }, smallSize, null, isStopped);
+      smallBuffer = remainder; // Registered for cleanup only now that allocation succeeded.
+      readIntoCacheBuffer(stream, smallSize, remainder);
+
+      LlapMetadataBuffer<Object>[] cacheData = new LlapMetadataBuffer[largeBuffers.length + 1];
+      System.arraycopy(largeBuffers, 0, cacheData, 0, largeBuffers.length);
+      cacheData[largeBuffers.length] = remainder;
+
+      done = true;
+      return new LlapMetadataBuffers<>(cacheData);
+    } finally {
+      if (!done) {
+        for (LlapMetadataBuffer<Object> buffer : largeBuffers) {
+          deallocateQuietly(buffer);
+        }
+        if (smallBuffer != null) {
+          deallocateQuietly(smallBuffer);
+        }
       }
     }
   }
 
+  /**
+   * Releases a footer buffer that was allocated but not yet published to the cache. The buffer must
+   * be unlocked (refCount 0, as deallocate asserts); callers here run before it is ever locked.
+   */
+  private void deallocateQuietly(MemoryBuffer buffer) {
+    try {
+      allocator.deallocate(buffer);
+    } catch (Throwable t) {
+      LlapIoImpl.LOG.error("Ignoring the cleanup error after another error", t);
+    }
+  }
+
   private static void readIntoCacheBuffer(
       InputStream stream, int length, MemoryBuffer dest) throws IOException {
     ByteBuffer bb = dest.getByteBufferRaw();

llap-server/src/test/org/apache/hadoop/hive/llap/cache/TestOrcMetadataCache.java → llap-server/src/test/org/apache/hadoop/hive/llap/cache/TestMetadataCache.java

@@ -20,6 +20,8 @@
 import static org.apache.hadoop.hive.llap.cache.LlapCacheableBuffer.INVALIDATE_OK;
 import static org.junit.Assert.*;
 
+import java.io.ByteArrayInputStream;
+import java.io.EOFException;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Random;
@@ -34,6 +36,7 @@
 import org.apache.hadoop.hive.common.io.DataCache;
 import org.apache.hadoop.hive.common.io.DiskRange;
 import org.apache.hadoop.hive.common.io.DiskRangeList;
+import org.apache.hadoop.hive.common.io.encoded.MemoryBuffer;
 
 import org.apache.hadoop.hive.conf.HiveConf;
 import org.apache.hadoop.hive.llap.IllegalCacheConfigurationException;
@@ -51,7 +54,9 @@
 import org.junit.Assert;
 import org.junit.Test;
 
-public class TestOrcMetadataCache {
+public class TestMetadataCache {
+  private static final int MAX_ALLOC = 64;
+
   private static class DummyCachePolicy implements LowLevelCachePolicy {
     int lockCount = 0, unlockCount = 0;
 
@@ -94,10 +99,12 @@ public void debugDumpShort(StringBuilder sb) {
 
   private static class DummyMemoryManager implements MemoryManager {
     private int allocs;
+    private long reservedBytes;
 
     @Override
     public void reserveMemory(long memoryToReserve, AtomicBoolean isStopped) {
       ++allocs;
+      reservedBytes += memoryToReserve;
     }
 
     @Override public long evictMemory(long memoryToEvict) {
@@ -106,21 +113,21 @@ public void reserveMemory(long memoryToReserve, AtomicBoolean isStopped) {
 
     @Override
     public void releaseMemory(long memUsage) {
+      reservedBytes -= memUsage;
     }
 
     @Override
     public void updateMaxSize(long maxSize) {
     }
+
+    long reservedBytes() {
+      return reservedBytes;
+    }
   }
 
   @Test
   public void testCaseSomePartialBuffersAreEvicted() {
-    final DummyMemoryManager mm = new DummyMemoryManager();
-    final DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    final LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    final BuddyAllocator alloc = new BuddyAllocator(false, false, 8, MAX_ALLOC, 1, 4096, 0, null, mm, metrics, null, true);
-    final MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+    final MetadataCache cache = newMetadataCache();
     final Object fileKey1 = new Object();
     final Random rdm = new Random();
     final ByteBuffer smallBuffer = ByteBuffer.allocate(2 * MAX_ALLOC);
@@ -138,14 +145,8 @@ public void testCaseSomePartialBuffersAreEvicted() {
   }
 
   @Test
-  public void testBuffers() throws Exception {
-    DummyMemoryManager mm = new DummyMemoryManager();
-    DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    BuddyAllocator alloc = new BuddyAllocator(
-        false, false, 8, MAX_ALLOC, 1, 4096, 0, null, mm, metrics, null, true);
-    MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+  public void testBuffers() {
+    MetadataCache cache = newMetadataCache();
     Object fileKey1 = new Object();
     Random rdm = new Random();
 
@@ -185,6 +186,104 @@ public void testBuffers() throws Exception {
     assertFalse(b0.incRef() > 0); // Should have also been thrown out.
   }
 
+  /**
+   * A footer of exactly maxAlloc: the split condition used a strict {@code <}, so no buffer was
+   * allocated and the InputStream put path returned an empty wrapper (NPE on lock).
+   */
+  @Test
+  public void testStreamFooterEqualToMaxAlloc() throws Exception {
+    verifyFooterFromStream(newMetadataCache(), MAX_ALLOC, true);
+  }
+
+  /**
+   * A footer larger than maxAlloc is split into a multi-buffer wrapper: an exact multiple of maxAlloc
+   * (only full chunks) and a multiple plus a remainder (full chunks + a smaller last chunk). The
+   * remainder buffer used to be allocated with the full footer length, exceeding the max allocation.
+   */
+  @Test
+  public void testStreamFooterLargerThanMaxAlloc() throws Exception {
+    MetadataCache cache = newMetadataCache();
+    verifyFooterFromStream(cache, 2 * MAX_ALLOC, false);
+    verifyFooterFromStream(cache, 2 * MAX_ALLOC + 3, false);
+  }
+
+  /**
+   * If reading the footer stream fails after some buffers were already allocated, every allocated
+   * buffer must be released back to the allocator. Otherwise the long-lived LLAP daemon leaks arena
+   * memory on every failed footer read (a truncated file, an IO error, or an oversized chunk).
+   */
+  @Test
+  public void testStreamFooterReadFailureReleasesBuffers() throws IOException {
+    assertFooterReadFailureReleasesBuffers(MAX_ALLOC);         // single-buffer path
+    assertFooterReadFailureReleasesBuffers(2 * MAX_ALLOC);     // multi-buffer, exact multiple
+    assertFooterReadFailureReleasesBuffers(2 * MAX_ALLOC + 3); // multi-buffer, with remainder
+  }
+
+  private void assertFooterReadFailureReleasesBuffers(int length) throws IOException {
+    final int arenaSize = 4096;
+    DummyMemoryManager mm = new DummyMemoryManager();
+    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
+    BuddyAllocator alloc = new BuddyAllocator(
+        false, false, 8, MAX_ALLOC, 1, arenaSize, 0, null, mm, metrics, null, true);
+    MetadataCache cache = new MetadataCache(alloc, mm, new DummyCachePolicy(), true, metrics);
+
+    Object fileKey = new Object();
+    // One byte short of the footer, so reading the last chunk hits EOF part-way through.
+    ByteArrayInputStream truncated = new ByteArrayInputStream(new byte[length - 1]);
+    try {
+      cache.putFileMetadata(fileKey, length, truncated, null, null);
+      fail("Expected an EOFException for a truncated footer stream of length " + length);
+    } catch (EOFException expected) {
+      // expected
+    }
+    assertEquals("Allocated buffers must be released after a failed footer read (length " + length + ")",
+        0, mm.reservedBytes());
+    assertNull("A failed put must not leave a cache entry", cache.getFileMetadata(fileKey));
+
+    // Independent check against the real allocator: leaked buffers keep FLAG_NEW_ALLOC and can never
+    // be force-discarded, so re-filling the whole arena succeeds only if nothing leaked.
+    MemoryBuffer[] wholeArena = new MemoryBuffer[arenaSize / MAX_ALLOC];
+    alloc.allocateMultiple(wholeArena, MAX_ALLOC);
+  }
+
+  private MetadataCache newMetadataCache() {
+    return newMetadataCache(4096);
+  }
+
+  private MetadataCache newMetadataCache(int arenaSize) {
+    return newMetadataCache(arenaSize, new DummyCachePolicy());
+  }
+
+  private MetadataCache newMetadataCache(int arenaSize, DummyCachePolicy cp) {
+    DummyMemoryManager mm = new DummyMemoryManager();
+    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
+    BuddyAllocator alloc = new BuddyAllocator(
+        false, false, 8, MAX_ALLOC, 1, arenaSize, 0, null, mm, metrics, null, true);
+    return new MetadataCache(alloc, mm, cp, true, metrics);
+  }
+
+  private void verifyFooterFromStream(MetadataCache cache, int length, boolean expectSingleBuffer)
+      throws IOException {
+    Object fileKey = new Object();
+    ByteBuffer footer = ByteBuffer.allocate(length);
+    new Random().nextBytes(footer.array());
+
+    LlapBufferOrBuffers result = cache.putFileMetadata(
+        fileKey, length, new ByteArrayInputStream(footer.array()), null, null);
+    cache.decRefBuffer(result);
+    assertEquals(expectSingleBuffer, result.getSingleLlapBuffer() != null);
+    assertEquals(footer, reassemble(result));
+
+    result = cache.getFileMetadata(fileKey);
+    assertEquals(footer, reassemble(result));
+    cache.decRefBuffer(result);
+  }
+
+  private ByteBuffer reassemble(LlapBufferOrBuffers result) {
+    LlapAllocatorBuffer single = result.getSingleLlapBuffer();
+    return single != null ? single.getByteBufferDup() : extractResultBbs(result);
+  }
+
   public ByteBuffer extractResultBbs(LlapBufferOrBuffers result) {
     int totalLen = 0;
     for (LlapAllocatorBuffer buf : result.getMultipleLlapBuffers()) {
@@ -199,14 +298,9 @@ public ByteBuffer extractResultBbs(LlapBufferOrBuffers result) {
   }
 
   @Test
-  public void testIncompleteCbs() throws Exception {
-    DummyMemoryManager mm = new DummyMemoryManager();
+  public void testIncompleteCbs() {
     DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    BuddyAllocator alloc = new BuddyAllocator(
-        false, false, 8, MAX_ALLOC, 1, 4096, 0, null, mm, metrics, null, true);
-    MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+    MetadataCache cache = newMetadataCache(4096, cp);
     DataCache.BooleanRef gotAllData = new DataCache.BooleanRef();
     Object fileKey1 = new Object();
 
@@ -239,13 +333,7 @@ public void testIncompleteCbs() throws Exception {
 
   @Test
   public void testGetOrcTailForPath() throws Exception {
-    DummyMemoryManager mm = new DummyMemoryManager();
-    DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    BuddyAllocator alloc = new BuddyAllocator(
-        false, false, 8, MAX_ALLOC, 1, 4 * 4096, 0, null, mm, metrics, null, true);
-    MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+    MetadataCache cache = newMetadataCache(4 * 4096);
 
     Path path = new Path("../data/files/alltypesorc");
     Configuration jobConf = new Configuration();
@@ -260,13 +348,7 @@ public void testGetOrcTailForPath() throws Exception {
 
   @Test
   public void testGetOrcTailForPathWithFileId() throws Exception {
-    DummyMemoryManager mm = new DummyMemoryManager();
-    DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    BuddyAllocator alloc = new BuddyAllocator(
-        false, false, 8, MAX_ALLOC, 1, 4 * 4096, 0, null, mm, metrics, null, true);
-    MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+    MetadataCache cache = newMetadataCache(4 * 4096);
 
     Path path = new Path("../data/files/alltypesorc");
     Configuration jobConf = new Configuration();
@@ -284,13 +366,7 @@ public void testGetOrcTailForPathWithFileId() throws Exception {
 
   @Test
   public void testGetOrcTailForPathWithFileIdChange() throws Exception {
-    DummyMemoryManager mm = new DummyMemoryManager();
-    DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    BuddyAllocator alloc = new BuddyAllocator(
-        false, false, 8, MAX_ALLOC, 1, 4 * 4096, 0, null, mm, metrics, null, true);
-    MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+    MetadataCache cache = newMetadataCache(4 * 4096);
 
     Path path = new Path("../data/files/alltypesorc");
     Configuration jobConf = new Configuration();
@@ -317,14 +393,8 @@ public void testGetOrcTailForPathCacheNotReady() throws Exception {
   }
 
   @Test
-  public void testProactiveEvictionMark() throws Exception {
-    DummyMemoryManager mm = new DummyMemoryManager();
-    DummyCachePolicy cp = new DummyCachePolicy();
-    final int MAX_ALLOC = 64;
-    LlapDaemonCacheMetrics metrics = LlapDaemonCacheMetrics.create("", "");
-    BuddyAllocator alloc = new BuddyAllocator(
-        false, false, 8, MAX_ALLOC, 1, 4096, 0, null, mm, metrics, null, true);
-    MetadataCache cache = new MetadataCache(alloc, mm, cp, true, metrics);
+  public void testProactiveEvictionMark() {
+    MetadataCache cache = newMetadataCache();
 
     long fn1 = 1;
     long fn2 = 2;