[BEAM-171] Fix JAXBCoder in the nested context

sdks/java/core/src/main/java/com/google/cloud/dataflow/sdk/coders/JAXBCoder.java

@@ -19,6 +19,8 @@
 
 import com.google.cloud.dataflow.sdk.util.CloudObject;
 import com.google.cloud.dataflow.sdk.util.Structs;
+import com.google.cloud.dataflow.sdk.util.VarInt;
+import com.google.common.io.ByteStreams;
 
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
@@ -71,13 +73,19 @@ public void encode(T value, OutputStream outStream, Context context)
  JAXBContext jaxbContext = JAXBContext.newInstance(jaxbClass);
  jaxbMarshaller = jaxbContext.createMarshaller();
  }
-
- jaxbMarshaller.marshal(value, new FilterOutputStream(outStream) {
- // JAXB closes the underyling stream so we must filter out those calls.
- @Override
- public void close() throws IOException {
+ if (!context.isWholeStream) {
+ try {
+ long size = getEncodedElementByteSize(value, Context.OUTER);
+ // record the number of bytes the XML consists of so when reading we only read the encoded
+ // value
+ VarInt.encode(size, outStream);
+ } catch (Exception e) {
+ throw new CoderException(
+ "An Exception occured while trying to get the size of an encoded representation", e);
  }
- });
+ }
+
+ jaxbMarshaller.marshal(value, new CloseIgnoringOutputStream(outStream));
  } catch (JAXBException e) {
  throw new CoderException(e);
  }
@@ -91,13 +99,13 @@ public T decode(InputStream inStream, Context context) throws CoderException, IO
  jaxbUnmarshaller = jaxbContext.createUnmarshaller();
  }
 
+ InputStream stream = inStream;
+ if (!context.isWholeStream) {
+ long limit = VarInt.decodeLong(inStream);
+ stream = ByteStreams.limit(inStream, limit);
+ }
  @SuppressWarnings("unchecked")
- T obj = (T) jaxbUnmarshaller.unmarshal(new FilterInputStream(inStream) {
- // JAXB closes the underyling stream so we must filter out those calls.
- @Override
- public void close() throws IOException {
- }
- });
+ T obj = (T) jaxbUnmarshaller.unmarshal(new CloseIgnoringInputStream(stream));
  return obj;
  } catch (JAXBException e) {
  throw new CoderException(e);
@@ -109,6 +117,30 @@ public String getEncodingId() {
  return getJAXBClass().getName();
  }
 
+ private static class CloseIgnoringInputStream extends FilterInputStream {
+
+ protected CloseIgnoringInputStream(InputStream in) {
+ super(in);
+ }
+
+ @Override
+ public void close() {
+ // Do nothing. JAXB closes the underlying stream so we must filter out those calls.
+ }
+ }
+
+ private static class CloseIgnoringOutputStream extends FilterOutputStream {
+
+ protected CloseIgnoringOutputStream(OutputStream out) {
+ super(out);
+ }
+
+ @Override
+ public void close() throws IOException {
+ // JAXB closes the underlying stream so we must filter out those calls.
+ }
+ }
+
  ////////////////////////////////////////////////////////////////////////////////////
  // JSON Serialization details below
 

sdks/java/core/src/test/java/com/google/cloud/dataflow/sdk/coders/JAXBCoderTest.java

@@ -19,11 +19,18 @@
 
 import com.google.cloud.dataflow.sdk.testing.CoderProperties;
 import com.google.cloud.dataflow.sdk.util.CoderUtils;
+import com.google.common.collect.ImmutableList;
+
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.List;
+
 import javax.xml.bind.annotation.XmlRootElement;
 
 /** Unit tests for {@link JAXBCoder}. */
@@ -79,13 +86,62 @@ public boolean equals(Object obj) {
  }
 
  @Test
- public void testEncodeDecode() throws Exception {
+ public void testEncodeDecodeOuter() throws Exception {
  JAXBCoder<TestType> coder = JAXBCoder.of(TestType.class);
 
  byte[] encoded = CoderUtils.encodeToByteArray(coder, new TestType("abc", 9999));
  Assert.assertEquals(new TestType("abc", 9999), CoderUtils.decodeFromByteArray(coder, encoded));
  }
 
+ @Test
+ public void testEncodeDecodeNested() throws Exception {
+ JAXBCoder<TestType> jaxbCoder = JAXBCoder.of(TestType.class);
+ TestCoder nesting = new TestCoder(jaxbCoder);
+
+ byte[] encoded = CoderUtils.encodeToByteArray(nesting, new TestType("abc", 9999));
+ Assert.assertEquals(
+ new TestType("abc", 9999), CoderUtils.decodeFromByteArray(nesting, encoded));
+ }
+
+ /**
+ * A coder that surrounds the value with two values, to demonstrate nesting.
+ */
+ private static class TestCoder extends StandardCoder<TestType> {
+ private final JAXBCoder<TestType> jaxbCoder;
+ public TestCoder(JAXBCoder<TestType> jaxbCoder) {
+ this.jaxbCoder = jaxbCoder;
+ }
+
+ @Override
+ public void encode(TestType value, OutputStream outStream, Context context)
+ throws CoderException, IOException {
+ Context subContext = context.nested();
+ VarIntCoder.of().encode(3, outStream, subContext);
+ jaxbCoder.encode(value, outStream, subContext);
+ VarLongCoder.of().encode(22L, outStream, subContext);
+ }
+
+ @Override
+ public TestType decode(InputStream inStream, Context context)
+ throws CoderException, IOException {
+ Context subContext = context.nested();
+ VarIntCoder.of().decode(inStream, subContext);
+ TestType result = jaxbCoder.decode(inStream, subContext);
+ VarLongCoder.of().decode(inStream, subContext);
+ return result;
+ }
+
+ @Override
+ public List<? extends Coder<?>> getCoderArguments() {
+ return ImmutableList.of(jaxbCoder);
+ }
+
+ @Override
+ public void verifyDeterministic() throws NonDeterministicException {
+ jaxbCoder.verifyDeterministic();
+ }
+ }
+
  @Test
  public void testEncodable() throws Exception {
  CoderProperties.coderSerializable(JAXBCoder.of(TestType.class));