fix(#760): remove future on exception

Change-Id: I9c876701b14b72d476f7a700465c05f6b870b5c1

Author: mapan

core/src/main/java/io/aiven/kafka/tieredstorage/fetch/cache/ChunkCache.java

@@ -47,6 +47,8 @@
 import com.github.benmanes.caffeine.cache.Weigher;
 
 public abstract class ChunkCache<T> implements ChunkManager, Configurable {
+    private static final org.slf4j.Logger log = org.slf4j.LoggerFactory.getLogger(ChunkCache.class);
+
     public static final String METRIC_GROUP = "chunk-cache-metrics";
     public static final String THREAD_POOL_METRIC_GROUP = "chunk-cache-thread-pool-metrics";
 
@@ -79,44 +81,34 @@ public InputStream getChunk(final ObjectKey objectKey,
         final var currentChunk = manifest.chunkIndex().chunks().get(chunkId);
         startPrefetching(objectKey, manifest, currentChunk.originalPosition + currentChunk.originalSize);
         final ChunkKey chunkKey = new ChunkKey(objectKey.value(), chunkId);
-        final AtomicReference<InputStream> result = new AtomicReference<>();
-        try {
-            return cache.asMap()
-                .compute(chunkKey, (key, val) -> {
-                    final CompletableFuture<T> future = CompletableFuture.supplyAsync(() -> {
-                        if (val == null) {
-                            statsCounter.recordMiss();
-                            try {
-                                final InputStream chunk =
-                                        chunkManager.getChunk(objectKey, manifest, chunkId);
-                                final T t = this.cacheChunk(chunkKey, chunk);
-                                result.getAndSet(cachedChunkToInputStream(t));
-                                return t;
-                            } catch (final StorageBackendException | IOException e) {
-                                throw new CompletionException(e);
-                            }
-                        } else {
-                            statsCounter.recordHit();
-                            try {
-                                final T cachedChunk = val.get();
-                                result.getAndSet(cachedChunkToInputStream(cachedChunk));
-                                return cachedChunk;
-                            } catch (final InterruptedException | ExecutionException e) {
-                                throw new CompletionException(e);
-                            }
-                        }
-                    }, executor);
 
-                    future.whenComplete((r, ex) -> {
-                        if (ex != null) {
-                            cache.asMap().remove(key, future);
-                        }
-                    });
+        final CompletableFuture<T> future = cache.asMap().compute(chunkKey, (key, val) -> {
+            if (val != null) {
+                statsCounter.recordHit();
+                return val;
+            }
 
-                    return future;
-                })
-                .thenApplyAsync(t -> result.get())
-                .get(getTimeout.toMillis(), TimeUnit.MILLISECONDS);
+            statsCounter.recordMiss();
+            final CompletableFuture<T> created = CompletableFuture.supplyAsync(() -> {
+                try (InputStream chunk = chunkManager.getChunk(objectKey, manifest, chunkId)) {
+                    return this.cacheChunk(chunkKey, chunk);
+                } catch (final StorageBackendException | IOException e) {
+                    throw new CompletionException(e);
+                }
+            });
+
+            created.whenComplete((r, ex) -> {
+                if (ex != null) {
+                    cache.asMap().remove(key, created);
+                }
+            });
+
+            return created;
+        });
+
+        try {
+            T t = future.get(getTimeout.toMillis(), TimeUnit.MILLISECONDS);
+            return this.cachedChunkToInputStream(t);
         } catch (final ExecutionException e) {
             // Unwrap previously wrapped exceptions if possible.
             final Throwable cause = e.getCause();