skoppu22 commented on code in PR #206: URL: https://github.com/apache/cassandra-analytics/pull/206#discussion_r3776230765
########## cassandra-analytics-core/src/main/java/org/apache/cassandra/spark/data/S3CassandraDataLayer.java: ########## @@ -0,0 +1,1769 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.cassandra.spark.data; + +import com.esotericsoftware.kryo.Kryo; +import com.esotericsoftware.kryo.io.Input; +import com.esotericsoftware.kryo.io.Output; + +import com.google.common.annotations.VisibleForTesting; +import com.google.common.cache.Cache; +import com.google.common.cache.CacheBuilder; +import com.google.common.collect.Range; + +import org.apache.cassandra.analytics.stats.Stats; +import org.apache.cassandra.bridge.CassandraBridge; +import org.apache.cassandra.bridge.CassandraBridgeFactory; +import org.apache.cassandra.bridge.CassandraVersion; +import org.apache.cassandra.clients.ExecutorHolder; +import org.apache.cassandra.spark.config.SchemaFeature; +import org.apache.cassandra.spark.config.SchemaFeatureSet; +import org.apache.cassandra.spark.data.backup.BackupReader; +import org.apache.cassandra.spark.data.backup.BackupReaderRegistry; +import org.apache.cassandra.spark.sparksql.RowBuilder; +import org.apache.cassandra.spark.sparksql.SnapshotTimestampDecorator; +import org.apache.cassandra.spark.data.partitioner.CassandraInstance; +import org.apache.cassandra.spark.data.partitioner.CassandraRing; +import org.apache.cassandra.spark.data.partitioner.ConsistencyLevel; +import org.apache.cassandra.spark.data.partitioner.Partitioner; +import org.apache.cassandra.spark.data.partitioner.TokenPartitioner; +import org.apache.cassandra.spark.sparksql.SparkCustomMetricsStats; +import org.apache.cassandra.spark.utils.TimeProvider; +import org.apache.cassandra.spark.utils.S3SnapshotTimeProvider; +import org.apache.cassandra.spark.utils.ScalaFunctions; +import org.apache.cassandra.spark.utils.streaming.BufferingInputStream; +import org.apache.cassandra.spark.utils.streaming.CassandraFileSource; +import org.apache.cassandra.spark.utils.streaming.StreamConsumer; +import org.apache.cassandra.spark.common.S3SizingFactory; + +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.types.DataType; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.util.ShutdownHookManager; + +import org.apache.commons.lang.NotImplementedException; + +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.io.InputStream; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.io.Serializable; +import java.math.BigInteger; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.function.Predicate; +import java.util.stream.Collectors; +import java.util.stream.Stream; +import java.util.OptionalLong; + +import org.apache.cassandra.bridge.TokenRange; +import org.apache.cassandra.spark.utils.RangeUtils; + +/** + * S3-backed CassandraDataLayer. The concrete backup format is provided by a pluggable + * {@link BackupReader} resolved via {@link BackupReaderRegistry} using the + * {@code backupReaderType} option (no default; callers must register a factory). + * <p> + * Assumes that Murmur3Partitioner is used. The backup reader is expected to return a list of + * Cassandra instances per individual vnode. + */ +public class S3CassandraDataLayer extends PartitionedDataLayer implements Serializable +{ + private static final long serialVersionUID = 1997L; + + private static final Logger LOGGER = LoggerFactory.getLogger(S3CassandraDataLayer.class); + + /** + * JVM-wide intern cache that canonicalizes {@link BackupReader} instances per executor, so + * all tasks reading the same manifest+S3 identity share a single reader (and its + * implementation-specific cache, e.g. {@code sstableInfoCache}) instead of one copy per + * deserialized task. + * <p> + * Key: {@code (cluster, keyspace, table, datacenter, earliestEpoch, latestEpoch, + * manifestFingerprint, s3Region, s3Bucket, s3EndpointOverride, s3CredentialsFingerprint, + * s3HttpMaxConcurrency)}. The fingerprint (SHA-256 over sorted {@code (nodeId, epoch)} + * pairs) is the authoritative manifest identity — without it, two manifest sets sharing + * the same {@code (min, max)} epochs alias and silently read stale SSTables (real failure + * mode with 3+ nodes when a middle node rolls independently). S3 identity fields prevent a + * caller with a different {@code s3Config} from reading through the wrong endpoint. + * <p> + * Credentials in the key isolate IAM principals and force a fresh reader after static-key + * rotation. For prod (EMR instance role / IRSA / IMDSv2) the access keys are null and the + * fingerprint collapses to a constant; STS rotation happens inside the SDK and does not + * invalidate the key. + * <p> + * Values are weak ({@link CacheBuilder#weakValues()}), so canonical readers are GC'd once + * no layer references them. {@link Cache#get(Object, java.util.concurrent.Callable)} is + * the atomic install-or-return primitive and pins the returned value across the call. + * <p> + * BackupReader read methods receive the task {@link Stats}, preserving S3 GET/HEAD metric + * attribution when tasks share a canonical reader. + */ + private static final class ReaderInternCache + { + // weakValues: entries auto-evict once no layer references the canonical reader. + // Reachability is the correct lifecycle signal here; do not add time-based eviction. + private static final Cache<Key, BackupReader> CACHE = + CacheBuilder.newBuilder().weakValues().build(); + + private static BackupReader canonicalize(String clusterName, + String keyspace, + String table, + String datacenter, + long earliestSnapshotEpochSecond, + long latestSnapshotEpochSecond, + @NotNull BackupReader fresh) + { + // Bypass when manifest identity isn't fully materialized: production sets all three + // (both epochs and fingerprint) inside initializeS3BackupReader. Hitting any of + // these branches means we'd otherwise install a partially-keyed entry. + if (earliestSnapshotEpochSecond <= 0 || latestSnapshotEpochSecond <= 0) + { + return fresh; + } + S3ClientConfig fingerprintConfig = fresh.s3Config(); + String fingerprintBucket = fresh.bucket(); + if (fingerprintConfig == null || fingerprintBucket == null) + { + // Mock readers in reflection-driven tests land here. + return fresh; + } + String manifestFingerprint = fresh.getManifestFingerprint(clusterName); + if (manifestFingerprint == null || manifestFingerprint.isEmpty()) + { + return fresh; + } + + Key key = Key.from(clusterName, keyspace, table, datacenter, + earliestSnapshotEpochSecond, latestSnapshotEpochSecond, + manifestFingerprint, + fingerprintConfig, fingerprintBucket); + + try + { + // Cache.get(key, loader) is atomic install-or-return; losing-candidate fresh + // readers never publish and are GC-eligible immediately on return. + BackupReader canonical = CACHE.get(key, () -> { + LOGGER.info("ReaderInternCache: installed canonical BackupReader " + + "cluster={} keyspace={} table={} datacenter={} earliestEpoch={} latestEpoch={} " + + "manifestFingerprint={} region={} bucket={} endpoint={} maxConcurrency={} identity={}", + clusterName, keyspace, table, datacenter, + earliestSnapshotEpochSecond, latestSnapshotEpochSecond, + manifestFingerprint, + fingerprintConfig.s3Region(), fingerprintBucket, + fingerprintConfig.s3EndpointOverride(), + fingerprintConfig.s3HttpMaxConcurrency(), + System.identityHashCode(fresh)); + return fresh; + }); + + if (canonical != fresh) + { + LOGGER.debug("ReaderInternCache: reused canonical BackupReader " + + "cluster={} keyspace={} table={} datacenter={} earliestEpoch={} latestEpoch={} " + + "manifestFingerprint={} region={} bucket={} canonicalIdentity={} discardedFreshIdentity={}", + clusterName, keyspace, table, datacenter, + earliestSnapshotEpochSecond, latestSnapshotEpochSecond, + manifestFingerprint, + fingerprintConfig.s3Region(), fingerprintBucket, + System.identityHashCode(canonical), System.identityHashCode(fresh)); + + // Defense in depth against a future Key regression that aliases buckets. + if (!fingerprintBucket.equals(canonical.bucket())) + { + LOGGER.error("ReaderInternCache: bucket mismatch on canonical reader for " + + "key={}. Canonical bucket={} fresh bucket={}. Replacing canonical " + + "with fresh reader to avoid wrong-bucket reads.", + key, canonical.bucket(), fingerprintBucket); + CACHE.put(key, fresh); + return fresh; + } + } + return canonical; + } + catch (java.util.concurrent.ExecutionException e) + { + // Loader does not throw checked exceptions; unreachable today. + throw new RuntimeException("ReaderInternCache loader unexpectedly threw", e.getCause()); + } + } + + @VisibleForTesting + static void clearForTesting() + { + CACHE.invalidateAll(); + // Drain weak-ref eviction queue so sizeForTesting() is stable. + CACHE.cleanUp(); + } + + @VisibleForTesting + static long sizeForTesting() + { + CACHE.cleanUp(); + return CACHE.size(); + } + + private static final class Key + { + private final String clusterName; + private final String keyspace; + private final String table; + private final String datacenter; + private final long earliestSnapshotEpochSecond; + private final long latestSnapshotEpochSecond; + // SHA-256 over sorted (nodeId, autosnapEpoch) pairs. Disambiguates manifest sets + // that share the same (min, max) epochs but differ on a middle node's rotation. + private final String manifestFingerprint; + private final String s3Region; + private final String s3Bucket; + @Nullable + private final String s3EndpointOverride; + // "<accessKey>|<secretHash>", mirroring S3ClientCache.getCacheKey. Null/empty + // access keys normalize to "default", empty secrets to hash 0; raw secret never + // enters the key. Prod (EMR/IRSA/IMDSv2) collapses to a constant "default|0". + private final String s3CredentialsFingerprint; + private final int s3HttpMaxConcurrency; + + private Key(String clusterName, String keyspace, String table, String datacenter, + long earliestSnapshotEpochSecond, long latestSnapshotEpochSecond, + String manifestFingerprint, + String s3Region, String s3Bucket, @Nullable String s3EndpointOverride, + String s3CredentialsFingerprint, int s3HttpMaxConcurrency) + { + this.clusterName = clusterName; + this.keyspace = keyspace; + this.table = table; + this.datacenter = datacenter; + this.earliestSnapshotEpochSecond = earliestSnapshotEpochSecond; + this.latestSnapshotEpochSecond = latestSnapshotEpochSecond; + this.manifestFingerprint = manifestFingerprint; + this.s3Region = s3Region; + this.s3Bucket = s3Bucket; + this.s3EndpointOverride = s3EndpointOverride; + this.s3CredentialsFingerprint = s3CredentialsFingerprint; + this.s3HttpMaxConcurrency = s3HttpMaxConcurrency; + } + + static Key from(String clusterName, String keyspace, String table, String datacenter, + long earliestSnapshotEpochSecond, long latestSnapshotEpochSecond, + String manifestFingerprint, + S3ClientConfig s3Config, String bucket) + { + return new Key(clusterName, keyspace, table, datacenter, + earliestSnapshotEpochSecond, latestSnapshotEpochSecond, + manifestFingerprint, + s3Config.s3Region(), bucket, s3Config.s3EndpointOverride(), + credentialsFingerprint(s3Config), + s3Config.s3HttpMaxConcurrency()); + } + + // Mirrors S3ClientCache.getCacheKey credential portion: accessKey|secretHash. + private static String credentialsFingerprint(S3ClientConfig s3Config) + { + String accessKey = s3Config.s3AccessKeyId(); + String secret = s3Config.s3SecretAccessKey(); + String normalizedAccessKey = (accessKey != null && !accessKey.isEmpty()) ? accessKey : "default"; + int secretHash = (secret != null && !secret.isEmpty()) ? secret.hashCode() : 0; + return normalizedAccessKey + "|" + secretHash; + } + + @Override + public boolean equals(Object o) + { + if (this == o) + { + return true; + } + if (!(o instanceof Key)) + { + return false; + } + Key other = (Key) o; + return earliestSnapshotEpochSecond == other.earliestSnapshotEpochSecond + && latestSnapshotEpochSecond == other.latestSnapshotEpochSecond + && s3HttpMaxConcurrency == other.s3HttpMaxConcurrency + && Objects.equals(clusterName, other.clusterName) + && Objects.equals(keyspace, other.keyspace) + && Objects.equals(table, other.table) + && Objects.equals(datacenter, other.datacenter) + && Objects.equals(manifestFingerprint, other.manifestFingerprint) + && Objects.equals(s3Region, other.s3Region) + && Objects.equals(s3Bucket, other.s3Bucket) + && Objects.equals(s3EndpointOverride, other.s3EndpointOverride) + && Objects.equals(s3CredentialsFingerprint, other.s3CredentialsFingerprint); + } + + @Override + public int hashCode() + { + return Objects.hash(clusterName, keyspace, table, datacenter, + earliestSnapshotEpochSecond, latestSnapshotEpochSecond, + manifestFingerprint, + s3Region, s3Bucket, s3EndpointOverride, + s3CredentialsFingerprint, s3HttpMaxConcurrency); + } + + @Override + public String toString() + { + // Error-logging only; raw credentials never appear (already a hash). + return "ReaderInternCache.Key{cluster=" + clusterName + + " keyspace=" + keyspace + + " table=" + table + + " dc=" + datacenter + + " earliestEpoch=" + earliestSnapshotEpochSecond + + " latestEpoch=" + latestSnapshotEpochSecond + + " manifestFingerprint=" + manifestFingerprint + + " region=" + s3Region + + " bucket=" + s3Bucket + + " endpoint=" + s3EndpointOverride + + " credsFingerprint=" + s3CredentialsFingerprint + + " maxConcurrency=" + s3HttpMaxConcurrency + + "}"; + } + } + } + + /** + * Test-only: clear the JVM-wide reader intern cache between tests. Production must never + * call this — it will not free memory (canonical readers stay referenced by live layers) + * and the next deserialization will install a duplicate. + */ + @VisibleForTesting + public static void clearReaderInternCacheForTesting() + { + ReaderInternCache.clearForTesting(); + } + + @VisibleForTesting + public static long readerInternCacheSizeForTesting() + { + return ReaderInternCache.sizeForTesting(); + } + + /** + * Test-only entry to {@link ReaderInternCache#canonicalize}, bypassing layer construction. + * Layer constructors register a Spark shutdown hook that pins the layer (and reader) for + * JVM lifetime, which would defeat weak-value GC assertions. + * + * @param clusterName logical cluster identity + * @param keyspace Cassandra keyspace + * @param table Cassandra table + * @param datacenter datacenter + * @param earliestSnapshotEpochSecond earliest contributing snapshot epoch (seconds) + * @param latestSnapshotEpochSecond latest contributing snapshot epoch (seconds) + * @param fresh freshly-constructed candidate reader to canonicalize + * @return the canonical {@link BackupReader} (may be {@code fresh}, or a previously interned instance) + */ + @VisibleForTesting + public static BackupReader canonicalizeForTesting(String clusterName, + String keyspace, + String table, + String datacenter, + long earliestSnapshotEpochSecond, + long latestSnapshotEpochSecond, + BackupReader fresh) + { + return ReaderInternCache.canonicalize(clusterName, keyspace, table, datacenter, + earliestSnapshotEpochSecond, latestSnapshotEpochSecond, + fresh); + } + + private String clusterName; + private String keyspace; + private String table; + private String s3Region; + private String s3Bucket; + @Nullable + private String s3EndpointOverride; + @Nullable + private String s3AccessKeyId; + @Nullable + private String s3SecretAccessKey; + + protected transient CassandraBridge bridge; + + private CassandraRing ring; + private TokenPartitioner tokenPartitioner; + protected CqlTable cqlTable; + + @Nullable + protected String lastModifiedTimestampField; + @Nullable + protected String snapshotTimestampField; + protected List<SchemaFeature> requestedFeatures; + protected int sstableS3ReadTimeoutSeconds; + protected long latestSnapshotEpochSecond; + + // Data.db ranged-GET buffer sizes. Carried as instance fields (not just on S3DataSourceClientConfig) + // so they survive Spark serialization to executors. + private long dataChunkBufferSize = org.apache.cassandra.spark.utils.Properties.DEFAULT_S3_DATA_CHUNK_BUFFER_SIZE; + private long dataMaxBufferSize = org.apache.cassandra.spark.utils.Properties.DEFAULT_S3_DATA_MAX_BUFFER_SIZE; + + // Switch for Data.db ranged-GET delivery. Default true: Data.db reads use the + // AsyncResponseTransformer.toPublisher() streaming path. When false, + // AsyncResponseTransformer.toBytes() is used (single materialized byte[] per ranged GET). + // Non-Data file types and mutable metadata reads always use their existing paths regardless of this + // flag. Carried as instance field so it survives Spark serialization to executors. + private boolean sstableDataPublisherReadEnabled = true; + + // SSTable metadata cache sizes forwarded to {@code SSTableCache} via JVM sysprops. Carried here so + // executor-side deserialization can re-apply them; defaults mirror S3DataSourceClientConfig. + private int sstableCacheSummaryMaxEntries = 32768; + private int sstableCacheIndexMaxEntries = 16384; + private int sstableCacheStatsMaxEntries = 16384; + private int sstableCacheFilterMaxEntries = 16384; + private int sstableCacheCompressionInfoMaxEntries = 16384; + + private boolean sstableTokenIndexEnabled = false; + private int sstableTokenIndexPrebuildPartitions = 0; + private int sstableTokenIndexPrebuildPerTaskConcurrency = 4; + private transient SSTableTokenIndex sstableTokenIndex; + + private BackupReader s3BackupReader = null; + protected transient TimeProvider timeProvider; + private transient Stats stats; + private S3DataSourceClientConfig s3Config; + // Selects the BackupReaderFactory. Non-final so readObject can reassign it on executors. + private String backupReaderType; + + public S3CassandraDataLayer(@NotNull S3DataSourceClientConfig config) + { + super(config.consistencyLevel(), config.datacenter()); + this.s3Config = config; + this.clusterName = config.clusterName(); + this.keyspace = config.keyspace(); + this.table = config.table(); + this.s3Region = config.s3Region(); + this.s3Bucket = config.s3Bucket(); + this.s3EndpointOverride = config.s3EndpointOverride(); + this.s3AccessKeyId = config.s3AccessKeyId(); + this.s3SecretAccessKey = config.s3SecretAccessKey(); + this.sstableS3ReadTimeoutSeconds = config.sstableS3ReadTimeoutSeconds(); + this.dataChunkBufferSize = config.s3DataChunkBufferSize(); + this.dataMaxBufferSize = config.s3DataMaxBufferSize(); + this.sstableDataPublisherReadEnabled = config.sstableDataPublisherReadEnabled(); + this.sstableCacheSummaryMaxEntries = config.sstableCacheSummaryMaxEntries(); + this.sstableCacheIndexMaxEntries = config.sstableCacheIndexMaxEntries(); + this.sstableCacheStatsMaxEntries = config.sstableCacheStatsMaxEntries(); + this.sstableCacheFilterMaxEntries = config.sstableCacheFilterMaxEntries(); + this.sstableCacheCompressionInfoMaxEntries = config.sstableCacheCompressionInfoMaxEntries(); + this.sstableTokenIndexEnabled = config.sstableTokenIndexEnabled(); + this.sstableTokenIndexPrebuildPartitions = config.sstableTokenIndexPrebuildPartitions(); + this.sstableTokenIndexPrebuildPerTaskConcurrency = config.sstableTokenIndexPrebuildPerTaskConcurrency(); + this.backupReaderType = config.backupReaderType(); + + // Driver-side apply; executor side is covered from readObject / Kryo Serializer.read. + applySSTableCacheSystemProperties(); + + LOGGER.info("Initializing S3CassandraDataLayer for cluster={}, keyspace={}, table={}, " + + "dataChunkBufferSize={} bytes, dataMaxBufferSize={} bytes, " + + "sstableDataPublisherReadEnabled={}", + clusterName, keyspace, table, dataChunkBufferSize, dataMaxBufferSize, + sstableDataPublisherReadEnabled); + + // Initialize stats before initializing s3BackupReader such that stats can be passed to s3BackupReader + this.stats = new SparkCustomMetricsStats(); + + initializeS3BackupReader(); + + // list Cassandra instances in S3 bucket + final List<CassandraInstance> instances = s3BackupReader.instances(clusterName, config.keyspace(), config.table(), config.datacenter()); + // build CassandraRing and TokenPartitioner. Prefer rack-aware authoritative replica + // placement from the BackupReader; fall back to the naive (rack-unaware) ring when + // none is available. Exceptions from the reader signal a genuine integrity issue and + // must surface — see BackupReader#buildRackAwareReplicas for the contract. + final Partitioner partitioner = Partitioner.Murmur3Partitioner; Review Comment: Murmur3Partitioner is hardcoded here. Can we provide a hook to configure partitioners, something like `Partitioner partitioner = s3BackupReader.partitioner();` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
