Improve tracing of filesystem calls in CacheFileSystem

Use a tracing wrapper on external FS to track any FS calls made by cache implementations.
Removed redudant tracking of external FS calls from cache implementations.
Updated tests for tracing FS calls to be closer to reality.

Author: raunaqmorarka

lib/trino-filesystem-cache-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioInput.java

@@ -16,19 +16,13 @@
 import alluxio.client.file.URIStatus;
 import alluxio.client.file.cache.CacheManager;
 import alluxio.conf.AlluxioConfiguration;
-import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.Tracer;
 import io.trino.filesystem.TrinoInput;
 import io.trino.filesystem.TrinoInputFile;
 
 import java.io.EOFException;
 import java.io.IOException;
 
-import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_LOCATION;
-import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_READ_POSITION;
-import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_READ_SIZE;
-import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_KEY;
-import static io.trino.filesystem.tracing.Tracing.withTracing;
 import static java.lang.Math.min;
 import static java.util.Objects.checkFromIndexSize;
 import static java.util.Objects.requireNonNull;
@@ -39,9 +33,7 @@ public class AlluxioInput
     private final TrinoInputFile inputFile;
     private final long fileLength;
     private final AlluxioCacheStats statistics;
-    private final String cacheKey;
     private final AlluxioInputHelper helper;
-    private final Tracer tracer;
 
     private TrinoInput input;
     private boolean closed;
@@ -55,11 +47,9 @@ public AlluxioInput(
             AlluxioConfiguration configuration,
             AlluxioCacheStats statistics)
     {
-        this.tracer = requireNonNull(tracer, "tracer is null");
         this.inputFile = requireNonNull(inputFile, "inputFile is null");
         this.fileLength = requireNonNull(status, "status is null").getLength();
         this.statistics = requireNonNull(statistics, "statistics is null");
-        this.cacheKey = requireNonNull(cacheKey, "cacheKey is null");
         this.helper = new AlluxioInputHelper(tracer, inputFile.location(), cacheKey, status, cacheManager, configuration, statistics);
     }
 
@@ -90,22 +80,13 @@ private int doExternalRead(long position, byte[] buffer, int offset, int length)
             return 0;
         }
 
-        Span span = tracer.spanBuilder("Alluxio.readExternal")
-                .setAttribute(CACHE_KEY, cacheKey)
-                .setAttribute(CACHE_FILE_LOCATION, inputFile.location().toString())
-                .setAttribute(CACHE_FILE_READ_SIZE, (long) length)
-                .setAttribute(CACHE_FILE_READ_POSITION, position)
-                .startSpan();
-
-        return withTracing(span, () -> {
-            AlluxioInputHelper.PageAlignedRead aligned = helper.alignRead(position, length);
-            byte[] readBuffer = new byte[aligned.length()];
-            getInput().readFully(aligned.pageStart(), readBuffer, 0, readBuffer.length);
-            helper.putCache(aligned.pageStart(), aligned.pageEnd(), readBuffer, aligned.length());
-            System.arraycopy(readBuffer, aligned.pageOffset(), buffer, offset, length);
-            statistics.recordExternalRead(readBuffer.length);
-            return length;
-        });
+        AlluxioInputHelper.PageAlignedRead aligned = helper.alignRead(position, length);
+        byte[] readBuffer = new byte[aligned.length()];
+        getInput().readFully(aligned.pageStart(), readBuffer, 0, readBuffer.length);
+        helper.putCache(aligned.pageStart(), aligned.pageEnd(), readBuffer, aligned.length());
+        System.arraycopy(readBuffer, aligned.pageOffset(), buffer, offset, length);
+        statistics.recordExternalRead(readBuffer.length);
+        return length;
     }
 
     private TrinoInput getInput()

lib/trino-filesystem-cache-alluxio/src/main/java/io/trino/filesystem/alluxio/AlluxioInputStream.java

@@ -139,7 +139,7 @@ private int doExternalRead0(long readPosition, byte[] buffer, int offset, int le
             return 0;
         }
 
-        Span span = tracer.spanBuilder("Alluxio.readExternal")
+        Span span = tracer.spanBuilder("Alluxio.readExternalStream")
                 .setAttribute(CACHE_KEY, key)
                 .setAttribute(CACHE_FILE_LOCATION, inputFile.location().toString())
                 .setAttribute(CACHE_FILE_READ_SIZE, (long) length)

lib/trino-filesystem-cache-alluxio/src/test/java/io/trino/filesystem/alluxio/TestAlluxioCacheFileSystemAccessOperations.java

@@ -59,7 +59,11 @@
 import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_WRITE_POSITION;
 import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_FILE_WRITE_SIZE;
 import static io.trino.filesystem.tracing.CacheSystemAttributes.CACHE_KEY;
+import static io.trino.filesystem.tracing.FileSystemAttributes.FILE_LOCATION;
+import static io.trino.filesystem.tracing.FileSystemAttributes.FILE_READ_POSITION;
+import static io.trino.filesystem.tracing.FileSystemAttributes.FILE_READ_SIZE;
 import static io.trino.testing.MultisetAssertions.assertMultisetsEqual;
+import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toCollection;
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -128,7 +132,7 @@ public void testCache()
                         .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 11), readTimes)
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 11))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 11))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 11))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 11))
                         .build());
 
         byte[] modifiedContent = "modified content".getBytes(StandardCharsets.UTF_8);
@@ -140,7 +144,7 @@ public void testCache()
         readTimes = 7;
         assertCacheOperations(location, modifiedContent, readTimes,
                 ImmutableMultiset.<CacheOperationSpan>builder()
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 16))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 16))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 16))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 16))
                         .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 16), readTimes)
@@ -165,7 +169,7 @@ public void testPartialCacheHits()
                 ImmutableMultiset.<CacheOperationSpan>builder()
                         .add(new CacheOperationSpan("Alluxio.readCached", "memory:///partial", 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 0, PAGE_SIZE))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE))
                         .build());
@@ -175,7 +179,7 @@ public void testPartialCacheHits()
                         .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 0, PAGE_SIZE + 10))
                         .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), PAGE_SIZE, 10))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), PAGE_SIZE, PAGE_SIZE))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), PAGE_SIZE, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE))
                         .build());
@@ -214,7 +218,7 @@ public void testMultiPageExternalsReads()
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE))
                         .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), PAGE_SIZE + 1))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), PAGE_SIZE + 1))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), PAGE_SIZE * 2))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), PAGE_SIZE * 2))
                         .build());
         cacheKeyProvider.increaseCacheVersion();
@@ -223,7 +227,7 @@ public void testMultiPageExternalsReads()
                         .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, PAGE_SIZE))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), PAGE_SIZE, PAGE_SIZE))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 2 * PAGE_SIZE))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 2 * PAGE_SIZE))
                         .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 2 * PAGE_SIZE))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 2 * PAGE_SIZE))
                         .build());
@@ -276,7 +280,7 @@ public void testCacheWithMissingPage()
                         .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 12), readTimes)
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 12))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 12))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 12))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 12))
                         .build());
 
         TrinoInputFile inputFile = fileSystem.newInputFile(location);
@@ -289,7 +293,7 @@ public void testCacheWithMissingPage()
                 ImmutableMultiset.<CacheOperationSpan>builder()
                         .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 12))
                         .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 12))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 12))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 12))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 12))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 12))
                         .build());
@@ -313,7 +317,7 @@ public void testCacheWithCorruptedPage()
                         .addCopies(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 14), readTimes)
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 14))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 14))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 14))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 14))
                         .build());
 
         TrinoInputFile inputFile = fileSystem.newInputFile(location);
@@ -326,7 +330,7 @@ public void testCacheWithCorruptedPage()
                 ImmutableMultiset.<CacheOperationSpan>builder()
                         .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), 14))
                         .add(new CacheOperationSpan("AlluxioCacheManager.get", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 14))
-                        .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), 14))
+                        .add(new CacheOperationSpan("Input.readFully", location.toString(), 14))
                         .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), 14))
                         .add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), 0, 14))
                         .build());
@@ -366,7 +370,7 @@ private void assertUnCachedRead(Location location, int fileSize)
         ImmutableMultiset.Builder<CacheOperationSpan> builder = ImmutableMultiset.<CacheOperationSpan>builder()
                 .add(new CacheOperationSpan("Alluxio.readCached", location.toString(), fileSize))
                 .add(new CacheOperationSpan("Alluxio.writeCache", location.toString(), fileSize))
-                .add(new CacheOperationSpan("Alluxio.readExternal", location.toString(), fileSize));
+                .add(new CacheOperationSpan("Input.readFully", location.toString(), fileSize));
 
         for (int offset = 0; offset < fileSize; offset = offset + PAGE_SIZE) {
             builder.add(new CacheOperationSpan("AlluxioCacheManager.put", cacheKey(location, cacheKeyProvider.currentCacheVersion()), offset, PAGE_SIZE));
@@ -412,8 +416,9 @@ private void assertCacheOperations(Location location, byte[] content, int readTi
 
     private Multiset<CacheOperationSpan> getCacheOperations(List<SpanData> spans)
     {
-        return spans.stream().filter(span -> span.getName().startsWith("Alluxio"))
-                        .map(CacheOperationSpan::create)
+        return spans.stream()
+                .filter(span -> span.getName().startsWith("Input.") || span.getName().startsWith("Alluxio"))
+                .map(CacheOperationSpan::create)
                 .collect(toCollection(HashMultiset::create));
     }
 
@@ -429,18 +434,20 @@ public static CacheOperationSpan create(SpanData span)
             Attributes attributes = span.getAttributes();
 
             long length = switch (span.getName()) {
-                case "Alluxio.readCached", "Alluxio.readExternal", "AlluxioCacheManager.get" -> attributes.get(CACHE_FILE_READ_SIZE);
+                case "Alluxio.readCached", "Alluxio.readExternalStream", "AlluxioCacheManager.get" -> attributes.get(CACHE_FILE_READ_SIZE);
                 case "Alluxio.writeCache", "AlluxioCacheManager.put" -> attributes.get(CACHE_FILE_WRITE_SIZE);
+                case "Input.readFully" -> attributes.get(FILE_READ_SIZE);
                 default -> throw new IllegalArgumentException("Unrecognized span " + span.getName() + " [" + span.getAttributes() + "]");
             };
 
             long position = switch (span.getName()) {
-                case "Alluxio.readCached", "Alluxio.readExternal", "AlluxioCacheManager.get" -> attributes.get(CACHE_FILE_READ_POSITION);
+                case "Alluxio.readCached", "Alluxio.readExternalStream", "AlluxioCacheManager.get" -> attributes.get(CACHE_FILE_READ_POSITION);
                 case "Alluxio.writeCache", "AlluxioCacheManager.put" -> attributes.get(CACHE_FILE_WRITE_POSITION);
+                case "Input.readFully" -> attributes.get(FILE_READ_POSITION);
                 default -> throw new IllegalArgumentException("Unrecognized span  " + span.getName() + " [" + span.getAttributes() + "]");
             };
 
-            return new CacheOperationSpan(span.getName(), firstNonNull(attributes.get(CACHE_FILE_LOCATION), attributes.get(CACHE_KEY)), position, length);
+            return new CacheOperationSpan(span.getName(), getLocation(span), position, length);
         }
 
         @Override
@@ -454,4 +461,12 @@ private static String cacheKey(Location location, int cacheVersion)
     {
         return testingCacheKeyForLocation(location, cacheVersion);
     }
+
+    private static String getLocation(SpanData span)
+    {
+        if (span.getName().startsWith("Input.")) {
+            return requireNonNull(span.getAttributes().get(FILE_LOCATION));
+        }
+        return firstNonNull(span.getAttributes().get(CACHE_FILE_LOCATION), span.getAttributes().get(CACHE_KEY));
+    }
 }

lib/trino-filesystem-manager/src/main/java/io/trino/filesystem/manager/FileSystemModule.java

@@ -140,13 +140,14 @@ static TrinoFileSystemFactory createFileSystemFactory(
                 .orElseThrow(() -> new IllegalArgumentException("No factory for location: " + location));
 
         TrinoFileSystemFactory delegate = new SwitchingFileSystemFactory(loader);
+        delegate = new TracingFileSystemFactory(tracer, delegate);
         if (fileSystemCache.isPresent()) {
-            delegate = new CacheFileSystemFactory(tracer, delegate, fileSystemCache.orElseThrow(), keyProvider.orElseThrow());
+            return new CacheFileSystemFactory(tracer, delegate, fileSystemCache.orElseThrow(), keyProvider.orElseThrow());
         }
         // use MemoryFileSystemCache only when no other TrinoFileSystemCache is configured
-        else if (memoryFileSystemCache.isPresent()) {
-            delegate = new CacheFileSystemFactory(tracer, delegate, memoryFileSystemCache.orElseThrow(), keyProvider.orElseThrow());
+        if (memoryFileSystemCache.isPresent()) {
+            return new CacheFileSystemFactory(tracer, delegate, memoryFileSystemCache.orElseThrow(), keyProvider.orElseThrow());
         }
-        return new TracingFileSystemFactory(tracer, delegate);
+        return delegate;
     }
 }