Support complex type key in MessagePackGenerator

Author: komamitsu

msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessagePackFactory.java

@@ -37,6 +37,7 @@ public class MessagePackFactory
     private static final long serialVersionUID = 2578263992015504347L;
 
     private final MessagePack.PackerConfig packerConfig;
+    private boolean reuseResourceInGenerator = true;
 
     public MessagePackFactory()
     {
@@ -48,11 +49,16 @@ public MessagePackFactory(MessagePack.PackerConfig packerConfig)
         this.packerConfig = packerConfig;
     }
 
+    public void setReuseResourceInGenerator(boolean reuseResourceInGenerator)
+    {
+        this.reuseResourceInGenerator = reuseResourceInGenerator;
+    }
+
     @Override
     public JsonGenerator createGenerator(OutputStream out, JsonEncoding enc)
             throws IOException
     {
-        return new MessagePackGenerator(_generatorFeatures, _objectCodec, out, packerConfig);
+        return new MessagePackGenerator(_generatorFeatures, _objectCodec, out, packerConfig, reuseResourceInGenerator);
     }
 
     @Override

msgpack-jackson/src/main/java/org/msgpack/jackson/dataformat/MessagePackGenerator.java

@@ -22,6 +22,7 @@
 import com.fasterxml.jackson.core.base.GeneratorBase;
 import com.fasterxml.jackson.core.io.SerializedString;
 import com.fasterxml.jackson.core.json.JsonWriteContext;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.msgpack.core.MessagePack;
 import org.msgpack.core.MessagePacker;
 import org.msgpack.core.buffer.OutputStreamBufferOutput;
@@ -42,6 +43,8 @@ public class MessagePackGenerator
     private static final Charset DEFAULT_CHARSET = Charset.forName("UTF-8");
     private final MessagePacker messagePacker;
     private static ThreadLocal<OutputStreamBufferOutput> messageBufferOutputHolder = new ThreadLocal<OutputStreamBufferOutput>();
+    private final OutputStream output;
+    private final MessagePack.PackerConfig packerConfig;
     private LinkedList<StackItem> stack;
     private StackItem rootStackItem;
 
@@ -97,20 +100,35 @@ List<Object> getKeys()
         }
     }
 
-    public MessagePackGenerator(int features, ObjectCodec codec, OutputStream out, MessagePack.PackerConfig packerConfig)
+    public MessagePackGenerator(
+            int features,
+            ObjectCodec codec,
+            OutputStream out,
+            MessagePack.PackerConfig packerConfig,
+            boolean reuseResourceInGenerator)
             throws IOException
     {
         super(features, codec);
-        OutputStreamBufferOutput messageBufferOutput = messageBufferOutputHolder.get();
-        if (messageBufferOutput == null) {
-            messageBufferOutput = new OutputStreamBufferOutput(out);
+        this.output = out;
+
+        OutputStreamBufferOutput messageBufferOutput;
+        if (reuseResourceInGenerator) {
+            messageBufferOutput = messageBufferOutputHolder.get();
+            if (messageBufferOutput == null) {
+                messageBufferOutput = new OutputStreamBufferOutput(out);
+                messageBufferOutputHolder.set(messageBufferOutput);
+            }
+            else {
+                messageBufferOutput.reset(out);
+            }
         }
         else {
-            messageBufferOutput.reset(out);
+            messageBufferOutput = new OutputStreamBufferOutput(out);
         }
-        messageBufferOutputHolder.set(messageBufferOutput);
-
         this.messagePacker = packerConfig.newPacker(messageBufferOutput);
+
+        this.packerConfig = packerConfig;
+
         this.stack = new LinkedList<StackItem>();
     }
 
@@ -224,7 +242,11 @@ else if (v instanceof MessagePackExtensionType) {
             messagePacker.writePayload(extData);
         }
         else {
-            throw new IllegalArgumentException(v.toString());
+            messagePacker.flush();
+            MessagePackFactory messagePackFactory = new MessagePackFactory(packerConfig);
+            messagePackFactory.setReuseResourceInGenerator(false);
+            ObjectMapper objectMapper = new ObjectMapper(messagePackFactory);
+            output.write(objectMapper.writeValueAsBytes(v));
         }
     }
 

msgpack-jackson/src/test/java/org/msgpack/jackson/dataformat/MessagePackGeneratorTest.java

@@ -690,4 +690,25 @@ else if (keyName.equals("bigIntMap")) {
             }
         }
     }
+
+    @Test
+    public void testComplexTypeKey()
+            throws IOException
+    {
+        HashMap<TinyPojo, Integer> map = new HashMap<TinyPojo, Integer>();
+        TinyPojo pojo = new TinyPojo();
+        pojo.t = "foo";
+        map.put(pojo, 42);
+
+        ObjectMapper objectMapper = new ObjectMapper(new MessagePackFactory());
+        objectMapper.setSerializerFactory(new MessagePackSerializerFactory());
+        byte[] bytes = objectMapper.writeValueAsBytes(map);
+
+        MessageUnpacker unpacker = MessagePack.newDefaultUnpacker(bytes);
+        assertThat(unpacker.unpackMapHeader(), is(1));
+        assertThat(unpacker.unpackMapHeader(), is(1));
+        assertThat(unpacker.unpackString(), is("t"));
+        assertThat(unpacker.unpackString(), is("foo"));
+        assertThat(unpacker.unpackInt(), is(42));
+    }
 }