Revert "OAK-12070 - Reduce memory consumption of azure segment stores (#2699)"

This reverts commit 082c3a3.

Author: reschke

oak-segment-aws/src/main/java/org/apache/jackrabbit/oak/segment/aws/AwsSegmentArchiveReader.java

@@ -20,8 +20,7 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.util.Iterator;
-import java.util.NoSuchElementException;
+import java.util.UUID;
 
 import org.apache.jackrabbit.oak.commons.Buffer;
 import org.apache.jackrabbit.oak.segment.remote.AbstractRemoteSegmentArchiveReader;
@@ -32,35 +31,47 @@ public class AwsSegmentArchiveReader extends AbstractRemoteSegmentArchiveReader
 
     private final S3Directory directory;
 
+    private final String archiveName;
+
+    private final long length;
+
     AwsSegmentArchiveReader(S3Directory directory, String archiveName, IOMonitor ioMonitor) throws IOException {
-        super(ioMonitor, archiveName, createEntryIterable(directory, archiveName));
+        super(ioMonitor);
         this.directory = directory;
+        this.archiveName = archiveName;
+        this.length = computeArchiveIndexAndLength();
+    }
+
+    @Override
+    public long length() {
+        return length;
+    }
+
+    @Override
+    public String getName() {
+        return archiveName;
     }
 
-    private static Iterable<ArchiveEntry> createEntryIterable(S3Directory directory, String archiveName) throws IOException{
+    @Override
+    protected long computeArchiveIndexAndLength() throws IOException {
+        long length = 0;
         Buffer buffer = directory.readObjectToBuffer(archiveName + ".idx", OFF_HEAP);
-        return () -> new Iterator<>() {
-            @Override
-            public boolean hasNext() {
-                return buffer.hasRemaining();
-            }
-
-            @Override
-            public ArchiveEntry next() {
-                if (!hasNext()) {
-                    throw new NoSuchElementException();
-                }
-
-                long msb = buffer.getLong();
-                long lsb = buffer.getLong();
-                int position = buffer.getInt();
-                int contentLength = buffer.getInt();
-                int generation = buffer.getInt();
-                int fullGeneration = buffer.getInt();
-                boolean compacted = buffer.get() != 0;
-                return new ArchiveEntry(new RemoteSegmentArchiveEntry(msb, lsb, position, contentLength, generation, fullGeneration, compacted));
-            }
-        };
+        while (buffer.hasRemaining()) {
+            long msb = buffer.getLong();
+            long lsb = buffer.getLong();
+            int position = buffer.getInt();
+            int contentLength = buffer.getInt();
+            int generation = buffer.getInt();
+            int fullGeneration = buffer.getInt();
+            boolean compacted = buffer.get() != 0;
+
+            RemoteSegmentArchiveEntry indexEntry = new RemoteSegmentArchiveEntry(msb, lsb, position, contentLength,
+                    generation, fullGeneration, compacted);
+            index.put(new UUID(indexEntry.getMsb(), indexEntry.getLsb()), indexEntry);
+            length += contentLength;
+        }
+
+        return length;
     }
 
     @Override

oak-segment-azure/src/main/java/org/apache/jackrabbit/oak/segment/azure/AzureSegmentArchiveReader.java

@@ -17,47 +17,65 @@
 package org.apache.jackrabbit.oak.segment.azure;
 
 import com.azure.storage.blob.BlobContainerClient;
+import com.azure.storage.blob.models.BlobItem;
 import com.azure.storage.blob.models.BlobStorageException;
 import com.azure.storage.blob.models.ListBlobsOptions;
 import com.azure.storage.blob.specialized.BlockBlobClient;
 import org.apache.jackrabbit.oak.commons.Buffer;
 import org.apache.jackrabbit.oak.segment.remote.AbstractRemoteSegmentArchiveReader;
+import org.apache.jackrabbit.oak.segment.remote.RemoteSegmentArchiveEntry;
 import org.apache.jackrabbit.oak.segment.spi.monitor.IOMonitor;
-import org.jetbrains.annotations.NotNull;
 
 import java.io.File;
 import java.io.IOException;
 import java.util.Map;
+import java.util.UUID;
 
 import static org.apache.jackrabbit.oak.segment.azure.AzureUtilities.readBufferFully;
 
 public class AzureSegmentArchiveReader extends AbstractRemoteSegmentArchiveReader {
 
     private final BlobContainerClient blobContainerClient;
 
+    private final long length;
+
+    private final String archiveName;
+
     private final String archivePathPrefix;
 
-    AzureSegmentArchiveReader(BlobContainerClient blobContainerClient, String rootPrefix, String archiveName, IOMonitor ioMonitor) {
-        super(ioMonitor, AzureUtilities.ensureNoTrailingSlash(archiveName),
-                createEntryIterable(blobContainerClient, AzureUtilities.asAzurePrefix(rootPrefix, archiveName)));
+    AzureSegmentArchiveReader(BlobContainerClient blobContainerClient, String rootPrefix, String archiveName, IOMonitor ioMonitor) throws IOException {
+        super(ioMonitor);
         this.blobContainerClient = blobContainerClient;
+        this.archiveName = AzureUtilities.ensureNoTrailingSlash(archiveName);
         this.archivePathPrefix = AzureUtilities.asAzurePrefix(rootPrefix, archiveName);
+        this.length = computeArchiveIndexAndLength();
+    }
+
+    @Override
+    public long length() {
+        return length;
+    }
+
+    @Override
+    public String getName() {
+        return archiveName;
     }
 
-    private static Iterable<ArchiveEntry> createEntryIterable(BlobContainerClient blobContainerClient, @NotNull String archivePathPrefix) {
+    @Override
+    protected long computeArchiveIndexAndLength() throws IOException {
+        long length = 0;
         ListBlobsOptions listBlobsOptions = new ListBlobsOptions();
         listBlobsOptions.setPrefix(archivePathPrefix);
-        return AzureUtilities.getBlobs(blobContainerClient, listBlobsOptions).stream()
-                .map(blobItem -> {
-                    Map<String, String> metadata = blobItem.getMetadata();
-                    int length = blobItem.getProperties().getContentLength().intValue();
-                    if (AzureBlobMetadata.isSegment(metadata)) {
-                        return new ArchiveEntry(AzureBlobMetadata.toIndexEntry(metadata, length));
-                    } else {
-                        return new ArchiveEntry(length);
-                    }
-                })
-                ::iterator;
+        for (BlobItem blob : AzureUtilities.getBlobs(blobContainerClient, listBlobsOptions)) {
+            Map<String, String> metadata = blob.getMetadata();
+            if (AzureBlobMetadata.isSegment(metadata)) {
+                RemoteSegmentArchiveEntry indexEntry = AzureBlobMetadata.toIndexEntry(metadata, blob.getProperties().getContentLength().intValue());
+                index.put(new UUID(indexEntry.getMsb(), indexEntry.getLsb()), indexEntry);
+            }
+            length += blob.getProperties().getContentLength();
+        }
+
+        return length;
     }
 
     @Override

oak-segment-azure/src/main/java/org/apache/jackrabbit/oak/segment/azure/v8/AzureSegmentArchiveReaderV8.java

@@ -22,37 +22,54 @@
 import java.io.IOException;
 import java.net.URISyntaxException;
 import java.util.Map;
+import java.util.UUID;
 
 import com.microsoft.azure.storage.StorageException;
+import com.microsoft.azure.storage.blob.CloudBlob;
 import com.microsoft.azure.storage.blob.CloudBlobDirectory;
 import com.microsoft.azure.storage.blob.CloudBlockBlob;
 
 import org.apache.jackrabbit.oak.commons.Buffer;
 import org.apache.jackrabbit.oak.segment.azure.AzureBlobMetadata;
 import org.apache.jackrabbit.oak.segment.remote.AbstractRemoteSegmentArchiveReader;
+import org.apache.jackrabbit.oak.segment.remote.RemoteSegmentArchiveEntry;
 import org.apache.jackrabbit.oak.segment.spi.monitor.IOMonitor;
 
 public class AzureSegmentArchiveReaderV8 extends AbstractRemoteSegmentArchiveReader {
 
     private final CloudBlobDirectory archiveDirectory;
 
+    private final long length;
+
     protected AzureSegmentArchiveReaderV8(CloudBlobDirectory archiveDirectory, IOMonitor ioMonitor) throws IOException {
-        super(ioMonitor, AzureUtilitiesV8.getName(archiveDirectory), createEntryIterable(archiveDirectory));
+        super(ioMonitor);
         this.archiveDirectory = archiveDirectory;
+        this.length = computeArchiveIndexAndLength();
+    }
+
+    @Override
+    public long length() {
+        return length;
     }
 
-    private static Iterable<ArchiveEntry> createEntryIterable(CloudBlobDirectory archiveDirectory) throws IOException {
-        return AzureUtilitiesV8.getBlobs(archiveDirectory).stream()
-                .map(blob -> {
-                    Map<String, String> metadata = blob.getMetadata();
-                    int length = (int) blob.getProperties().getLength();
-                    if (AzureBlobMetadata.isSegment(metadata)) {
-                        return new ArchiveEntry(AzureBlobMetadata.toIndexEntry(metadata, length));
-                    } else {
-                        return new ArchiveEntry(length);
-                    }
-                })
-                ::iterator;
+    @Override
+    public String getName() {
+        return AzureUtilitiesV8.getName(archiveDirectory);
+    }
+
+    @Override
+    protected long computeArchiveIndexAndLength() throws IOException {
+        long length = 0;
+        for (CloudBlob blob : AzureUtilitiesV8.getBlobs(archiveDirectory)) {
+            Map<String, String> metadata = blob.getMetadata();
+            if (AzureBlobMetadata.isSegment(metadata)) {
+                RemoteSegmentArchiveEntry indexEntry = AzureBlobMetadata.toIndexEntry(metadata, (int) blob.getProperties().getLength());
+                index.put(new UUID(indexEntry.getMsb(), indexEntry.getLsb()), indexEntry);
+            }
+            length += blob.getProperties().getLength();
+        }
+
+        return length;
     }
 
     @Override

oak-segment-remote/src/main/java/org/apache/jackrabbit/oak/segment/remote/AbstractRemoteSegmentArchiveReader.java

@@ -30,53 +30,20 @@
 
 import java.io.File;
 import java.io.IOException;
-import java.util.Collections;
-import java.util.Comparator;
-import java.util.LinkedList;
+import java.util.ArrayList;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
-import java.util.stream.Collectors;
 
 public abstract class AbstractRemoteSegmentArchiveReader implements SegmentArchiveReader {
-
     protected final IOMonitor ioMonitor;
 
-    /**
-     * Unordered immutable map of segment UUIDs to their corresponding archive entries.
-     */
-    private final Map<UUID, RemoteSegmentArchiveEntry> index;
-
-    /**
-     * The name of the archive.
-     */
-    private final String archiveName;
+    protected final Map<UUID, RemoteSegmentArchiveEntry> index = new LinkedHashMap<>();
 
-    /**
-     * The total size of the archive in bytes.
-     */
-    private final long length;
-
-    protected AbstractRemoteSegmentArchiveReader(IOMonitor ioMonitor, String archiveName, Iterable<ArchiveEntry> entries) {
+    public AbstractRemoteSegmentArchiveReader(IOMonitor ioMonitor) throws IOException {
         this.ioMonitor = ioMonitor;
-        this.archiveName = archiveName;
-
-        IndexBuilder indexBuilder = new IndexBuilder();
-        entries.forEach(indexBuilder::addEntry);
-        this.index = indexBuilder.createIndex();
-        this.length = indexBuilder.getLength();
-    }
-
-    @Override
-    public @NotNull String getName() {
-        return archiveName;
-    }
-
-    @Override
-    public long length() {
-        return length;
     }
 
     @Override
@@ -106,16 +73,9 @@ public boolean containsSegment(long msb, long lsb) {
         return index.containsKey(new UUID(msb, lsb));
     }
 
-    @Override
-    public Set<UUID> getSegmentUUIDs() {
-        return Collections.unmodifiableSet(index.keySet());
-    }
-
     @Override
     public List<SegmentArchiveEntry> listSegments() {
-        return index.values().stream()
-                .sorted(Comparator.comparing(RemoteSegmentArchiveEntry::getPosition))
-                .collect(Collectors.toList());
+        return new ArrayList<>(index.values());
     }
 
     @Override
@@ -142,6 +102,12 @@ public int getEntrySize(int size) {
         return size;
     }
 
+    /**
+     * Populates the archive index, summing up each entry's length.
+     * @return length, the total length of the archive
+     */
+    protected abstract long computeArchiveIndexAndLength() throws IOException;
+
     /**
      * Reads the segment from the remote storage.
      * @param segmentFileName, the name of the segment (msb + lsb) prefixed by its position in the archive
@@ -166,53 +132,4 @@ public int getEntrySize(int size) {
     public boolean isRemote() {
         return true;
     }
-
-    protected static final class ArchiveEntry {
-
-        private final RemoteSegmentArchiveEntry entry;
-
-        private final int length;
-
-        public ArchiveEntry(RemoteSegmentArchiveEntry entry) {
-            this.entry = entry;
-            this.length = entry.getLength();
-        }
-
-        public ArchiveEntry(int length) {
-            this.entry = null;
-            this.length = length;
-        }
-
-        int getLength() {
-            return length;
-        }
-
-        RemoteSegmentArchiveEntry getRemoteSegmentArchiveEntry() {
-            return entry;
-        }
-    }
-
-    private static final class IndexBuilder {
-
-        private final List<Map.Entry<UUID, RemoteSegmentArchiveEntry>> entries = new LinkedList<>();
-
-        private long length = 0;
-
-        private void addEntry(ArchiveEntry entry) {
-            RemoteSegmentArchiveEntry archiveEntry = entry.getRemoteSegmentArchiveEntry();
-            if (archiveEntry != null) {
-                this.entries.add(Map.entry(archiveEntry.getUuid(), archiveEntry));
-            }
-            this.length += entry.getLength();
-        }
-
-        @SuppressWarnings("unchecked")
-        private Map<UUID, RemoteSegmentArchiveEntry> createIndex() {
-            return Map.ofEntries(entries.toArray(Map.Entry[]::new));
-        }
-
-        private long getLength() {
-            return length;
-        }
-    }
 }