Skip to content

Commit e495f4b

Browse files
sixpluszeroclaude
sixpluszero
and
claude
committed
[da-vinci][server] Add bounded hot transient record cache for large records
during ingestion Introduce a shared Caffeine cache per StoreIngestionTask that retains large transient records across consumer poll boundaries, reducing expensive DB lookups for chunked values in AA-replication and write-compute stores. The cache is gated by a per-store flag (transientRecordCacheEnabled) and server-level configs for max weight (default 32MB) and minimum value size (default 100KB). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 44476a0 commit e495f4b

14 files changed

Lines changed: 782 additions & 6 deletions

File tree

‎build.gradle‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,7 @@ subprojects {
327327
// when actually using the new protocol. Example to pin KME to v12 when introducing v13:
328328
// project(':internal:venice-common').file('src/main/resources/avro/KafkaMessageEnvelope/v12', PathValidation.DIRECTORY)
329329
def versionOverrides = [
330-
project(':internal:venice-common').file('src/main/resources/avro/StoreMetaValue/v40', PathValidation.DIRECTORY)
330+
project(':internal:venice-common').file('src/main/resources/avro/StoreMetaValue/v42', PathValidation.DIRECTORY)
331331
]
332332

333333
def schemaDirs = [sourceDir]

‎clients/da-vinci-client/src/main/java/com/linkedin/davinci/config/VeniceServerConfig.java‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,8 @@
132132
import static com.linkedin.venice.ConfigKeys.SERVER_INGESTION_OTEL_STATS_ENABLED;
133133
import static com.linkedin.venice.ConfigKeys.SERVER_INGESTION_TASK_MAX_IDLE_COUNT;
134134
import static com.linkedin.venice.ConfigKeys.SERVER_INGESTION_TASK_REUSABLE_OBJECTS_STRATEGY;
135+
import static com.linkedin.venice.ConfigKeys.SERVER_INGESTION_TRANSIENT_RECORD_CACHE_MAX_WEIGHT;
136+
import static com.linkedin.venice.ConfigKeys.SERVER_INGESTION_TRANSIENT_RECORD_CACHE_MIN_VALUE_SIZE;
135137
import static com.linkedin.venice.ConfigKeys.SERVER_KAFKA_CONSUMER_OFFSET_COLLECTION_ENABLED;
136138
import static com.linkedin.venice.ConfigKeys.SERVER_KAFKA_MAX_POLL_RECORDS;
137139
import static com.linkedin.venice.ConfigKeys.SERVER_LAG_BASED_REPLICA_AUTO_RESUBSCRIBE_ENABLED;
@@ -711,6 +713,8 @@ public class VeniceServerConfig extends VeniceClusterConfig {
711713
private final int lagMonitorCleanupCycle;
712714
private final boolean readQuotaInitializationFallbackEnabled;
713715
private final boolean ingestionProgressLoggingEnabled;
716+
private final long transientRecordCacheMaxWeight;
717+
private final int transientRecordCacheMinValueSize;
714718

715719
public VeniceServerConfig(VeniceProperties serverProperties) throws ConfigurationException {
716720
this(serverProperties, Collections.emptyMap());
@@ -1227,6 +1231,10 @@ public VeniceServerConfig(VeniceProperties serverProperties, Map<String, Map<Str
12271231
this.readQuotaInitializationFallbackEnabled =
12281232
serverProperties.getBoolean(SERVER_READ_QUOTA_INITIALIZATION_FALLBACK_ENABLED, true);
12291233
this.ingestionProgressLoggingEnabled = serverProperties.getBoolean(POSITIONAL_PROGRESS_LOGGING_ENABLED, false);
1234+
this.transientRecordCacheMaxWeight =
1235+
serverProperties.getLong(SERVER_INGESTION_TRANSIENT_RECORD_CACHE_MAX_WEIGHT, 32 * 1024 * 1024L);
1236+
this.transientRecordCacheMinValueSize =
1237+
serverProperties.getInt(SERVER_INGESTION_TRANSIENT_RECORD_CACHE_MIN_VALUE_SIZE, 100 * 1024);
12301238
}
12311239

12321240
List<Double> extractThrottleLimitFactorsFor(VeniceProperties serverProperties, String configKey) {
@@ -2225,4 +2233,12 @@ public boolean isReadQuotaInitializationFallbackEnabled() {
22252233
public boolean isIngestionProgressLoggingEnabled() {
22262234
return ingestionProgressLoggingEnabled;
22272235
}
2236+
2237+
public long getTransientRecordCacheMaxWeight() {
2238+
return transientRecordCacheMaxWeight;
2239+
}
2240+
2241+
public int getTransientRecordCacheMinValueSize() {
2242+
return transientRecordCacheMinValueSize;
2243+
}
22282244
}

‎clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/PartitionConsumptionState.java‎

Lines changed: 50 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import static java.util.concurrent.TimeUnit.MINUTES;
44

5+
import com.github.benmanes.caffeine.cache.Cache;
56
import com.linkedin.davinci.compression.KeyUrnCompressor;
67
import com.linkedin.davinci.compression.UrnDictV1;
78
import com.linkedin.davinci.helix.LeaderFollowerPartitionStateModel;
@@ -38,6 +39,7 @@
3839
import java.util.concurrent.CompletableFuture;
3940
import java.util.concurrent.Future;
4041
import java.util.concurrent.TimeUnit;
42+
import java.util.concurrent.atomic.AtomicLong;
4143
import java.util.concurrent.atomic.AtomicReference;
4244
import java.util.function.BooleanSupplier;
4345
import java.util.stream.Collectors;
@@ -208,6 +210,11 @@ enum LatchStatus {
208210
*/
209211
private final Map<ByteArrayKey, TransientRecord> transientRecordMap = new VeniceConcurrentHashMap<>();
210212

213+
/** Shared Caffeine cache across partitions for hot transient records. Nullable when disabled. */
214+
private final Cache<ByteArrayKey, TransientRecord> hotRecordCache;
215+
private final int minValueSizeForHotCache;
216+
private final AtomicLong hotRecordCacheHitCount = new AtomicLong(0);
217+
211218
/**
212219
* This field is used to track whether the last queued record has been fully processed or not.
213220
* For Leader role, it is redundant from {@literal ProducedRecord#persistedToDBFuture} since it is tracking
@@ -338,6 +345,17 @@ public PartitionConsumptionState(
338345
PubSubContext pubSubContext,
339346
boolean hybrid,
340347
Schema keySchema) {
348+
this(partitionReplica, offsetRecord, pubSubContext, hybrid, keySchema, null, 0);
349+
}
350+
351+
public PartitionConsumptionState(
352+
PubSubTopicPartition partitionReplica,
353+
OffsetRecord offsetRecord,
354+
PubSubContext pubSubContext,
355+
boolean hybrid,
356+
Schema keySchema,
357+
Cache<ByteArrayKey, TransientRecord> hotRecordCache,
358+
int minValueSizeForHotCache) {
341359
LOGGER.info("Creating PCS for replica: {}", partitionReplica);
342360

343361
this.partitionReplica = Objects.requireNonNull(partitionReplica, "TopicPartition cannot be null when creating PCS");
@@ -350,6 +368,8 @@ public PartitionConsumptionState(
350368
this.keySchema = keySchema;
351369
this.offsetRecord = offsetRecord;
352370
this.pubSubContext = pubSubContext;
371+
this.hotRecordCache = hotRecordCache;
372+
this.minValueSizeForHotCache = minValueSizeForHotCache;
353373
this.errorReported = false;
354374
this.lagCaughtUp = false;
355375
this.lagCaughtUpTimeInMs = 0;
@@ -768,10 +788,39 @@ public void setTransientRecord(
768788
}
769789

770790
transientRecordMap.put(ByteArrayKey.wrap(key), transientRecord);
791+
792+
if (hotRecordCache != null && valueLen >= minValueSizeForHotCache) {
793+
hotRecordCache.put(buildHotCacheKey(key), transientRecord);
794+
}
771795
}
772796

773797
public TransientRecord getTransientRecord(byte[] key) {
774-
return transientRecordMap.get(ByteArrayKey.wrap(key));
798+
TransientRecord record = transientRecordMap.get(ByteArrayKey.wrap(key));
799+
if (record != null) {
800+
return record;
801+
}
802+
if (hotRecordCache != null) {
803+
record = hotRecordCache.getIfPresent(buildHotCacheKey(key));
804+
if (record != null) {
805+
hotRecordCacheHitCount.incrementAndGet();
806+
}
807+
}
808+
return record;
809+
}
810+
811+
private ByteArrayKey buildHotCacheKey(byte[] key) {
812+
int partition = getPartition();
813+
byte[] compositeKey = new byte[4 + key.length];
814+
compositeKey[0] = (byte) (partition >>> 24);
815+
compositeKey[1] = (byte) (partition >>> 16);
816+
compositeKey[2] = (byte) (partition >>> 8);
817+
compositeKey[3] = (byte) partition;
818+
System.arraycopy(key, 0, compositeKey, 4, key.length);
819+
return ByteArrayKey.wrap(compositeKey);
820+
}
821+
822+
public long getHotRecordCacheHitCount() {
823+
return hotRecordCacheHitCount.get();
775824
}
776825

777826
/**

‎clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java‎

Lines changed: 35 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@
2222
import static java.util.concurrent.TimeUnit.MINUTES;
2323
import static java.util.concurrent.TimeUnit.SECONDS;
2424

25+
import com.github.benmanes.caffeine.cache.Cache;
26+
import com.github.benmanes.caffeine.cache.Caffeine;
2527
import com.linkedin.davinci.client.DaVinciRecordTransformer;
2628
import com.linkedin.davinci.client.DaVinciRecordTransformerConfig;
2729
import com.linkedin.davinci.client.DaVinciRecordTransformerRecordMetadata;
@@ -370,6 +372,10 @@ void setWriteComputeFailureCode(int code) {
370372

371373
private final boolean isActiveActiveReplicationEnabled;
372374

375+
/** Shared Caffeine cache for hot transient records across all partitions. Nullable when disabled. */
376+
private final Cache<ByteArrayKey, PartitionConsumptionState.TransientRecord> hotRecordCache;
377+
private final int minValueSizeForHotCache;
378+
373379
/**
374380
* This would be the number of partitions in the StorageEngine and in version topics
375381
*/
@@ -682,6 +688,22 @@ public StoreIngestionTask(
682688
this.localKafkaServer = this.kafkaProps.getProperty(KAFKA_BOOTSTRAP_SERVERS);
683689
this.localKafkaServerSingletonSet = Collections.singleton(localKafkaServer);
684690
this.isActiveActiveReplicationEnabled = version.isActiveActiveReplicationEnabled();
691+
692+
if (store.isTransientRecordCacheEnabled()
693+
&& (this.isActiveActiveReplicationEnabled || this.isWriteComputationEnabled)) {
694+
long maxWeight = serverConfig.getTransientRecordCacheMaxWeight();
695+
this.hotRecordCache = Caffeine.newBuilder()
696+
.maximumWeight(maxWeight)
697+
.weigher((ByteArrayKey k, PartitionConsumptionState.TransientRecord v) -> {
698+
return k.getContent().length + (v.getValue() != null ? v.getValueLen() : 0);
699+
})
700+
.build();
701+
this.minValueSizeForHotCache = serverConfig.getTransientRecordCacheMinValueSize();
702+
} else {
703+
this.hotRecordCache = null;
704+
this.minValueSizeForHotCache = 0;
705+
}
706+
685707
this.offsetLagDeltaRelaxEnabled = serverConfig.getOffsetLagDeltaRelaxFactorForFastOnlineTransitionInRestart() > 0;
686708
this.timeLagRelaxEnabled = serverConfig.getTimeLagThresholdForFastOnlineTransitionInRestartMinutes() > 0;
687709
this.metaStoreWriter = builder.getMetaStoreWriter();
@@ -2729,7 +2751,9 @@ private PartitionConsumptionState createAndInstallPartitionConsumptionState(
27292751
offsetRecord,
27302752
pubSubContext,
27312753
hybridStoreConfig.isPresent(),
2732-
schemaRepository.getKeySchema(storeName).getSchema());
2754+
schemaRepository.getKeySchema(storeName).getSchema(),
2755+
hotRecordCache,
2756+
minValueSizeForHotCache);
27332757
freshPcs.setCurrentVersionSupplier(isCurrentVersion);
27342758

27352759
boolean isFutureVersionReady = isFutureVersionReady(kafkaVersionTopic, storeRepository);
@@ -2771,7 +2795,9 @@ private PartitionConsumptionState createPlaceholderPartitionConsumptionState(
27712795
placeholderOffset,
27722796
pubSubContext,
27732797
hybridStoreConfig.isPresent(),
2774-
schemaRepository.getKeySchema(storeName).getSchema());
2798+
schemaRepository.getKeySchema(storeName).getSchema(),
2799+
hotRecordCache,
2800+
minValueSizeForHotCache);
27752801
pcs.setCurrentVersionSupplier(isCurrentVersion);
27762802

27772803
boolean isFutureVersionReady = isFutureVersionReady(kafkaVersionTopic, storeRepository);
@@ -3041,7 +3067,9 @@ private void resetOffset(int partition, PubSubTopicPartition topicPartition, boo
30413067
new OffsetRecord(partitionStateSerializer, pubSubContext),
30423068
pubSubContext,
30433069
hybridStoreConfig.isPresent(),
3044-
schemaRepository.getKeySchema(storeName).getSchema());
3070+
schemaRepository.getKeySchema(storeName).getSchema(),
3071+
hotRecordCache,
3072+
minValueSizeForHotCache);
30453073
consumptionState.setCurrentVersionSupplier(isCurrentVersion);
30463074
partitionConsumptionStateMap.put(partition, consumptionState);
30473075
storageUtilizationManager.initPartition(partition);
@@ -5342,6 +5370,10 @@ public boolean isActiveActiveReplicationEnabled() {
53425370
return this.isActiveActiveReplicationEnabled;
53435371
}
53445372

5373+
public long getHotRecordCacheEstimatedSize() {
5374+
return hotRecordCache != null ? hotRecordCache.estimatedSize() : 0;
5375+
}
5376+
53455377
/**
53465378
* Invoked by admin request to dump the requested partition consumption states
53475379
*/

‎clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/HostLevelIngestionStats.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -172,6 +172,8 @@ public class HostLevelIngestionStats extends AbstractVeniceStats {
172172
private final Sensor batchProcessingRequestLatencySensor;
173173
private final LongAdderRateGauge batchProcessingRequestErrorSensor;
174174

175+
private final Sensor hotRecordCacheHitCountSensor;
176+
175177
/**
176178
* @param totalStats the total stats singleton instance, or null if we are constructing the total stats
177179
*/
@@ -528,6 +530,12 @@ public HostLevelIngestionStats(
528530
totalStats,
529531
() -> totalStats.batchProcessingRequestLatencySensor,
530532
avgAndMax());
533+
534+
this.hotRecordCacheHitCountSensor = registerPerStoreAndTotalSensor(
535+
"hot_record_cache_hit_count",
536+
totalStats,
537+
() -> totalStats.hotRecordCacheHitCountSensor,
538+
new OccurrenceRate());
531539
}
532540

533541
private Measurable measurable(
@@ -699,6 +707,10 @@ public void recordWriteComputeCacheHitCount() {
699707
writeComputeCacheHitCount.record();
700708
}
701709

710+
public void recordHotRecordCacheHitCount() {
711+
hotRecordCacheHitCountSensor.record();
712+
}
713+
702714
public void recordWriteComputeLookupCount() {
703715
writeComputeLookupCount.record();
704716
}

0 commit comments

Comments
 (0)