Add suppport for mark/reset.

2 files changed, 69 insertions, 4 deletions

diff --git a/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java b/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java
index 748c2951a6a..1662ed5b46a 100644
--- a/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java
+++ b/jdisc_core/src/main/java/com/yahoo/jdisc/handler/UnsafeContentInputStream.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jdisc.handler;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.nio.ByteBuffer;
 import java.util.Objects;
@@ -19,6 +20,8 @@ public class UnsafeContentInputStream extends InputStream {
 
     private final ReadableContentChannel content;
     private ByteBuffer buf = ByteBuffer.allocate(0);
+    private byte [] marked;
+    private int readSinceMarked;
 
     /**
      * <p>Constructs a new ContentInputStream that reads from the given {@link ReadableContentChannel}.</p>
@@ -37,7 +40,15 @@ public class UnsafeContentInputStream extends InputStream {
         if (buf == null) {
             return -1;
         }
-        return ((int)buf.get()) & 0xFF;
+        byte b = buf.get();
+        if (marked != null) {
+            if (readSinceMarked < marked.length) {
+                marked[readSinceMarked++] = b;
+            } else {
+                marked = null;
+            }
+        }
+        return ((int)b) & 0xFF;
     }
 
     @Override
@@ -79,4 +90,28 @@ public class UnsafeContentInputStream extends InputStream {
 
         }
     }
+
+    @Override
+    public synchronized void mark(int readlimit) {
+        marked = new byte[readlimit];
+        readSinceMarked = 0;
+    }
+
+    @Override
+    public synchronized void reset() throws IOException {
+        if (marked == null) {
+            throw new IOException("mark has not been called, or too much has been read since marked.");
+        }
+        ByteBuffer newBuf = ByteBuffer.allocate(readSinceMarked + buf.remaining());
+        newBuf.put(marked, 0, readSinceMarked);
+        newBuf.put(buf);
+        newBuf.flip();
+        buf = newBuf;
+        marked = null;
+    }
+
+    @Override
+    public boolean markSupported() {
+        return true;
+    }
 }
diff --git a/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java b/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java
index c00fab6cb56..c96450c1bd2 100644
--- a/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java
+++ b/jdisc_core/src/test/java/com/yahoo/jdisc/handler/UnsafeContentInputStreamTestCase.java
@@ -1,6 +1,7 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.jdisc.handler;
 
+import com.yahoo.text.Utf8;
 import org.junit.Test;
 
 import java.io.BufferedReader;
@@ -8,11 +9,11 @@ import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.ByteBuffer;
-import java.util.concurrent.Future;
 
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
 
 /**
  * @author Simon Thoresen Hult
@@ -32,7 +33,37 @@ public class UnsafeContentInputStreamTestCase {
         assertNull(reader.readLine());
     }
 
-    @SuppressWarnings("deprecation")
+    @Test
+    public void testMark() throws IOException {
+        BufferedContentChannel channel = new BufferedContentChannel();
+        FastContentWriter writer = new FastContentWriter(channel);
+        writer.write("Hello ");
+        writer.write("World!");
+        writer.close();
+
+        InputStream stream = asInputStream(channel);
+        assertTrue(stream.markSupported());
+        int first = stream.read();
+        assertEquals('H', first);
+        stream.mark(10);
+        byte [] buf = new byte[8];
+        stream.read(buf);
+        assertEquals("ello Wor", Utf8.toString(buf));
+        stream.reset();
+        stream.mark(5);
+        buf = new byte [9];
+        stream.read(buf);
+        assertEquals("ello Worl", Utf8.toString(buf));
+        try {
+            stream.reset();
+            fail("UnsafeContentInputStream.reset expected to fail when your read past readLimit.");
+        } catch (IOException e) {
+            assertEquals("mark has not been called, or too much has been read since marked.", e.getMessage());
+        } catch (Throwable t) {
+            fail("Did not expect " + t);
+        }
+    }
+
     @Test
     public void requireThatCompletionsAreCalledWithDeprecatedContentWriter() throws IOException {
         BufferedContentChannel channel = new BufferedContentChannel();
@@ -63,7 +94,6 @@ public class UnsafeContentInputStreamTestCase {
         assertTrue(writer.isDone());
     }
 
-    @SuppressWarnings("deprecation")
     @Test
     public void requireThatCloseDrainsStreamWithDeprecatedContentWriter() {
         BufferedContentChannel channel = new BufferedContentChannel();