Fix spotbug

Author: Minh Nguyen

internal/venice-common/src/main/java/com/linkedin/venice/serialization/avro/AvroProtocolDefinition.java

@@ -235,7 +235,7 @@ private static Set<Byte> validateMagicBytes() {
    */
   final String className;
 
-  Schema schema;
+  final Schema schema;
 
   /**
    * Indicates whether the protocol version should prepend the payload, or whether
@@ -290,7 +290,6 @@ private static Set<Byte> validateMagicBytes() {
   /**
    * Constructor for protocols where there is no SpecificRecord class, because the
    * schema defines only a primitive type, such as byte[].
-   * SpecificRecord class can be set through setClass() method later.
    */
   AvroProtocolDefinition(int currentProtocolVersion, Schema schema, String name) {
     this.magicByte = Optional.empty();
@@ -300,25 +299,11 @@ private static Set<Byte> validateMagicBytes() {
     this.schema = schema;
   }
 
-  /**
-   * Set the SpecificRecord class for this protocol.
-   * This method should be used only when the constructor without SpecificRecord class is used.
-   *
-   */
-  public void setClass(Class<? extends SpecificRecord> specificRecordClass) {
-    if (!this.className.equals(specificRecordClass.getSimpleName())) {
-      throw new VeniceException(
-          "SpecificRecord class is already set to " + this.className + ", cannot set it to "
-              + specificRecordClass.getSimpleName());
-    }
-    this.schema = SpecificData.get().getSchema(specificRecordClass);
-  }
-
   public <T extends SpecificRecord> InternalAvroSpecificSerializer<T> getSerializer() {
     if (magicByte.isPresent() || protocolVersionStoredInHeader) {
       return new InternalAvroSpecificSerializer<>(this);
     }
-    return new InternalAvroSpecificSerializer<>(this, 0);
+    return new InternalAvroSpecificSerializer<>(this, 0, this.getCurrentProtocolVersionSchema());
   }
 
   public int getCurrentProtocolVersion() {

internal/venice-common/src/main/java/com/linkedin/venice/serialization/avro/InternalAvroSpecificSerializer.java

@@ -104,18 +104,26 @@ private static class ReusableObjects {
 
   private final BiConsumer<Integer, Schema> newSchemaEncountered;
 
+  protected InternalAvroSpecificSerializer(
+      AvroProtocolDefinition protocolDef,
+      Integer payloadOffsetOverride,
+      Schema compiledProtocolSchema) {
+    this(protocolDef, payloadOffsetOverride, (schemaId, schema) -> {}, compiledProtocolSchema);
+  }
+
   protected InternalAvroSpecificSerializer(AvroProtocolDefinition protocolDef) {
     this(protocolDef, null);
   }
 
   protected InternalAvroSpecificSerializer(AvroProtocolDefinition protocolDef, Integer payloadOffsetOverride) {
-    this(protocolDef, payloadOffsetOverride, (schemaId, schema) -> {});
+    this(protocolDef, payloadOffsetOverride, (schemaId, schema) -> {}, protocolDef.getCurrentProtocolVersionSchema());
   }
 
   protected InternalAvroSpecificSerializer(
       AvroProtocolDefinition protocolDef,
       Integer payloadOffsetOverride,
-      BiConsumer<Integer, Schema> newSchemaEncountered) {
+      BiConsumer<Integer, Schema> newSchemaEncountered,
+      Schema compiledProtocolSchema) {
     // Magic byte handling
     if (protocolDef.getMagicByte().isPresent()) {
       this.magicByte = protocolDef.getMagicByte().get();
@@ -125,6 +133,8 @@ protected InternalAvroSpecificSerializer(
       this.MAGIC_BYTE_LENGTH = 0;
     }
 
+    this.compiledProtocol = compiledProtocolSchema;
+
     // Protocol version handling
     this.PROTOCOL_VERSION_OFFSET = MAGIC_BYTE_OFFSET + MAGIC_BYTE_LENGTH;
     if (protocolDef.protocolVersionStoredInHeader) {
@@ -156,9 +166,8 @@ protected InternalAvroSpecificSerializer(
       }
       this.PAYLOAD_OFFSET = payloadOffsetOverride;
     }
-    this.compiledProtocol = protocolDef.getCurrentProtocolVersionSchema();
 
-    Map<Integer, Schema> protocolSchemaMap = Utils.getAllSchemasFromResources(protocolDef);
+    Map<Integer, Schema> protocolSchemaMap = Utils.getAllSchemasFromResources(protocolDef, compiledProtocolSchema);
     int minimumSchemaId = protocolSchemaMap.keySet()
         .stream()
         .min(Integer::compareTo)

internal/venice-common/src/main/java/com/linkedin/venice/serialization/avro/KafkaValueSerializer.java

@@ -20,6 +20,6 @@ public KafkaValueSerializer() {
   }
 
   public KafkaValueSerializer(BiConsumer<Integer, Schema> newSchemaEncountered) {
-    super(AvroProtocolDefinition.KAFKA_MESSAGE_ENVELOPE, null, newSchemaEncountered);
+    super(AvroProtocolDefinition.KAFKA_MESSAGE_ENVELOPE, null, newSchemaEncountered, null);
   }
 }

internal/venice-common/src/main/java/com/linkedin/venice/utils/Utils.java

@@ -442,6 +442,12 @@ public static Schema getSchemaFromResource(String resourcePath) throws IOExcepti
   }
 
   public static Map<Integer, Schema> getAllSchemasFromResources(AvroProtocolDefinition protocolDef) {
+    return getAllSchemasFromResources(protocolDef, protocolDef.getCurrentProtocolVersionSchema());
+  }
+
+  public static Map<Integer, Schema> getAllSchemasFromResources(
+      AvroProtocolDefinition protocolDef,
+      Schema compiledSchema) {
     final int SENTINEL_PROTOCOL_VERSION_USED_FOR_UNDETECTABLE_COMPILED_SCHEMA =
         InternalAvroSpecificSerializer.SENTINEL_PROTOCOL_VERSION_USED_FOR_UNDETECTABLE_COMPILED_SCHEMA;
     final int SENTINEL_PROTOCOL_VERSION_USED_FOR_UNVERSIONED_PROTOCOL =
@@ -480,7 +486,7 @@ public static Map<Integer, Schema> getAllSchemasFromResources(AvroProtocolDefini
       try {
         Schema schema = Utils.getSchemaFromResource(versionPath);
         protocolSchemaMap.put(version, schema);
-        if (schema.equals(protocolDef.getCurrentProtocolVersionSchema())) {
+        if (schema.equals(compiledSchema)) {
           compiledProtocolVersion = (byte) version;
           break;
         }
@@ -523,7 +529,7 @@ public static Map<Integer, Schema> getAllSchemasFromResources(AvroProtocolDefini
     if (intendedCurrentProtocol == null) {
       throw new VeniceException(
           "Failed to get schema for current version: " + currentProtocolVersion + " class: " + className);
-    } else if (!intendedCurrentProtocol.equals(protocolDef.getCurrentProtocolVersionSchema())) {
+    } else if (!intendedCurrentProtocol.equals(compiledSchema)) {
       throw new VeniceException(
           "The intended protocol version (" + currentProtocolVersion
               + ") does not match the compiled protocol version (" + compiledProtocolVersion + ").");

internal/venice-common/src/test/java/com/linkedin/venice/serialization/avro/AvroProtocolDefinitionTest.java

@@ -1,12 +1,8 @@
 package com.linkedin.venice.serialization.avro;
 
-import static com.linkedin.venice.serialization.avro.AvroProtocolDefinition.ADMIN_OPERATION;
 import static com.linkedin.venice.serialization.avro.AvroProtocolDefinition.KAFKA_MESSAGE_ENVELOPE;
 import static com.linkedin.venice.serialization.avro.AvroProtocolDefinition.SERVER_ADMIN_RESPONSE;
 
-import com.linkedin.venice.exceptions.VeniceException;
-import com.linkedin.venice.kafka.protocol.KafkaMessageEnvelope;
-import org.apache.avro.specific.SpecificRecord;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -17,21 +13,4 @@ public void testGetSerializer() {
     Assert.assertNotNull(KAFKA_MESSAGE_ENVELOPE.getSerializer());
     Assert.assertNotNull(SERVER_ADMIN_RESPONSE.getSerializer());
   }
-
-  @Test
-  public void testSetClass() {
-    // No error should be thrown when setting the correct class
-    Class<? extends SpecificRecord> KafkaMessageEnvelopeClass = KafkaMessageEnvelope.class;
-    KAFKA_MESSAGE_ENVELOPE.setClass(KafkaMessageEnvelopeClass);
-
-    // An error should be thrown when setting an incorrect class
-    try {
-      ADMIN_OPERATION.setClass(KafkaMessageEnvelopeClass);
-      Assert.fail("Expected an IllegalArgumentException to be thrown");
-    } catch (VeniceException e) {
-      Assert.assertEquals(
-          e.getMessage(),
-          "SpecificRecord class is already set to AdminOperation, cannot set it to KafkaMessageEnvelope");
-    }
-  }
 }

internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionWithProtocolRollbackIntegrationTest.java

@@ -1,7 +1,10 @@
 package com.linkedin.venice.controller.kafka.consumer;
 
 import static com.linkedin.venice.ConfigKeys.*;
+import static com.linkedin.venice.utils.AvroSchemaUtils.*;
 
+import com.linkedin.avroutil1.compatibility.AvroCompatibilityHelper;
+import com.linkedin.avroutil1.compatibility.FieldBuilder;
 import com.linkedin.venice.controller.VeniceParentHelixAdmin;
 import com.linkedin.venice.controller.kafka.AdminTopicUtils;
 import com.linkedin.venice.controller.kafka.protocol.admin.AdminOperation;
@@ -35,7 +38,6 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
-import java.util.Map;
 import java.util.Properties;
 import java.util.concurrent.TimeUnit;
 import org.apache.avro.Schema;
@@ -61,15 +63,14 @@
 public class AdminConsumptionWithProtocolRollbackIntegrationTest {
   private static final int TIMEOUT = 1 * Time.MS_PER_MINUTE;
 
-  private final AdminOperationSerializer adminOperationSerializer = new AdminOperationSerializer();
+  private AdminOperationSerializer adminOperationSerializer;
 
   private static final String owner = "test_owner";
   private static final String keySchema = "\"string\"";
   private static final String valueSchema = "\"string\"";
   private static final int adminConsumptionMaxWorkerPoolSize = 3;
 
   private VeniceTwoLayerMultiRegionMultiClusterWrapper venice;
-  private AdminConsumerService adminConsumerService;
   private ControllerClient parentControllerClient;
   private VeniceParentHelixAdmin parentHelixAdmin;
   private VeniceWriter<byte[], byte[], byte[]> writer;
@@ -102,7 +103,6 @@ public void setUp() {
     VeniceControllerWrapper parentController = venice.getParentControllers().get(0);
     parentControllerClient = new ControllerClient(venice.getClusterNames()[0], parentController.getControllerUrl());
     clusterName = venice.getClusterNames()[0];
-    adminConsumerService = parentController.getAdminConsumerServiceByCluster(clusterName);
     parentHelixAdmin = (VeniceParentHelixAdmin) parentController.getVeniceAdmin();
     PubSubTopicRepository pubSubTopicRepository = parentHelixAdmin.getPubSubTopicRepository();
     TopicManager topicManager = parentHelixAdmin.getTopicManager();
@@ -114,6 +114,7 @@ public void setUp() {
         pubSubBrokerWrapper.getPubSubClientsFactory().getProducerAdapterFactory();
     writer = IntegrationTestPushUtils.getVeniceWriterFactory(pubSubBrokerWrapper, pubSubProducerAdapterFactory)
         .createVeniceWriter(new VeniceWriterOptions.Builder(adminTopic.getName()).build());
+    adminOperationSerializer = parentHelixAdmin.getAdminOperationSerializer();
   }
 
   @AfterClass
@@ -193,51 +194,48 @@ private void addSchemaToSchemaSystemStore(int schemaId, Schema schema) {
    * and register it to schema system store with new schema id
    */
   private int createNewSchemaAndRegister(int originalSchemaId) {
-    Schema schema = AdminOperationSerializer.getSchema(originalSchemaId);
+    Schema schema = adminOperationSerializer.getSchema(originalSchemaId);
 
     int newSchemaId = originalSchemaId + 1;
     Schema newSchema = createNewSchema(schema);
     // We use the same schema and by-pass compatibility check for testing purpose
     addSchemaToSchemaSystemStore(newSchemaId, newSchema);
 
-    AdminOperationSerializer adminOperationSerializer = parentHelixAdmin.getAdminOperationSerializer();
     adminOperationSerializer.addSchema(newSchemaId, newSchema);
     return newSchemaId;
   }
 
   /**
    * Create a new schema by adding a new field to the original schema
    */
-  private Schema createNewSchema(Schema originalSchema) {
-    // Convert schema to string
-    Schema newSchema = Schema.createRecord(
+  private Schema createNewSchema(final Schema originalSchema) {
+    if (originalSchema == null) {
+      throw new IllegalArgumentException("Original schema cannot be null");
+    }
+
+    // Create a new record schema with defensive copies
+    final Schema newSchema = Schema.createRecord(
         originalSchema.getName(),
         originalSchema.getDoc(),
         originalSchema.getNamespace(),
         originalSchema.isError());
 
-    // Create a deep copy of all fields
-    List<Schema.Field> newFields = new ArrayList<>();
+    // Deep copy fields
+    final List<Schema.Field> newFields = new ArrayList<>();
     for (Schema.Field field: originalSchema.getFields()) {
-      Schema.Field newField = new Schema.Field(
-          field.name(),
-          field.schema(), // This creates a new field with the same schema
-          field.doc(),
-          field.defaultVal());
-      // Copy all properties
-      for (Map.Entry<String, Object> prop: field.getObjectProps().entrySet()) {
-        newField.addProp(prop.getKey(), prop.getValue());
-      }
-      newFields.add(newField);
+      // Defensive copy of schema (Avro schemas are immutable, so this is safe)
+      FieldBuilder newField = deepCopySchemaFieldWithoutFieldProps(field);
+      copyFieldProperties(newField, field);
+      newFields.add(newField.build());
     }
 
-    // Add the new field
-    Schema.Field testField =
-        new Schema.Field("testField", Schema.create(Schema.Type.INT), "Documentation for testField", 0);
+    // Add extra field
+    final Schema.Field testField = AvroCompatibilityHelper
+        .createSchemaField("testField", Schema.create(Schema.Type.INT), "Documentation for testField", 0);
     newFields.add(testField);
 
-    // Set all fields at once
-    newSchema.setFields(newFields);
+    // Set fields using defensive copy
+    newSchema.setFields(Collections.unmodifiableList(new ArrayList<>(newFields)));
 
     return newSchema;
   }
@@ -302,4 +300,26 @@ private byte[] getStoreUpdateMessage(
   private int nextExecutionId() {
     return executionId++;
   }
+
+  private static FieldBuilder deepCopySchemaFieldWithoutFieldProps(Schema.Field field) {
+    FieldBuilder fieldBuilder = AvroCompatibilityHelper.newField(null)
+        .setName(field.name())
+        .setSchema(field.schema())
+        .setDoc(field.doc())
+        .setOrder(field.order());
+    // set default as AvroCompatibilityHelper builder might drop defaults if there is type mismatch
+    if (field.hasDefaultValue()) {
+      fieldBuilder.setDefault(getFieldDefault(field));
+    }
+    return fieldBuilder;
+  }
+
+  private static void copyFieldProperties(FieldBuilder fieldBuilder, Schema.Field field) {
+    AvroCompatibilityHelper.getAllPropNames(field).forEach(k -> {
+      String propValue = AvroCompatibilityHelper.getFieldPropAsJsonString(field, k);
+      if (propValue != null) {
+        fieldBuilder.addProp(k, propValue);
+      }
+    });
+  }
 }

services/venice-controller/src/main/java/com/linkedin/venice/controller/init/SystemSchemaInitializationRoutine.java

@@ -20,6 +20,7 @@
 import java.util.Map;
 import java.util.Optional;
 import org.apache.avro.Schema;
+import org.apache.avro.specific.SpecificData;
 import org.apache.avro.specific.SpecificRecord;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -30,6 +31,7 @@ public class SystemSchemaInitializationRoutine implements ClusterLeaderInitializ
   private static final String DEFAULT_KEY_SCHEMA_STR = "\"int\"";
 
   private final AvroProtocolDefinition protocolDefinition;
+  private final Schema protocolSchema;
   private final VeniceControllerMultiClusterConfig multiClusterConfigs;
   private final VeniceHelixAdmin admin;
   private final Optional<Schema> keySchema;
@@ -69,16 +71,14 @@ public SystemSchemaInitializationRoutine(
       boolean autoRegisterDerivedComputeSchema,
       Class<? extends SpecificRecord> specificRecordClass) {
     this.protocolDefinition = protocolDefinition;
+    this.protocolSchema = specificRecordClass != null
+        ? SpecificData.get().getSchema(specificRecordClass)
+        : protocolDefinition.getCurrentProtocolVersionSchema();
     this.multiClusterConfigs = multiClusterConfigs;
     this.admin = admin;
     this.keySchema = keySchema;
     this.storeMetadataUpdate = storeMetadataUpdate;
     this.autoRegisterDerivedComputeSchema = autoRegisterDerivedComputeSchema;
-    // Set the specific record class if provided. This allows us to override class and schema for protocol definitions
-    // that is in venice-controller and cannot reference in venice-common.
-    if (specificRecordClass != null) {
-      protocolDefinition.setClass(specificRecordClass);
-    }
   }
 
   /**
@@ -89,7 +89,7 @@ public void execute(String clusterToInit) {
     String intendedCluster = multiClusterConfigs.getSystemSchemaClusterName();
     if (intendedCluster.equals(clusterToInit)) {
       String systemStoreName = protocolDefinition.getSystemStoreName();
-      Map<Integer, Schema> protocolSchemaMap = Utils.getAllSchemasFromResources(protocolDefinition);
+      Map<Integer, Schema> protocolSchemaMap = Utils.getAllSchemasFromResources(protocolDefinition, protocolSchema);
 
       // Sanity check to make sure the store is not already created in another cluster.
       try {