Merge pull request #761 from haileyajohnson/s3-perf

[5.x] fix s3 zarr performance issue

Author: lesserwhirls

cdm/core/src/main/java/ucar/nc2/internal/iosp/hdf4/H4header.java

@@ -58,6 +58,10 @@ public class H4header implements HdfHeaderIF {
   private static final long maxHeaderPos = 500000; // header's gotta be within this
 
   static boolean isValidFile(ucar.unidata.io.RandomAccessFile raf) throws IOException {
+    // fail fast on directory
+    if (raf.isDirectory()) {
+      return false;
+    }
     long pos = 0;
     long size = raf.length();
 

cdm/core/src/main/java/ucar/nc2/internal/iosp/hdf5/H5headerNew.java

@@ -121,6 +121,10 @@ public static void setDebugFlags(ucar.nc2.util.DebugFlags debugFlag) {
   private static final boolean transformReference = true;
 
   public static boolean isValidFile(RandomAccessFile raf) throws IOException {
+    // fail fast on directory
+    if (raf.isDirectory()) {
+      return false;
+    }
     // For HDF5, we need to search forward
     long filePos = 0;
     long size = raf.length();

cdm/core/src/main/java/ucar/nc2/internal/iosp/netcdf3/N3headerNew.java

@@ -30,6 +30,10 @@ public class N3headerNew {
   public static boolean debugHeaderSize; // see NetcdfFile.setDebugFlags
 
   public static boolean isValidFile(ucar.unidata.io.RandomAccessFile raf) throws IOException {
+    // fail fast on directory
+    if (raf.isDirectory()) {
+      return false;
+    }
     switch (NetcdfFileFormat.findNetcdfFormatType(raf)) {
       case NETCDF3:
       case NETCDF3_64BIT_OFFSET:

cdm/core/src/main/java/ucar/nc2/stream/NcStreamIosp.java

@@ -43,6 +43,10 @@ public class NcStreamIosp extends AbstractIOServiceProvider {
   private static final boolean debug = false;
 
   public boolean isValidFile(RandomAccessFile raf) throws IOException {
+    // fail fast on directory
+    if (raf.isDirectory()) {
+      return false;
+    }
     raf.seek(0);
     if (!readAndTest(raf, NcStream.MAGIC_START))
       return false; // must start with these 4 bytes

cdm/s3/src/main/java/thredds/filesystem/s3/ControllerS3.java

@@ -325,7 +325,7 @@ public MFile next() {
       MFile nextMFile = null;
       try {
         CdmS3Uri newUri = bucketUri.resolveNewKey(object.key());
-        nextMFile = new MFileS3(newUri);
+        nextMFile = new MFileS3(newUri, object.size(), object.lastModified().toEpochMilli());
       } catch (URISyntaxException e) {
         logger.warn("Cannot create MFile for {} in bucket {}", object.key(), bucketUri.getBucket(), e);
       }

cdm/s3/src/main/java/thredds/inventory/s3/MFileS3.java

@@ -34,24 +34,41 @@ public class MFileS3 implements MFile {
   private final String key;
   private final String delimiter;
 
+  private long length;
+  private long lastMod;
+
   private Object auxInfo;
 
+
   public MFileS3(String s3Uri) throws IOException {
+    this(s3Uri, -1, -1);
+  }
+
+  public MFileS3(CdmS3Uri s3Uri) {
+    this(s3Uri, -1, -1);
+  }
+
+  public MFileS3(String s3Uri, long len, long lm) throws IOException {
     try {
       cdmS3Uri = new CdmS3Uri(s3Uri);
     } catch (URISyntaxException e) {
       throw new IOException("Unable to create a CdmS3Uri from: " + s3Uri, e);
     }
     key = getKey();
     delimiter = getDelimiter();
+    length = len;
+    lastMod = lm;
+
     // This can take some time, so wait to execute until the first time headObjectResponse is accessed
     this.headObjectResponse = () -> getHeadObjectResponse();
   }
 
-  public MFileS3(CdmS3Uri s3Uri) {
+  public MFileS3(CdmS3Uri s3Uri, long len, long lm) {
     cdmS3Uri = s3Uri;
     key = getKey();
     delimiter = getDelimiter();
+    length = len;
+    lastMod = lm;
     // This can take some time, so wait to execute until the first time headObjectResponse is accessed
     this.headObjectResponse = () -> getHeadObjectResponse();
   }
@@ -87,12 +104,25 @@ private HeadObjectResponse getHeadObjectResponse() {
 
   @Override
   public long getLastModified() {
-    return headObjectResponse.get().lastModified().toEpochMilli();
+    // negative values indicate unavailable, get from head request
+    return this.lastMod < 0 ? updateLastModified() : this.lastMod;
+  }
+
+  /** Update last modified by fetching from a head request */
+  public long updateLastModified() {
+    this.lastMod = headObjectResponse.get().lastModified().toEpochMilli();
+    return this.lastMod;
   }
 
   @Override
   public long getLength() {
-    return headObjectResponse.get().contentLength();
+    return this.length < 0 ? updateLength() : this.length;
+  }
+
+  /** Update file length by fetching from a head request */
+  public long updateLength() {
+    this.length = headObjectResponse.get().contentLength();
+    return this.length;
   }
 
   @Override

cdm/s3/src/test/java/thredds/inventory/s3/TestMFileS3.java

@@ -93,24 +93,21 @@ public void justBucketOsdc() throws IOException {
   //
   @Test
   public void bucketAndKeyAws() throws IOException {
-    long lastModified = 1532465845000L;
-    checkWithBucketAndKey(AWS_G16_S3_OBJECT_1, G16_OBJECT_KEY_1, null, lastModified);
-    checkWithBucketAndKey(AWS_G16_S3_OBJECT_1 + DELIMITER_FRAGMENT, G16_NAME_1, "/", lastModified);
+    checkWithBucketAndKey(AWS_G16_S3_OBJECT_1, G16_OBJECT_KEY_1, null);
+    checkWithBucketAndKey(AWS_G16_S3_OBJECT_1 + DELIMITER_FRAGMENT, G16_NAME_1, "/");
   }
 
   @Test
   public void bucketAndKeyGcs() throws IOException {
-    long lastModified = 1504051532000L;
-    checkWithBucketAndKey(GCS_G16_S3_OBJECT_1, G16_OBJECT_KEY_1, null, lastModified);
-    checkWithBucketAndKey(GCS_G16_S3_OBJECT_1 + DELIMITER_FRAGMENT, G16_NAME_1, "/", lastModified);
+    checkWithBucketAndKey(GCS_G16_S3_OBJECT_1, G16_OBJECT_KEY_1, null);
+    checkWithBucketAndKey(GCS_G16_S3_OBJECT_1 + DELIMITER_FRAGMENT, G16_NAME_1, "/");
   }
 
   @Test
   @Category(NotPullRequest.class)
   public void bucketAndKeyOsdc() throws IOException {
-    long lastModified = 1611593614000L;
-    checkWithBucketAndKey(OSDC_G16_S3_OBJECT_1, OSDC_G16_OBJECT_KEY_1, null, lastModified);
-    checkWithBucketAndKey(OSDC_G16_S3_OBJECT_1 + DELIMITER_FRAGMENT, G16_NAME_1, "/", lastModified);
+    checkWithBucketAndKey(OSDC_G16_S3_OBJECT_1, OSDC_G16_OBJECT_KEY_1, null);
+    checkWithBucketAndKey(OSDC_G16_S3_OBJECT_1 + DELIMITER_FRAGMENT, G16_NAME_1, "/");
   }
 
   @Test
@@ -184,8 +181,7 @@ private void checkWithBucket(String cdmS3Uri) throws IOException {
     assertThat(parent).isNull();
   }
 
-  private void checkWithBucketAndKey(String cdmS3Uri, String expectedName, String delimiter, long expectedLastModified)
-      throws IOException {
+  private void checkWithBucketAndKey(String cdmS3Uri, String expectedName, String delimiter) throws IOException {
     logger.info("Checking {}", cdmS3Uri);
     MFile mFile = new MFileS3(cdmS3Uri);
     assertThat(mFile.getPath()).isEqualTo(cdmS3Uri);
@@ -197,7 +193,6 @@ private void checkWithBucketAndKey(String cdmS3Uri, String expectedName, String
       assertThat(mFile.getParent()).isNull();
     }
     assertThat(mFile.isDirectory()).isFalse();
-    assertThat(mFile.getLastModified()).isEqualTo(expectedLastModified);
     assertThat(mFile.getLength()).isEqualTo(G16_OBJECT_1_SIZE);
   }
 

cdm/zarr/src/main/java/ucar/nc2/iosp/zarr/ZarrHeader.java

@@ -29,6 +29,7 @@ public class ZarrHeader {
 
   private final RandomAccessDirectory rootRaf;
   private final Group.Builder rootGroup;
+  private final String rootLocation;
 
   private List<Attribute> attrs;
 
@@ -37,6 +38,7 @@ public class ZarrHeader {
   public ZarrHeader(RandomAccessDirectory raf, Group.Builder rootGroup) {
     this.rootRaf = raf;
     this.rootGroup = rootGroup;
+    this.rootLocation = ZarrPathUtils.trimLocation(this.rootRaf.getLocation());
     this.attrs = null;
   }
 
@@ -46,8 +48,7 @@ public ZarrHeader(RandomAccessDirectory raf, Group.Builder rootGroup) {
    * @throws IOException
    */
   public void read() throws IOException {
-    String location = ZarrPathUtils.trimLocation(this.rootRaf.getLocation());
-    List<RandomAccessDirectoryItem> items = this.rootRaf.getFilesInPath(location);
+    List<RandomAccessDirectoryItem> items = this.rootRaf.getFilesInPath(this.rootLocation);
     // because items are ordered alphabetically, it is safe to assume groups will be created before subgroups and vars
     // However, we need to assure we make attrs before their groups and vars
     // assumes no files will be read between .zarray and .zattrs
@@ -97,7 +98,7 @@ private void makeGroup(RandomAccessDirectoryItem item) throws ZarrFormatExceptio
     // make new Group
     Group.Builder group = Group.builder();
     String location = ZarrPathUtils.trimLocation(item.getLocation());
-    if (location.equals(this.rootRaf.getLocation() + ZarrKeys.ZGROUP)) {
+    if (location.equals(this.rootLocation + ZarrKeys.ZGROUP)) {
       group = this.rootGroup;
     }
     // set Group name
@@ -199,7 +200,7 @@ private List<Attribute> makeAttributes(RandomAccessDirectoryItem item) throws IO
    */
   private Group.Builder findGroup(String location) throws ZarrFormatException {
     // set Group parent
-    String groupName = ZarrPathUtils.getParentGroupNameFromPath(location, rootRaf.getLocation());
+    String groupName = ZarrPathUtils.getParentGroupNameFromPath(location, this.rootLocation);
     return this.rootGroup.findGroupNested(groupName).orElseThrow(ZarrFormatException::new);
   }
 

cdm/zarr/src/test/java/ucar/nc2/iosp/zarr/TestZarrIosp.java

@@ -38,19 +38,19 @@ public class TestZarrIosp {
   @Test
   public void testIsValidFile() throws IOException {
     ZarrIosp iosp = new ZarrIosp();
+    // zarr stores
     assertThat(iosp.isValidFile(NetcdfFiles.getRaf(DIRECTORY_STORE_URI, -1))).isTrue();
     assertThat(iosp.isValidFile(NetcdfFiles.getRaf(ZIP_STORE_URI, -1))).isTrue();
-    // TODO: S3 too slow for large set of objects
-    // assertThat(iosp.isValidFile(NetcdfFiles.getRaf(OBJECT_STORE_ZARR_URI, -1))).isTrue();
-
+    assertThat(iosp.isValidFile(NetcdfFiles.getRaf(OBJECT_STORE_ZARR_URI, -1))).isTrue();
+    // not zarr stores
     assertThat(iosp.isValidFile(NetcdfFiles.getRaf(DIRECTORY_STORE_URI + "/.zgroup", -1))).isFalse();
   }
 
   @Test
   public void testBuildNcfile() throws IOException {
     _testBuildNcfile(DIRECTORY_STORE_URI);
     _testBuildNcfile(ZIP_STORE_URI);
-    // _testBuildNcfile(OBJECT_STORE_ZARR_URI);
+    _testBuildNcfile(OBJECT_STORE_ZARR_URI);
   }
 
   private void _testBuildNcfile(String location) throws IOException {