Physical IO feature retry improvement

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

@@ -25,11 +25,12 @@
 import lombok.NonNull;
 import software.amazon.s3.analyticsaccelerator.common.Metrics;
 import software.amazon.s3.analyticsaccelerator.common.Preconditions;
-import software.amazon.s3.analyticsaccelerator.retry.RetryPolicy;
-import software.amazon.s3.analyticsaccelerator.retry.RetryStrategy;
-import software.amazon.s3.analyticsaccelerator.retry.SeekableInputStreamRetryStrategy;
 import software.amazon.s3.analyticsaccelerator.util.BlockKey;
 import software.amazon.s3.analyticsaccelerator.util.MetricKey;
+import software.amazon.s3.analyticsaccelerator.util.OpenStreamInformation;
+import software.amazon.s3.analyticsaccelerator.util.retry.DefaultRetryStrategyImpl;
+import software.amazon.s3.analyticsaccelerator.util.retry.RetryPolicy;
+import software.amazon.s3.analyticsaccelerator.util.retry.RetryStrategy;
 
 /**
  * Represents a block of data from an object stream, identified by a {@link BlockKey} and a
@@ -54,7 +55,8 @@ public class Block implements Closeable {
   private final Metrics aggregatingMetrics;
   private final long readTimeout;
   private final int retryCount;
-  private final RetryStrategy<Void> retryStrategy;
+  private final RetryStrategy retryStrategy;
+  private final OpenStreamInformation openStreamInformation;
 
   /**
    * A synchronization aid that allows threads to wait until the block's data is available.
@@ -83,14 +85,16 @@ public class Block implements Closeable {
    * @param aggregatingMetrics blobstore metrics
    * @param readTimeout read timeout in milliseconds
    * @param retryCount number of retries
+   * @param openStreamInformation contains stream information
    */
   public Block(
       @NonNull BlockKey blockKey,
       long generation,
       @NonNull BlobStoreIndexCache indexCache,
       @NonNull Metrics aggregatingMetrics,
       long readTimeout,
-      int retryCount) {
+      int retryCount,
+      @NonNull OpenStreamInformation openStreamInformation) {
     Preconditions.checkArgument(
         0 <= generation, "`generation` must be non-negative; was: %s", generation);
 
@@ -100,6 +104,7 @@ public Block(
     this.aggregatingMetrics = aggregatingMetrics;
     this.readTimeout = readTimeout;
     this.retryCount = retryCount;
+    this.openStreamInformation = openStreamInformation;
     this.retryStrategy = createRetryStrategy();
   }
 
@@ -110,16 +115,23 @@ public Block(
    * @throws RuntimeException if all retries fails and an error occurs
    */
   @SuppressWarnings("unchecked")
-  private RetryStrategy<Void> createRetryStrategy() {
+  private RetryStrategy createRetryStrategy() {
+    RetryStrategy base = new DefaultRetryStrategyImpl();
+    RetryStrategy provided = this.openStreamInformation.getRetryStrategy();
+
+    if (provided != null) {
+      base = base.merge(provided);
+    }
+
     if (this.readTimeout > 0) {
-      RetryPolicy<Void> timeoutRetries =
-          RetryPolicy.<Void>builder()
+      RetryPolicy timeoutPolicy =
+          RetryPolicy.builder()
               .handle(InterruptedException.class, TimeoutException.class, IOException.class)
               .withMaxRetries(this.retryCount)
               .build();
-      return new SeekableInputStreamRetryStrategy<>(timeoutRetries);
+      base = base.amend(timeoutPolicy);
     }
-    return new SeekableInputStreamRetryStrategy<>();
+    return base;
   }
 
   /**

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

@@ -58,6 +58,7 @@ public class BlockManager implements Closeable {
   private final BlockStore blockStore;
   private final SequentialReadProgression sequentialReadProgression;
   private final RangeOptimiser rangeOptimiser;
+  private final OpenStreamInformation openStreamInformation;
 
   private static final String OPERATION_MAKE_RANGE_AVAILABLE = "block.manager.make.range.available";
   private static final Logger LOG = LoggerFactory.getLogger(BlockManager.class);
@@ -92,6 +93,7 @@ public BlockManager(
     this.aggregatingMetrics = aggregatingMetrics;
     this.indexCache = indexCache;
     this.blockStore = new BlockStore(indexCache, aggregatingMetrics, configuration);
+    this.openStreamInformation = openStreamInformation;
     this.streamReader =
         new StreamReader(
             objectClient,
@@ -215,7 +217,8 @@ public synchronized void makeRangeAvailable(long pos, long len, ReadMode readMod
                       this.indexCache,
                       this.aggregatingMetrics,
                       this.configuration.getBlockReadTimeout(),
-                      this.configuration.getBlockReadRetryCount());
+                      this.configuration.getBlockReadRetryCount(),
+                      this.openStreamInformation);
               // Add block to the store for future reference
               blockStore.add(block);
               blocksToFill.add(block);

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

@@ -41,13 +41,13 @@
 import software.amazon.s3.analyticsaccelerator.request.Range;
 import software.amazon.s3.analyticsaccelerator.request.ReadMode;
 import software.amazon.s3.analyticsaccelerator.request.Referrer;
-import software.amazon.s3.analyticsaccelerator.retry.RetryPolicy;
-import software.amazon.s3.analyticsaccelerator.retry.RetryStrategy;
-import software.amazon.s3.analyticsaccelerator.retry.SeekableInputStreamRetryStrategy;
 import software.amazon.s3.analyticsaccelerator.util.MetricKey;
 import software.amazon.s3.analyticsaccelerator.util.ObjectKey;
 import software.amazon.s3.analyticsaccelerator.util.OpenStreamInformation;
 import software.amazon.s3.analyticsaccelerator.util.StreamAttributes;
+import software.amazon.s3.analyticsaccelerator.util.retry.DefaultRetryStrategyImpl;
+import software.amazon.s3.analyticsaccelerator.util.retry.RetryPolicy;
+import software.amazon.s3.analyticsaccelerator.util.retry.RetryStrategy;
 
 /**
  * {@code StreamReader} is responsible for asynchronously reading a range of bytes from an object in
@@ -67,7 +67,7 @@ public class StreamReader implements Closeable {
   private final Telemetry telemetry;
   private final PhysicalIOConfiguration physicalIOConfiguration;
 
-  private final RetryStrategy<ObjectContent> retryStrategy;
+  private final RetryStrategy retryStrategy;
 
   private static final String OPERATION_GET_OBJECT = "s3.stream.get";
   private static final String OPERATION_STREAM_READ = "s3.stream.read";
@@ -113,16 +113,23 @@ public StreamReader(
    * @throws RuntimeException if all retries fails and an error occurs
    */
   @SuppressWarnings("unchecked")
-  private RetryStrategy<ObjectContent> createRetryStrategy() {
+  private RetryStrategy createRetryStrategy() {
+    RetryStrategy base = new DefaultRetryStrategyImpl();
+    RetryStrategy provided = this.openStreamInformation.getRetryStrategy();
+
+    if (provided != null) {
+      base = base.merge(provided);
+    }
+
     if (this.physicalIOConfiguration.getBlockReadTimeout() > 0) {
-      RetryPolicy<ObjectContent> timeoutRetries =
-          RetryPolicy.<ObjectContent>builder()
+      RetryPolicy timeoutPolicy =
+          RetryPolicy.builder()
               .handle(InterruptedException.class, TimeoutException.class, ExecutionException.class)
               .withMaxRetries(this.physicalIOConfiguration.getBlockReadRetryCount())
               .build();
-      return new SeekableInputStreamRetryStrategy<>(timeoutRetries);
+      base = base.amend(timeoutPolicy);
     }
-    return new SeekableInputStreamRetryStrategy<>();
+    return base;
   }
 
   /**

input-stream/src/test/java/software/amazon/s3/analyticsaccelerator/io/physical/data/BlockStoreTest.java

@@ -93,7 +93,8 @@ public void test__blockStore__getBlockAfterAddBlock() {
             mock(BlobStoreIndexCache.class),
             mock(Metrics.class),
             DEFAULT_READ_TIMEOUT,
-            DEFAULT_RETRY_COUNT));
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT));
 
     // Then: getBlock can retrieve the same block
     Optional<Block> b = blockStore.getBlock(4);
@@ -149,7 +150,13 @@ public void test__blockStore__getBlockByIndex() {
         new BlockKey(objectKey, new Range(8192, 16383)); // Assuming readBufferSize is 8KB
     Block block =
         new Block(
-            blockKey, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
     blockStore.add(block);
 
     // When: getBlockByIndex is called with the correct index
@@ -187,10 +194,22 @@ public void test__blockStore__add_duplicateBlock() {
     BlockKey blockKey = new BlockKey(objectKey, new Range(0, 8191));
     Block block1 =
         new Block(
-            blockKey, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
     Block block2 =
         new Block(
-            blockKey, 1, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            1,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
 
     // When: The first block is added
     blockStore.add(block1);
@@ -219,7 +238,8 @@ public void test__blockStore__remove() throws IOException {
                 mockIndexCache,
                 mockMetrics,
                 DEFAULT_READ_TIMEOUT,
-                DEFAULT_RETRY_COUNT));
+                DEFAULT_RETRY_COUNT,
+                OpenStreamInformation.DEFAULT));
     block.setData(new byte[] {1, 2, 3, 4, 5});
     when(block.isDataReady()).thenReturn(true);
     blockStore.add(block);
@@ -245,7 +265,13 @@ public void test__blockStore__remove_nonExistentBlock() {
     BlockKey blockKey = new BlockKey(objectKey, new Range(0, 8191));
     Block block =
         new Block(
-            blockKey, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
 
     // When: An attempt is made to remove the block
     blockStore.remove(block);
@@ -266,7 +292,8 @@ public void test__blockStore__remove_dataNotReady() throws IOException {
                 mockIndexCache,
                 mockMetrics,
                 DEFAULT_READ_TIMEOUT,
-                DEFAULT_RETRY_COUNT));
+                DEFAULT_RETRY_COUNT,
+                OpenStreamInformation.DEFAULT));
     when(block.isDataReady()).thenReturn(false);
     blockStore.add(block);
 
@@ -292,10 +319,22 @@ public void test__blockStore__getMissingBlockIndexesInRange() {
 
     Block block1 =
         new Block(
-            blockKey1, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey1,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
     Block block2 =
         new Block(
-            blockKey2, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey2,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
 
     blockStore.add(block1);
     blockStore.add(block2);
@@ -355,7 +394,13 @@ public void test__blockStore__isEmpty() {
     BlockKey blockKey = new BlockKey(objectKey, new Range(0, 8191));
     Block block =
         new Block(
-            blockKey, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
     blockStore.add(block);
 
     // Then: isEmpty returns false
@@ -387,7 +432,8 @@ public void test__blockStore__concurrentAddRemove() throws InterruptedException
                     mockIndexCache,
                     mockMetrics,
                     DEFAULT_READ_TIMEOUT,
-                    DEFAULT_RETRY_COUNT);
+                    DEFAULT_RETRY_COUNT,
+                    OpenStreamInformation.DEFAULT);
             blockStore.add(block);
             blockStore.remove(block);
             latch.countDown();
@@ -406,7 +452,13 @@ public void test__blockStore__getBlock_atRangeBoundaries() {
     BlockKey blockKey = new BlockKey(objectKey, new Range(0, 8191));
     Block block =
         new Block(
-            blockKey, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
     blockStore.add(block);
 
     // At start of range
@@ -438,7 +490,13 @@ public void test__blockStore__getBlock_positionOutsideAnyBlock() {
     BlockKey blockKey = new BlockKey(objectKey, new Range(0, 8191));
     Block block =
         new Block(
-            blockKey, 0, mockIndexCache, mockMetrics, DEFAULT_READ_TIMEOUT, DEFAULT_RETRY_COUNT);
+            blockKey,
+            0,
+            mockIndexCache,
+            mockMetrics,
+            DEFAULT_READ_TIMEOUT,
+            DEFAULT_RETRY_COUNT,
+            OpenStreamInformation.DEFAULT);
     blockStore.add(block);
 
     Optional<Block> outsideBlock = blockStore.getBlock(100_000);