Memory manager

input-stream/src/integrationTest/java/software/amazon/s3/analyticsaccelerator/access/ConcurrencyCorrectnessTest.java

@@ -38,6 +38,7 @@ void testSequentialReads(
       StreamReadPatternKind streamReadPattern,
       AALInputStreamConfigurationKind configuration)
       throws IOException, InterruptedException, ExecutionException {
+
     testAALReadConcurrency(
         s3ClientKind,
         s3Object,
@@ -47,6 +48,17 @@ void testSequentialReads(
         CONCURRENCY_ITERATIONS);
   }
 
+  @ParameterizedTest
+  @MethodSource("evictionTests")
+  void testReadsWithEviction(
+      S3ClientKind s3ClientKind,
+      StreamReadPatternKind streamReadPattern,
+      AALInputStreamConfigurationKind configuration)
+      throws IOException, InterruptedException, ExecutionException {
+
+    testAALReadConcurrencyWithEviction(s3ClientKind, streamReadPattern, configuration);
+  }
+
   @ParameterizedTest
   @MethodSource("skippingReads")
   void testSkippingReads(
@@ -134,4 +146,9 @@ static Stream<Arguments> etagTests() {
         parquetPatterns(),
         getS3SeekableInputStreamConfigurations());
   }
+
+  static Stream<Arguments> evictionTests() {
+    return argumentsForEvictionTest(
+        getS3ClientKinds(), parquetPatterns(), getS3SeekableInputStreamConfigurations());
+  }
 }

input-stream/src/integrationTest/java/software/amazon/s3/analyticsaccelerator/access/IntegrationTestBase.java

@@ -323,6 +323,68 @@ protected void testAALReadConcurrency(
     }
   }
 
+  protected void testAALReadConcurrencyWithEviction(
+      @NonNull S3ClientKind s3ClientKind,
+      @NonNull StreamReadPatternKind streamReadPatternKind,
+      @NonNull AALInputStreamConfigurationKind AALInputStreamConfigurationKind)
+      throws IOException, InterruptedException, ExecutionException {
+
+    final Map<S3Object, Crc32CChecksum> s3ObjectCrc32CChecksumMap = new HashMap<>();
+    // Read using the standard S3 async client. We do this once, to calculate the checksums
+    for (S3Object obj : S3Object.smallAndMediumObjects()) {
+      System.out.println("Getting checksum for " + obj);
+      StreamReadPattern streamReadPattern = streamReadPatternKind.getStreamReadPattern(obj);
+      Crc32CChecksum directChecksum = new Crc32CChecksum();
+      executeReadPatternDirectly(s3ClientKind, obj, streamReadPattern, Optional.of(directChecksum));
+      s3ObjectCrc32CChecksumMap.put(obj, directChecksum);
+    }
+
+    List<S3Object> s3Objects = new ArrayList<>();
+    s3Objects.addAll(S3Object.smallObjects());
+    s3Objects.addAll(S3Object.smallAndMediumObjects());
+    s3Objects.addAll(S3Object.mediumObjects());
+
+    // Create the s3DATClientStreamReader - that creates the shared state
+    try (S3AALClientStreamReader s3AALClientStreamReader =
+        this.createS3AALClientStreamReader(s3ClientKind, AALInputStreamConfigurationKind)) {
+      // Create the thread pool
+      ExecutorService executorService = Executors.newFixedThreadPool(10);
+
+      List<Future<?>> resultFutures = new ArrayList<>();
+
+      for (S3Object obj : s3Objects) {
+        Future<?> future =
+            executorService.submit(
+                () -> {
+                  try {
+                    // This will create a new stream every time, but all streams will share state
+                    Crc32CChecksum datChecksum = new Crc32CChecksum();
+                    executeReadPatternOnAAL(
+                        obj,
+                        s3AALClientStreamReader,
+                        streamReadPatternKind.getStreamReadPattern(obj),
+                        Optional.of(datChecksum));
+
+                    // Assert checksums
+                    assertChecksums(s3ObjectCrc32CChecksumMap.get(obj), datChecksum);
+                  } catch (Throwable t) {
+                    throw new RuntimeException(t);
+                  }
+                });
+        resultFutures.add(future);
+      }
+
+      // wait for each future to propagate errors
+      for (Future<?> future : resultFutures) {
+        // This should throw an exception, if a thread threw one, including assertions
+        future.get();
+      }
+      // Shutdown. Wait for termination indefinitely - we expect it to always complete
+      executorService.shutdown();
+      assertTrue(executorService.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS));
+    }
+  }
+
   /**
    * Stream read patterns to test on
    *
@@ -404,6 +466,21 @@ static Stream<Arguments> argumentsFor(
         }
       }
     }
+    return results.stream();
+  }
+
+  static Stream<Arguments> argumentsForEvictionTest(
+      List<S3ClientKind> clients,
+      List<StreamReadPatternKind> readPatterns,
+      List<AALInputStreamConfigurationKind> configurations) {
+    ArrayList<Arguments> results = new ArrayList<>();
+    for (S3ClientKind client : clients) {
+      for (StreamReadPatternKind readPattern : readPatterns) {
+        for (AALInputStreamConfigurationKind configuration : configurations) {
+          results.add(Arguments.of(client, readPattern, configuration));
+        }
+      }
+    }
 
     return results.stream();
   }

input-stream/src/main/java/software/amazon/s3/analyticsaccelerator/S3SeekableInputStreamFactory.java

@@ -25,6 +25,7 @@
 import software.amazon.s3.analyticsaccelerator.io.logical.impl.ParquetColumnPrefetchStore;
 import software.amazon.s3.analyticsaccelerator.io.logical.impl.ParquetLogicalIOImpl;
 import software.amazon.s3.analyticsaccelerator.io.physical.data.BlobStore;
+import software.amazon.s3.analyticsaccelerator.io.physical.data.MemoryManager;
 import software.amazon.s3.analyticsaccelerator.io.physical.data.MetadataStore;
 import software.amazon.s3.analyticsaccelerator.io.physical.impl.PhysicalIOImpl;
 import software.amazon.s3.analyticsaccelerator.request.ObjectClient;
@@ -53,6 +54,7 @@ public class S3SeekableInputStreamFactory implements AutoCloseable {
   private final BlobStore objectBlobStore;
   private final Telemetry telemetry;
   private final ObjectFormatSelector objectFormatSelector;
+  private final MemoryManager memoryManager;
 
   /**
    * Creates a new instance of {@link S3SeekableInputStreamFactory}. This factory should be used to
@@ -75,6 +77,7 @@ public S3SeekableInputStreamFactory(
     this.objectFormatSelector = new ObjectFormatSelector(configuration.getLogicalIOConfiguration());
     this.objectBlobStore =
         new BlobStore(objectClient, telemetry, configuration.getPhysicalIOConfiguration());
+    this.memoryManager = new MemoryManager(objectBlobStore);
   }
 
   /**
@@ -125,7 +128,12 @@ LogicalIO createLogicalIO(S3URI s3URI, StreamContext streamContext) throws IOExc
         return new ParquetLogicalIOImpl(
             s3URI,
             new PhysicalIOImpl(
-                s3URI, objectMetadataStore, objectBlobStore, telemetry, streamContext),
+                s3URI,
+                objectMetadataStore,
+                objectBlobStore,
+                telemetry,
+                streamContext,
+                memoryManager),
             telemetry,
             configuration.getLogicalIOConfiguration(),
             parquetColumnPrefetchStore);
@@ -134,7 +142,12 @@ LogicalIO createLogicalIO(S3URI s3URI, StreamContext streamContext) throws IOExc
         return new DefaultLogicalIOImpl(
             s3URI,
             new PhysicalIOImpl(
-                s3URI, objectMetadataStore, objectBlobStore, telemetry, streamContext),
+                s3URI,
+                objectMetadataStore,
+                objectBlobStore,
+                telemetry,
+                streamContext,
+                memoryManager),
             telemetry);
     }
   }

input-stream/src/main/java/software/amazon/s3/analyticsaccelerator/io/physical/PhysicalIOConfiguration.java

@@ -30,7 +30,7 @@
 @Builder
 @EqualsAndHashCode
 public class PhysicalIOConfiguration {
-  private static final int DEFAULT_CAPACITY_BLOB_STORE = 50;
+  private static final long DEFAULT_MAX_MEMORY_LIMIT_AAL = Long.MAX_VALUE;
   private static final int DEFAULT_CAPACITY_METADATA_STORE = 50;
   private static final boolean DEFAULT_USE_SINGLE_CACHE = true;
   private static final long DEFAULT_BLOCK_SIZE_BYTES = 8 * ONE_MB;
@@ -40,10 +40,13 @@ public class PhysicalIOConfiguration {
   private static final double DEFAULT_SEQUENTIAL_PREFETCH_BASE = 2.0;
   private static final double DEFAULT_SEQUENTIAL_PREFETCH_SPEED = 1.0;
 
-  /** Capacity, in blobs. {@link PhysicalIOConfiguration#DEFAULT_CAPACITY_BLOB_STORE} by default. */
-  @Builder.Default private int blobStoreCapacity = DEFAULT_CAPACITY_BLOB_STORE;
+  /**
+   * Max memory to be used by library. {@link PhysicalIOConfiguration#DEFAULT_MAX_MEMORY_LIMIT_AAL}
+   * by default.
+   */
+  @Builder.Default private long maxMemoryLimitAAL = DEFAULT_MAX_MEMORY_LIMIT_AAL;
 
-  private static final String BLOB_STORE_CAPACITY_KEY = "blobstore.capacity";
+  private static final String MAX_MEMORY_LIMIT_AAL_KEY = "aal.maxmemory";
 
   /**
    * Capacity, in blobs. {@link PhysicalIOConfiguration#DEFAULT_CAPACITY_METADATA_STORE} by default.
@@ -104,8 +107,8 @@ public class PhysicalIOConfiguration {
    */
   public static PhysicalIOConfiguration fromConfiguration(ConnectorConfiguration configuration) {
     return PhysicalIOConfiguration.builder()
-        .blobStoreCapacity(
-            configuration.getInt(BLOB_STORE_CAPACITY_KEY, DEFAULT_CAPACITY_BLOB_STORE))
+        .maxMemoryLimitAAL(
+            configuration.getLong(MAX_MEMORY_LIMIT_AAL_KEY, DEFAULT_MAX_MEMORY_LIMIT_AAL))
         .metadataStoreCapacity(
             configuration.getInt(METADATA_STORE_CAPACITY_KEY, DEFAULT_CAPACITY_METADATA_STORE))
         .blockSizeBytes(configuration.getLong(BLOCK_SIZE_BYTES_KEY, DEFAULT_BLOCK_SIZE_BYTES))
@@ -123,7 +126,7 @@ public static PhysicalIOConfiguration fromConfiguration(ConnectorConfiguration c
   /**
    * Constructs {@link PhysicalIOConfiguration}.
    *
-   * @param blobStoreCapacity The capacity of the BlobStore
+   * @param maxMemoryLimitAAL The memory limit for library
    * @param metadataStoreCapacity The capacity of the MetadataStore
    * @param blockSizeBytes Block size, in bytes
    * @param readAheadBytes Read ahead, in bytes
@@ -136,15 +139,15 @@ public static PhysicalIOConfiguration fromConfiguration(ConnectorConfiguration c
    */
   @Builder
   private PhysicalIOConfiguration(
-      int blobStoreCapacity,
+      long maxMemoryLimitAAL,
       int metadataStoreCapacity,
       long blockSizeBytes,
       long readAheadBytes,
       long maxRangeSizeBytes,
       long partSizeBytes,
       double sequentialPrefetchBase,
       double sequentialPrefetchSpeed) {
-    Preconditions.checkArgument(blobStoreCapacity > 0, "`blobStoreCapacity` must be positive");
+    Preconditions.checkArgument(maxMemoryLimitAAL > 0, "`blobStoreCapacity` must be positive");
     Preconditions.checkArgument(
         metadataStoreCapacity > 0, "`metadataStoreCapacity` must be positive");
     Preconditions.checkArgument(blockSizeBytes > 0, "`blockSizeBytes` must be positive");
@@ -156,7 +159,7 @@ private PhysicalIOConfiguration(
     Preconditions.checkArgument(
         sequentialPrefetchSpeed > 0, "`sequentialPrefetchSpeed` must be positive");
 
-    this.blobStoreCapacity = blobStoreCapacity;
+    this.maxMemoryLimitAAL = maxMemoryLimitAAL;
     this.metadataStoreCapacity = metadataStoreCapacity;
     this.blockSizeBytes = blockSizeBytes;
     this.readAheadBytes = readAheadBytes;
@@ -171,7 +174,7 @@ public String toString() {
     final StringBuilder builder = new StringBuilder();
 
     builder.append("PhysicalIO configuration:\n");
-    builder.append("\tblobStoreCapacity: " + blobStoreCapacity + "\n");
+    builder.append("\tmaxMemoryLimitAAL: " + maxMemoryLimitAAL + "\n");
     builder.append("\tmetadataStoreCapacity: " + metadataStoreCapacity + "\n");
     builder.append("\tblockSizeBytes: " + blockSizeBytes + "\n");
     builder.append("\treadAheadBytes: " + readAheadBytes + "\n");

input-stream/src/main/java/software/amazon/s3/analyticsaccelerator/io/physical/data/Blob.java

@@ -17,6 +17,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
 import lombok.NonNull;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -41,6 +42,7 @@ public class Blob implements Closeable {
   private final BlockManager blockManager;
   private final ObjectMetadata metadata;
   private final Telemetry telemetry;
+  private AtomicInteger activeReaders;
 
   /**
    * Construct a new Blob.
@@ -60,6 +62,32 @@ public Blob(
     this.metadata = metadata;
     this.blockManager = blockManager;
     this.telemetry = telemetry;
+    this.activeReaders = new AtomicInteger(0);
+  }
+
+  /**
+   * Updates and returns the active readers of this blob
+   *
+   * @param readers The delta to be added to the current active readers of this blob
+   * @return the updated active readers for this blob
+   */
+  public int updateActiveReaders(int readers) {
+    return activeReaders.addAndGet(readers);
+  }
+
+  /** @return the current active readers of this blob */
+  public int getActiveReaders() {
+    return activeReaders.get();
+  }
+
+  /** @return the object key of this blob */
+  public ObjectKey getObjectKey() {
+    return objectKey;
+  }
+
+  /** @return the blockManager of this blob */
+  public BlockManager getBlockManager() {
+    return blockManager;
   }
 
   /**
@@ -71,8 +99,13 @@ public Blob(
    */
   public int read(long pos) throws IOException {
     Preconditions.checkArgument(pos >= 0, "`pos` must be non-negative");
-    blockManager.makePositionAvailable(pos, ReadMode.SYNC);
-    return blockManager.getBlock(pos).get().read(pos);
+    try {
+      blockManager.makePositionAvailable(pos, ReadMode.SYNC);
+      int bytesRead = blockManager.getBlock(pos).get().read(pos);
+      return bytesRead;
+    } finally {
+      updateActiveReaders(-1);
+    }
   }
 
   /**
@@ -92,34 +125,37 @@ public int read(byte[] buf, int off, int len, long pos) throws IOException {
     Preconditions.checkArgument(0 <= len, "`len` must not be negative");
     Preconditions.checkArgument(off < buf.length, "`off` must be less than size of buffer");
 
-    blockManager.makeRangeAvailable(pos, len, ReadMode.SYNC);
+    try {
+      blockManager.makeRangeAvailable(pos, len, ReadMode.SYNC);
 
-    long nextPosition = pos;
-    int numBytesRead = 0;
+      long nextPosition = pos;
+      int numBytesRead = 0;
 
-    while (numBytesRead < len && nextPosition < contentLength()) {
-      final long nextPositionFinal = nextPosition;
-      Block nextBlock =
-          blockManager
-              .getBlock(nextPosition)
-              .orElseThrow(
-                  () ->
-                      new IllegalStateException(
-                          String.format(
-                              "This block (for position %s) should have been available.",
-                              nextPositionFinal)));
+      while (numBytesRead < len && nextPosition < contentLength()) {
+        final long nextPositionFinal = nextPosition;
+        Block nextBlock =
+            blockManager
+                .getBlock(nextPosition)
+                .orElseThrow(
+                    () ->
+                        new IllegalStateException(
+                            String.format(
+                                "This block (for position %s) should have been available.",
+                                nextPositionFinal)));
 
-      int bytesRead = nextBlock.read(buf, off + numBytesRead, len - numBytesRead, nextPosition);
+        int bytesRead = nextBlock.read(buf, off + numBytesRead, len - numBytesRead, nextPosition);
 
-      if (bytesRead == -1) {
-        return numBytesRead;
-      }
+        if (bytesRead == -1) {
+          return numBytesRead;
+        }
 
-      numBytesRead = numBytesRead + bytesRead;
-      nextPosition += bytesRead;
+        numBytesRead = numBytesRead + bytesRead;
+        nextPosition += bytesRead;
+      }
+      return numBytesRead;
+    } finally {
+      updateActiveReaders(-1);
     }
-
-    return numBytesRead;
   }
 
   /**