Add suppport for mark/reset.

3 files changed, 91 insertions, 57 deletions

diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/FeedReaderFactory.java b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/FeedReaderFactory.java
index 419e61fa1c0..749f569d718 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/FeedReaderFactory.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/FeedReaderFactory.java
@@ -16,6 +16,7 @@ import java.io.InputStream;
  * @author dybis
  */
 public class FeedReaderFactory {
+    private static final int MARK_READLIMIT = 200;
 
     /**
      * Creates FeedReader
@@ -30,17 +31,17 @@ public class FeedReaderFactory {
             FeedParams.DataFormat dataFormat)  {
         switch (dataFormat) {
             case XML_UTF8:
+                byte [] peek = new byte[MARK_READLIMIT];
+                int bytesPeeked = 0;
                 try {
+                    if (inputStream.markSupported()) {
+                        inputStream.mark(MARK_READLIMIT);
+                        bytesPeeked = inputStream.read(peek);
+                        inputStream.reset();
+                    }
                     return new VespaXMLFeedReader(inputStream, docTypeManager);
                 } catch (Exception e) {
-                    byte [] peek;
-                    try {
-                        peek = new byte[Math.min(200, inputStream.available())];
-                        inputStream.read(peek);
-                    } catch (IOException  io) {
-                        peek = new byte [0];
-                    }
-                    throw new RuntimeException("Could not create VespaXMLFeedReader. First characters are: " + Utf8.toString(peek), e);
+                    throw new RuntimeException("Could not create VespaXMLFeedReader. First characters are: '" + Utf8.toString(peek, 0, bytesPeeked) + "'", e);
                 }
             case JSON_UTF8:
                 return new JsonFeedReader(inputStream, docTypeManager);
diff --git a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStream.java b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStream.java
index c9f255f026e..c8ae79deebd 100644
--- a/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStream.java
+++ b/vespaclient-container-plugin/src/main/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStream.java
@@ -13,6 +13,7 @@ public class ByteLimitedInputStream extends InputStream {
 
     private final InputStream wrappedStream;
     private int remaining;
+    private int remainingWhenMarked;
 
     public ByteLimitedInputStream(InputStream wrappedStream, int limit) {
         this.wrappedStream = wrappedStream;
@@ -78,4 +79,21 @@ public class ByteLimitedInputStream extends InputStream {
         }
     }
 
+    @Override
+    public synchronized void mark(int readlimit) {
+        wrappedStream.mark(readlimit);
+        remainingWhenMarked = remaining;
+    }
+
+    @Override
+    public synchronized void reset() throws IOException {
+        wrappedStream.reset();
+        remaining = remainingWhenMarked;
+    }
+
+    @Override
+    public boolean markSupported() {
+        return wrappedStream.markSupported();
+    }
+
 }
diff --git a/vespaclient-container-plugin/src/test/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStreamTestCase.java b/vespaclient-container-plugin/src/test/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStreamTestCase.java
index 3aa3cdcb3a8..3dd8145ec73 100644
--- a/vespaclient-container-plugin/src/test/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStreamTestCase.java
+++ b/vespaclient-container-plugin/src/test/java/com/yahoo/vespa/http/server/util/ByteLimitedInputStreamTestCase.java
@@ -8,8 +8,8 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
 
-import static org.hamcrest.core.Is.is;
-import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 /**
  * @author <a href="mailto:einarmr@yahoo-inc.com">Einar M R Rosenvinge</a>
@@ -29,63 +29,78 @@ public class ByteLimitedInputStreamTestCase {
     public void requireThatBasicsWork() throws IOException {
         ByteLimitedInputStream stream = create("abcdefghijklmnopqr".getBytes(StandardCharsets.US_ASCII), 9);
 
-        assertThat(stream.available(), is(9));
-        assertThat(stream.read(), is(97));
-        assertThat(stream.available(), is(8));
-        assertThat(stream.read(), is(98));
-        assertThat(stream.available(), is(7));
-        assertThat(stream.read(), is(99));
-        assertThat(stream.available(), is(6));
-        assertThat(stream.read(), is(100));
-        assertThat(stream.available(), is(5));
-        assertThat(stream.read(), is(101));
-        assertThat(stream.available(), is(4));
-        assertThat(stream.read(), is(102));
-        assertThat(stream.available(), is(3));
-        assertThat(stream.read(), is(103));
-        assertThat(stream.available(), is(2));
-        assertThat(stream.read(), is(104));
-        assertThat(stream.available(), is(1));
-        assertThat(stream.read(), is(105));
-        assertThat(stream.available(), is(0));
-        assertThat(stream.read(), is(-1));
-        assertThat(stream.available(), is(0));
-        assertThat(stream.read(), is(-1));
-        assertThat(stream.available(), is(0));
-        assertThat(stream.read(), is(-1));
-        assertThat(stream.available(), is(0));
-        assertThat(stream.read(), is(-1));
-        assertThat(stream.available(), is(0));
-        assertThat(stream.read(), is(-1));
-        assertThat(stream.available(), is(0));
+        assertEquals(9, stream.available());
+        assertEquals(97, stream.read());
+        assertEquals(8, stream.available());
+        assertEquals(98, stream.read());
+        assertEquals(7, stream.available());
+        assertEquals(99, stream.read());
+        assertEquals(6, stream.available());
+        assertEquals(100, stream.read());
+        assertEquals(5, stream.available());
+        assertEquals(101, stream.read());
+        assertEquals(4, stream.available());
+        assertEquals(102, stream.read());
+        assertEquals(3, stream.available());
+        assertEquals(103, stream.read());
+        assertEquals(2, stream.available());
+        assertEquals(104, stream.read());
+        assertEquals(1, stream.available());
+        assertEquals(105, stream.read());
+        assertEquals(0, stream.available());
+        assertEquals(-1, stream.read());
+        assertEquals(0, stream.available());
+        assertEquals(-1, stream.read());
+        assertEquals(0, stream.available());
+        assertEquals(-1, stream.read());
+        assertEquals(0, stream.available());
+        assertEquals(-1, stream.read());
+        assertEquals(0, stream.available());
+        assertEquals(-1, stream.read());
+        assertEquals(0, stream.available());
     }
 
     @Test
     public void requireThatChunkedReadWorks() throws IOException {
         ByteLimitedInputStream stream = create("abcdefghijklmnopqr".getBytes(StandardCharsets.US_ASCII), 9);
 
-        assertThat(stream.available(), is(9));
+        assertEquals(9, stream.available());
         byte[] toBuf = new byte[4];
-        assertThat(stream.read(toBuf), is(4));
-        assertThat(toBuf[0], is((byte) 97));
-        assertThat(toBuf[1], is((byte) 98));
-        assertThat(toBuf[2], is((byte) 99));
-        assertThat(toBuf[3], is((byte) 100));
-        assertThat(stream.available(), is(5));
+        assertEquals(4, stream.read(toBuf));
+        assertEquals(97, toBuf[0]);
+        assertEquals(98, toBuf[1]);
+        assertEquals(99, toBuf[2]);
+        assertEquals(100, toBuf[3]);
+        assertEquals(5, stream.available());
 
-        assertThat(stream.read(toBuf), is(4));
-        assertThat(toBuf[0], is((byte) 101));
-        assertThat(toBuf[1], is((byte) 102));
-        assertThat(toBuf[2], is((byte) 103));
-        assertThat(toBuf[3], is((byte) 104));
-        assertThat(stream.available(), is(1));
+        assertEquals(4, stream.read(toBuf));
+        assertEquals(101, toBuf[0]);
+        assertEquals(102, toBuf[1]);
+        assertEquals(103, toBuf[2]);
+        assertEquals(104, toBuf[3]);
+        assertEquals(1, stream.available());
 
-        assertThat(stream.read(toBuf), is(1));
-        assertThat(toBuf[0], is((byte) 105));
-        assertThat(stream.available(), is(0));
+        assertEquals(1, stream.read(toBuf));
+        assertEquals(105, toBuf[0]);
+        assertEquals(0, stream.available());
 
-        assertThat(stream.read(toBuf), is(-1));
-        assertThat(stream.available(), is(0));
+        assertEquals(-1, stream.read(toBuf));
+        assertEquals(0, stream.available());
+    }
+
+    @Test
+    public void requireMarkWorks() throws IOException {
+        InputStream stream = create("abcdefghijklmnopqr".getBytes(StandardCharsets.US_ASCII), 9);
+        assertEquals(97, stream.read());
+        assertTrue(stream.markSupported());
+        stream.mark(5);
+        assertEquals(98, stream.read());
+        assertEquals(99, stream.read());
+        stream.reset();
+        assertEquals(98, stream.read());
+        assertEquals(99, stream.read());
+        assertEquals(100, stream.read());
+        assertEquals(101, stream.read());
     }
 
 }