IPROTO-81 SerializationContext.registerMarshaller should unregister t…

…he previous Java type (if any)

core/src/main/java/org/infinispan/protostream/impl/SerializationContextImpl.java

@@ -182,18 +182,21 @@ public void registerMarshaller(BaseMarshaller<?> marshaller) {
          throw new IllegalArgumentException("marshaller argument cannot be null");
       }
       BaseMarshallerDelegate<?> marshallerDelegate = makeMarshallerDelegate(marshaller);
-      marshallersByName.put(marshaller.getTypeName(), marshallerDelegate);
+      BaseMarshallerDelegate<?> existingDelegate = marshallersByName.put(marshaller.getTypeName(), marshallerDelegate);
+      if (existingDelegate != null) {
+         marshallersByClass.remove(existingDelegate.getMarshaller().getJavaClass());
+      }
       marshallersByClass.put(marshaller.getJavaClass(), marshallerDelegate);
    }
 
    private BaseMarshallerDelegate<?> makeMarshallerDelegate(BaseMarshaller<?> marshaller) {
       if (marshaller.getJavaClass().isEnum() && !(marshaller instanceof EnumMarshaller)) {
-         throw new IllegalArgumentException("Invalid marshaller (the produced class is a Java Enum but the marshaller is not an EnumMarshaller) : " + marshaller);
+         throw new IllegalArgumentException("Invalid marshaller (the produced class is a Java Enum but the marshaller is not an EnumMarshaller) : " + marshaller.getClass().getName());
       }
       // we try to validate first that a message descriptor exists
       if (marshaller instanceof EnumMarshaller) {
          if (!marshaller.getJavaClass().isEnum()) {
-            throw new IllegalArgumentException("Invalid enum marshaller (the produced class is not a Java Enum) : " + marshaller);
+            throw new IllegalArgumentException("Invalid enum marshaller (the produced class is not a Java Enum) : " + marshaller.getClass().getName());
          }
          EnumDescriptor enumDescriptor = getEnumDescriptor(marshaller.getTypeName());
          return new EnumMarshallerDelegate<>((EnumMarshaller<?>) marshaller, enumDescriptor);
@@ -211,7 +214,7 @@ public void unregisterMarshaller(BaseMarshaller<?> marshaller) {
       }
       BaseMarshallerDelegate<?> marshallerDelegate = marshallersByName.get(marshaller.getTypeName());
       if (marshallerDelegate == null || marshallerDelegate.getMarshaller() != marshaller) {
-         throw new IllegalArgumentException("The given marshaller was not previously registered");
+         throw new IllegalArgumentException("The given marshaller was not previously registered with this SerializationContext");
       }
       marshallersByName.remove(marshaller.getTypeName());
       marshallersByClass.remove(marshaller.getJavaClass());
@@ -258,7 +261,7 @@ public <T> BaseMarshallerDelegate<T> getMarshallerDelegate(String descriptorFull
       if (marshallerDelegate == null) {
          BaseMarshaller<?> marshaller = getMarshallerFromProvider(descriptorFullName);
          if (marshaller == null) {
-            throw new IllegalArgumentException("No marshaller registered for " + descriptorFullName);
+            throw new IllegalArgumentException("No marshaller registered for Protobuf type " + descriptorFullName);
          }
          marshallerDelegate = (BaseMarshallerDelegate<T>) makeMarshallerDelegate(marshaller);
       }
@@ -270,7 +273,7 @@ public <T> BaseMarshallerDelegate<T> getMarshallerDelegate(Class<T> javaClass) {
       if (marshallerDelegate == null) {
          BaseMarshaller<?> marshaller = getMarshallerFromProvider(javaClass);
          if (marshaller == null) {
-            throw new IllegalArgumentException("No marshaller registered for " + javaClass.getName());
+            throw new IllegalArgumentException("No marshaller registered for Java type " + javaClass.getName());
          }
          marshallerDelegate = (BaseMarshallerDelegate<T>) makeMarshallerDelegate(marshaller);
       }

core/src/test/java/org/infinispan/protostream/annotations/impl/ProtoSchemaBuilderTest.java

@@ -3,9 +3,11 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 
 import java.io.File;
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.StandardCopyOption;
 import java.time.Instant;
@@ -17,10 +19,12 @@
 
 import org.infinispan.protostream.DescriptorParserException;
 import org.infinispan.protostream.FileDescriptorSource;
+import org.infinispan.protostream.MessageMarshaller;
 import org.infinispan.protostream.ProtobufUtil;
 import org.infinispan.protostream.SerializationContext;
 import org.infinispan.protostream.annotations.ProtoEnumValue;
 import org.infinispan.protostream.annotations.ProtoField;
+import org.infinispan.protostream.annotations.ProtoMessage;
 import org.infinispan.protostream.annotations.ProtoSchemaBuilder;
 import org.infinispan.protostream.annotations.ProtoSchemaBuilderException;
 import org.infinispan.protostream.annotations.impl.testdomain.Simple;
@@ -108,7 +112,7 @@ static abstract class AbstractType {
     * Abstract field types in a message class are not accepted.
     */
    @Test
-   public void tesAbstractClass() throws Exception {
+   public void testAbstractClass() throws Exception {
       exception.expect(ProtoSchemaBuilderException.class);
       exception.expectMessage("The type org.infinispan.protostream.annotations.impl.ProtoSchemaBuilderTest$MessageWithAbstractFieldType$AbstractType of field 'testField1' of class org.infinispan.protostream.annotations.impl.ProtoSchemaBuilderTest$MessageWithAbstractFieldType should not be abstract.");
 
@@ -379,6 +383,76 @@ public void testDuplicateEnumValueName() throws Exception {
             .build(ctx);
    }
 
+   @ProtoMessage(name = "User")
+   static class AnotherUser {
+
+      @ProtoField(number = 1, defaultValue = "1")
+      public byte gender;
+   }
+
+   @Test
+   public void testReplaceExistingMarshallerWithAnnotations() throws Exception {
+      SerializationContext ctx = createContext();
+
+      assertTrue(ctx.canMarshall("sample_bank_account.User"));
+      assertTrue(ctx.canMarshall(org.infinispan.protostream.domain.User.class));
+
+      // replace 'sample_bank_account.User' with a new definition and also generate and register a new marshaller for it
+      ProtoSchemaBuilder protoSchemaBuilder = new ProtoSchemaBuilder();
+      protoSchemaBuilder
+            .fileName("sample_bank_account/bank.proto")
+            .packageName("sample_bank_account")
+            .addClass(AnotherUser.class)
+            .build(ctx);
+
+      assertTrue(ctx.canMarshall("sample_bank_account.User"));
+      assertTrue(ctx.canMarshall(AnotherUser.class));
+      // this 'sample_bank_account.User' definition does not have a 'name' field
+      assertNull(ctx.getMessageDescriptor("sample_bank_account.User").findFieldByName("name"));
+
+      // the old Java type it was mapping to is no longer marshallable
+      assertFalse(ctx.canMarshall(org.infinispan.protostream.domain.User.class));
+   }
+
+   @Test
+   public void testReplaceExistingMarshaller() throws Exception {
+      SerializationContext ctx = createContext();
+
+      assertTrue(ctx.canMarshall("sample_bank_account.User"));
+      assertTrue(ctx.canMarshall(org.infinispan.protostream.domain.User.class));
+
+      MessageMarshaller<AnotherUser> anotherUserMarshaller = new MessageMarshaller<AnotherUser>() {
+
+         @Override
+         public AnotherUser readFrom(ProtoStreamReader reader) throws IOException {
+            int gender = reader.readInt("gender");
+            AnotherUser anotherUser = new AnotherUser();
+            anotherUser.gender = (byte) gender;
+            return anotherUser;
+         }
+
+         @Override
+         public void writeTo(ProtoStreamWriter writer, AnotherUser user) throws IOException {
+            writer.writeInt("gender", user.gender);
+         }
+
+         @Override
+         public String getTypeName() {
+            return "sample_bank_account.User";
+         }
+
+         @Override
+         public Class<AnotherUser> getJavaClass() {
+            return AnotherUser.class;
+         }
+      };
+      ctx.registerMarshaller(anotherUserMarshaller);
+
+      assertTrue(ctx.canMarshall("sample_bank_account.User"));
+      assertTrue(ctx.canMarshall(AnotherUser.class));
+      assertFalse(ctx.canMarshall(org.infinispan.protostream.domain.User.class));
+   }
+
    static class MessageWithAllFieldTypes {
 
       @ProtoField(number = 1, defaultValue = "A")