[server][vpj][admin-tool] Plumb KME SchemaReader through consumer-side dictionary fetch and topic dump paths

[sushantmane]

Summary

Defense-in-depth follow-up to #2808 (DoL stamp vtp header fix). Plumbs a KME SchemaReader through every production callsite that opens its own PubSubMessageDeserializer, so a consumer that lands on a header-less control-message record can still resolve an unknown KME protocol version through a schema reader instead of throwing Received Protocol Version 'N' which is not supported. Useful insurance against any future writer regression that drops the vtp header.

Depends on #2808. Once that lands, this branch will be rebased to drop the writer change (which is the first commit on this branch).

Consumer-side plumbing

The following production paths now use a KafkaValueSerializer whose setSchemaReader is wired with a KmeSchemaReader, so when the in-memory protocol-version cache misses and vtp is unavailable, the serializer fetches the unknown schema through the reader:

Callsite	Schema source
KafkaInputUtils.getCompressor (ZSTD_WITH_DICT branch, fetches the SoP dictionary from the source VT)	newer.kme.schemas.* broadcast on the VPJ job conf
AbstractInputRecordProcessor.readDictionaryFromKafka	same broadcast
AbstractPartitionWriter ZSTD branch	same broadcast
SparkPubSubPartitionReaderFactory.createReader	same broadcast
VTConsistencyCheckerJob.createConsumer	same broadcast
AdminTool.getConsumer	KmeSchemaReader.fromControllerClient(ControllerClient) fetches all KME schemas from the venice_system_store_kafka_message_envelope system store
StoreIngestionTask.getNewStoreVersionState (SoP-null fallback path)	reuses the host's kafkaMessageEnvelopeSchemaReader, newly plumbed through StoreIngestionTaskFactory.Builder.setKafkaMessageEnvelopeSchemaReader

VeniceChangelogConsumerImpl already wires a KME SchemaReader (added in #2177) — no change. TopicMetadataFetcher constructs a deserializer but never invokes the deserialize path (only uses position/topic-metadata APIs on the underlying consumer) — left as-is.

Graceful fallback semantics

Where the plumbing exists but the upstream config wasn't actually broadcast, the helpers log a warning and fall back to the jar-only deserializer instead of throwing. The on-wire vtp header bootstrap still applies on the fallback path. Specific behaviors:

• KafkaInputUtils.buildSchemaAwareDeserializer(VeniceProperties): if SYSTEM_SCHEMA_READER_ENABLED is false (older VPJ job conf), logs a warning and returns PubSubMessageDeserializer.createDefaultDeserializer().
• AdminTool.buildAdminToolDeserializer(): if controllerClient is null (e.g., admin-message dump that runs before cluster discovery), logs and returns PubSubMessageDeserializer.createOptimizedDeserializer(). If the controller fetch itself fails (network, ACL, etc.), the exception is caught and the same fallback is used with a warn log.
• StoreIngestionTask.getNewStoreVersionState SoP-null branch: if the host's kafkaMessageEnvelopeSchemaReader wasn't plumbed through StoreIngestionTaskFactory.Builder, logs and uses createDefaultDeserializer().

PubSubMessageDeserializer.createWithSchemaReader and createOptimizedWithSchemaReader themselves do Objects.requireNonNull(schemaReader, …) — that's an API contract.

Refactor: KmeSchemaReader moved to venice-common

KMESchemaReaderForKafkaInputFormat lived under clients/venice-push-job/.../hadoop/input/kafka. The class was a general SchemaReader over KME protocol-version schemas — its name was an artifact of where it lived, not what it did. Moved to internal/venice-common/.../schema/KmeSchemaReader so non-VPJ callers (admin tool, server, controller-side code) can construct one without a build dependency on venice-push-job. Renamed to KmeSchemaReader to follow Java naming conventions for acronyms.

A new KmeSchemaReader.fromControllerClient(ControllerClient) static factory pulls all KME schemas from the venice_system_store_kafka_message_envelope system store via controllerClient.getAllValueSchema(…) and merges them with the schemas baked into the jar's resources. Used by AdminTool here.

Strict-to-graceful iteration

The first iteration of this PR used strict throws (IllegalStateException) at the consumer-side fallback points to surface every callsite that hadn't plumbed a SchemaReader. CI confirmed the gaps exactly where expected: integration-test VPJ jobs don't set SYSTEM_SCHEMA_READER_ENABLED=true by default, and various admin-tool entry points run before cluster discovery has wired controllerClient. The follow-up commits soften those throws to logged fallbacks. The Objects.requireNonNull on createWithSchemaReader / createOptimizedWithSchemaReader remains — that's an API contract, not a missing-plumbing concern.

Related

• [server] Attach vtp protocol-schema header on DoL stamp segment-start record #2808 — the minimal DoL stamp vtp header fix this PR builds on. Merge that first.
• [protocol] Add KafkaMessageEnvelope schema v14 to carry upstream record timestamp #2778 added KME schema v14 to the protocol.
• [compat][server] Activate KME v14 and stamp upstreamMessageTimestamp on leader produces #2780 activated KME v14 on the writer.
• [server][dvc] Stop pinning ~16KB Avro schema header per queued PubSub message #2762 added vtp header stripping after successful deserialize.
• [cdc] Add KME schema reader in CDC to support schema evolution #2177 wired the KME SchemaReader into the CDC changelog consumer.

Testing Done

New unit tests (all pass locally):

• KmeSchemaReaderTest (new file, 8 tests): exercises the constructor merge with empty and non-empty newer-schemas maps; getKeySchema, getValueSchema(id), getLatestValueSchema(Id); the update-schema methods that should throw VeniceUnsupportedOperationException; close is a no-op; fromControllerClient happy path with a mocked ControllerClient returning a populated MultiSchemaResponse; and fromControllerClient error path where the response carries an error string.
• PubSubMessageDeserializerTest (+4 tests): createWithSchemaReader(null) and createOptimizedWithSchemaReader(null) both throw NullPointerException via requireNonNull; both successfully construct a deserializer with a non-null mocked SchemaReader and expose a KafkaValueSerializer.
• KafkaInputUtilsTest (+2 tests): buildSchemaAwareDeserializer falls back to a jar-only deserializer when SYSTEM_SCHEMA_READER_ENABLED is unset; builds a SchemaReader-backed deserializer when the system flag is enabled and newer.kme.schemas.* entries are present in the job conf.
• StoreIngestionTaskTest (+2 tests): buildDictionaryFetchDeserializer returns a SchemaReader-backed deserializer when the host plumbed kafkaMessageEnvelopeSchemaReader; returns the jar-only fallback (with a warn log) when the SchemaReader wasn't wired.

Existing tests confirmed still passing locally: all VeniceWriterUnitTest, KafkaInputUtilsTest, StoreIngestionTaskTest.testGetNewStoreVersionState*, PubSubMessageDeserializerTest.

Compile-clean across internal:venice-common, clients:venice-push-job, clients:venice-admin-tool, clients:da-vinci-client, services:venice-server (main + test sources).

[Copilot]

Pull request overview

This PR fixes a forward-compatibility bootstrap gap where VeniceWriter.sendDoLStamp() could emit the first record on a fresh version topic with empty headers (missing vtp), and extends “defense-in-depth” by plumbing a KME-backed SchemaReader into additional consumer/deserializer construction paths so unknown KME protocol versions can be resolved even when vtp is absent.

Changes:

• Producer: route DoL-stamp produce through getHeaders(...) so segment-start vtp is attached.
• Consumers/tools: introduce KmeSchemaReader in venice-common and add PubSubMessageDeserializer factories that accept a SchemaReader; update several VPJ/Spark/admin-tool/server ingestion call sites to use schema-aware deserializers.
• Tests: add a regression unit test asserting DoL stamps carry vtp.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
internal/venice-common/src/test/java/com/linkedin/venice/writer/VeniceWriterUnitTest.java	Adds regression test ensuring DoL stamp carries vtp header.
internal/venice-common/src/main/java/com/linkedin/venice/writer/VeniceWriter.java	Fixes DoL-stamp send path to attach headers via getHeaders(...).
internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java	Moves/renames KME schema reader into venice-common and adds a controller-backed factory.
internal/venice-common/src/main/java/com/linkedin/venice/pubsub/api/PubSubMessageDeserializer.java	Adds schema-reader-backed factory methods for production-safe deserialization.
clients/venice-push-job/src/main/java/com/linkedin/venice/spark/input/pubsub/SparkPubSubPartitionReaderFactory.java	Switches Spark PubSub consumption to a schema-aware deserializer.
clients/venice-push-job/src/main/java/com/linkedin/venice/spark/consistency/VTConsistencyCheckerJob.java	Switches VT consistency checker consumer to a schema-aware deserializer.
clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/task/datawriter/AbstractPartitionWriter.java	Wires schema-aware deserializer into ZSTD dictionary fetch path.
clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/task/datawriter/AbstractInputRecordProcessor.java	Wires schema-aware deserializer into dictionary fetch helper.
clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/input/kafka/KafkaInputUtils.java	Renames schema reader usage and adds buildSchemaAwareDeserializer(...) helper.
clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java	Switches admin-tool topic consumer to a controller-backed KME schema-aware deserializer.
clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskFactory.java	Plumbs a KME SchemaReader through the SIT factory builder.
clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java	Uses the plumbed schema reader for SoP-null dictionary fallback; adds strict null checks.
clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java	Passes the host’s KME schema reader into SIT builder.

Comments suppressed due to low confidence (3)

internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:56

• KmeSchemaReader(Map) claims to throw IllegalArgumentException when newerKmeSchemas is null, but it will currently NPE on newerKmeSchemas.entrySet(). Add an explicit null check (e.g., Objects.requireNonNull) so callers get a deterministic failure with a clear message.
  internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:54
• In this constructor, newerKmeSchemas are added into SchemaData before the resource-loaded schemas (see the subsequent addValueSchema loops). Since SchemaData.addValueSchema overwrites by id, any overlapping ids will end up using the resource schema, which contradicts the “merged on top”/runtime-precedence intent. Consider loading resource schemas first and then overlaying newerKmeSchemas (or explicitly preventing overwrites).
  internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:77
• fromControllerClient() assumes controllerClient is non-null and will NPE otherwise. Since this is a public factory that will be used across modules, add an explicit null check with a clear message (or document the precondition in the Javadoc).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 4 comments.

Comments suppressed due to low confidence (3)

internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:61

• In KmeSchemaReader(...), the merge order is reversed vs the class Javadoc: SchemaData.addValueSchema overwrites by ID, so adding newerKmeSchemas first and resource schemas second means jar resources win on duplicate IDs. If the intent is "newer schemas override jar-bundled schemas", add resource schemas first, then overlay newerKmeSchemas (or explicitly handle duplicates).
  internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:56
• The constructor Javadoc says it throws IllegalArgumentException when newerKmeSchemas is null, but the implementation doesn't validate it and will NPE on entrySet(). Consider Objects.requireNonNull(newerKmeSchemas, ...) to match the contract and fail fast with a clear message.
  internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:85
• fromControllerClient iterates response.getSchemas() without a null check. MultiSchemaResponse.schemas can be null if the controller returns a success response but omits the field, which would NPE here. Consider treating null as empty and/or throwing a VeniceException with a clear message.

[sushantmane]

Split into two PRs. The minimal DoL stamp vtp header fix is now #2808 — please review and merge that first. This PR (now scoped to the consumer-side KmeSchemaReader plumbing) will be rebased to drop the writer change once #2808 lands.

[Copilot]

Pull request overview

Copilot reviewed 17 out of 17 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:86

• fromControllerClient iterates response.getSchemas() without a null check. MultiSchemaResponse.schemas can be null (it’s nullable in the model and some callers explicitly guard against null), which would trigger an NPE here. Consider treating null as an empty schema list (jar-only reader) or throwing a clear VeniceException indicating a malformed controller response.
  internal/venice-common/src/main/java/com/linkedin/venice/schema/KmeSchemaReader.java:61
• KmeSchemaReader(Map) doesn’t enforce the documented null contract: passing null will currently NPE at newerKmeSchemas.entrySet() even though the Javadoc says IllegalArgumentException. Also, the merge order is opposite of the stated behavior (“newer schemas … merged on top of jar resources”): jar resource schemas are added after the runtime map and will overwrite any duplicate IDs from newerKmeSchemas. Add an explicit null guard (e.g., requireNonNull) and consider loading jar resources first, then applying runtime schemas so the precedence matches the documentation/intent.