SOLR-3715: fix FastOutputStream, FastWriter when buffer sizes can change

git-svn-id: https://svn.apache.org/repos/asf/lucene/dev/trunk@1374480 13f79535-47bb-0310-9956-ffa450edef68

solr/core/src/java/org/apache/solr/update/MemOutputStream.java

@@ -0,0 +1,36 @@
+package org.apache.solr.update;
+
+import org.apache.solr.common.util.FastOutputStream;
+
+import java.io.IOException;
+import java.util.LinkedList;
+import java.util.List;
+
+/** @lucene.internal */
+public class MemOutputStream extends FastOutputStream {
+  public List<byte[]> buffers = new LinkedList<byte[]>();
+  public MemOutputStream(byte[] tempBuffer) {
+    super(null, tempBuffer, 0);
+  }
+
+  @Override
+  public void flush(byte[] arr, int offset, int len) throws IOException {
+    if (arr == buf && offset==0 && len==buf.length) {
+      buffers.add(buf);  // steal the buffer
+      buf = new byte[8192];
+    } else if (len > 0) {
+      byte[] newBuf = new byte[len];
+      System.arraycopy(arr, offset, newBuf, 0, len);
+      buffers.add(newBuf);
+    }
+  }
+
+  public void writeAll(FastOutputStream fos) throws IOException {
+    for (byte[] buffer : buffers) {
+      fos.write(buffer);
+    }
+    if (pos > 0) {
+      fos.write(buf, 0, pos);
+    }
+  }
+}

solr/core/src/java/org/apache/solr/update/TransactionLog.java

@@ -775,31 +775,3 @@ class ChannelFastInputStream extends FastInputStream {
 }
 
 
-class MemOutputStream extends FastOutputStream {
-  public List<byte[]> buffers = new LinkedList<byte[]>();
-  public MemOutputStream(byte[] tempBuffer) {
-    super(null, tempBuffer, 0);
-  }
-
-  @Override
-  public void flush(byte[] arr, int offset, int len) throws IOException {
-    if (arr == buf && offset==0 && len==buf.length) {
-      buffers.add(buf);  // steal the buffer
-      buf = new byte[8192];
-    } else if (len > 0) {
-      byte[] newBuf = new byte[len];
-      System.arraycopy(arr, offset, newBuf, 0, len);
-      buffers.add(newBuf);
-    }
-  }
-
-  public void writeAll(FastOutputStream fos) throws IOException {
-    for (byte[] buffer : buffers) {
-      fos.write(buffer);
-    }
-    if (pos > 0) {
-      fos.write(buf, 0, pos);
-    }
-  }
-}
-

solr/core/src/java/org/apache/solr/util/FastWriter.java

@@ -28,7 +28,7 @@ public class FastWriter extends Writer {
   // it won't cause double buffering.
   private static final int BUFSIZE = 8192;
   protected final Writer sink;
-  protected final char[] buf;
+  protected char[] buf;
   protected int pos;
 
   public FastWriter(Writer w) {
@@ -69,42 +69,64 @@ public class FastWriter extends Writer {
   }
 
   @Override
-  public void write(char cbuf[], int off, int len) throws IOException {
-    int space = buf.length - pos;
-    if (len < space) {
-      System.arraycopy(cbuf, off, buf, pos, len);
-      pos += len;
-    } else if (len<BUFSIZE) {
-      // if the data to write is small enough, buffer it.
-      System.arraycopy(cbuf, off, buf, pos, space);
+  public void write(char arr[], int off, int len) throws IOException {
+    for(;;) {
+      int space = buf.length - pos;
+
+      if (len <= space) {
+        System.arraycopy(arr, off, buf, pos, len);
+        pos += len;
+        return;
+      } else if (len > buf.length) {
+        if (pos>0) {
+          flush(buf,0,pos);  // flush
+          pos=0;
+        }
+        // don't buffer, just write to sink
+        flush(arr, off, len);
+        return;
+      }
+
+      // buffer is too big to fit in the free space, but
+      // not big enough to warrant writing on its own.
+      // write whatever we can fit, then flush and iterate.
+
+      System.arraycopy(arr, off, buf, pos, space);
       flush(buf, 0, buf.length);
-      pos = len-space;
-      System.arraycopy(cbuf, off+space, buf, 0, pos);
-    } else {
-      flush(buf,0,pos);  // flush
-      pos=0;
-      // don't buffer, just write to sink
-      flush(cbuf, off, len);
+      pos = 0;
+      off += space;
+      len -= space;
     }
   }
 
   @Override
   public void write(String str, int off, int len) throws IOException {
-    int space = buf.length - pos;
-    if (len < space) {
-      str.getChars(off, off+len, buf, pos);
-      pos += len;
-    } else if (len<BUFSIZE) {
-      // if the data to write is small enough, buffer it.
+    for(;;) {
+      int space = buf.length - pos;
+
+      if (len <= space) {
+        str.getChars(off, off+len, buf, pos);
+        pos += len;
+        return;
+      } else if (len > buf.length) {
+        if (pos>0) {
+          flush(buf,0,pos);  // flush
+          pos=0;
+        }
+        // don't buffer, just write to sink
+        flush(str, off, len);
+        return;
+      }
+
+      // buffer is too big to fit in the free space, but
+      // not big enough to warrant writing on its own.
+      // write whatever we can fit, then flush and iterate.
+
       str.getChars(off, off+space, buf, pos);
       flush(buf, 0, buf.length);
-      str.getChars(off+space, off+len, buf, 0);
-      pos = len-space;
-    } else {
-      flush(buf,0,pos);  // flush
-      pos=0;
-      // don't buffer, just write to sink
-      flush(str, off, len);
+      pos = 0;
+      off += space;
+      len -= space;
     }
   }
 

solr/core/src/test/org/apache/solr/util/TestFastOutputStream.java

@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.solr.util;
+
+import org.apache.lucene.util.LuceneTestCase;
+import org.apache.lucene.util._TestUtil;
+import org.apache.solr.update.MemOutputStream;
+
+import java.util.HashSet;
+import java.util.Locale;
+import java.util.Random;
+import java.util.Set;
+import java.util.TimeZone;
+
+public class TestFastOutputStream extends LuceneTestCase {
+
+  Random rand;
+  byte[] arr;
+
+  public void testRandomWrites() throws Exception {
+    rand = random();
+
+    arr = new byte[20000];
+    for (int i=0; i<arr.length; i++) {
+      arr[i] = (byte)rand.nextInt();
+    }
+
+    for (int i=0; i<1000; i++) {
+      doRandomWrites();
+    }
+
+  }
+
+  public void doRandomWrites() throws Exception {
+    int bufSize = ( rand.nextBoolean() ? rand.nextInt(10) : rand.nextInt(20000) )+1;
+    MemOutputStream out = new MemOutputStream(new byte[bufSize]);
+
+    int hash = 0;
+    long written = 0;
+    int iter = rand.nextInt(10)+1;
+    for (int i=0; i<iter; i++) {
+      int off = rand.nextInt(arr.length);
+      int len = off < arr.length ? rand.nextInt(arr.length - off) : 0;
+      out.write(arr, off, len);
+      hash = incHash(hash, arr, off, len);
+      written += len;
+
+      int pos = rand.nextInt(arr.length);
+      out.write(arr[pos]);
+      hash = incHash(hash, arr, pos, 1);
+      written += 1;
+    }
+
+    out.close();
+
+    int hash2 = 0;
+    for (byte[] buffer : out.buffers) {
+      hash2 = incHash(hash2, buffer, 0, buffer.length);
+    }
+
+    assertEquals(hash, hash2);
+    assertEquals(written, out.written());
+    assertEquals(written, out.size());
+  }
+
+
+  public int incHash(int hash, byte[] arr, int off, int len) {
+    for (int i=off; i<off+len; i++) {
+      hash = hash * 31 + arr[i];
+    }
+    return hash;
+  }
+}
+

solr/core/src/test/org/apache/solr/util/TestFastWriter.java

@@ -0,0 +1,128 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.solr.util;
+
+import org.apache.lucene.util.LuceneTestCase;
+import org.apache.solr.common.util.FastOutputStream;
+import org.apache.solr.update.MemOutputStream;
+
+import java.io.IOException;
+import java.util.LinkedList;
+import java.util.List;
+import java.util.Random;
+
+
+class MemWriter extends FastWriter {
+  public List<char[]> buffers = new LinkedList<char[]>();
+
+  Random r;
+  public MemWriter(char[] tempBuffer, Random r) {
+    super(null, tempBuffer, 0);
+    this.r = r;
+  }
+
+  @Override
+  public void flush(char[] arr, int offset, int len) throws IOException {
+    if (arr == buf && offset==0 && len==buf.length) {
+      buffers.add(buf);  // steal the buffer
+      buf = new char[r.nextInt(9000)+1];
+    } else if (len > 0) {
+      char[] newBuf = new char[len];
+      System.arraycopy(arr, offset, newBuf, 0, len);
+      buffers.add(newBuf);
+    }
+  }
+
+  @Override
+  public void flush(String str, int offset, int len) throws IOException {
+    if (len == 0) return;
+    buffers.add( str.substring(offset, offset+len).toCharArray() );
+  }
+}
+
+
+
+public class TestFastWriter extends LuceneTestCase {
+
+  Random rand;
+  char[] arr;
+  String s;
+
+  public void testRandomWrites() throws Exception {
+    rand = random();
+
+    arr = new char[20000];
+    for (int i=0; i<arr.length; i++) {
+      arr[i] = (char)rand.nextInt();
+    }
+    s = new String(arr);
+
+    for (int i=0; i<1000; i++) {
+      doRandomWrites();
+    }
+  }
+
+
+  public void doRandomWrites() throws Exception {
+    int bufSize = ( rand.nextBoolean() ? rand.nextInt(10) : rand.nextInt(20000) )+1;
+    MemWriter out = new MemWriter(new char[bufSize], rand);
+
+    int hash = 0;
+    long written = 0;
+    int iter = rand.nextInt(20)+1;
+    for (int i=0; i<iter; i++) {
+      int which = rand.nextInt(3);
+
+
+      int off = rand.nextInt(arr.length);
+      int len = off < arr.length ? rand.nextInt(arr.length - off) : 0;
+
+
+
+      if (which == 0) {
+        out.write(arr, off, len);
+      } else if (which == 1) {
+        out.write(s, off, len);
+      } else {
+        len = 1;
+        out.write(arr[off]);
+      }
+
+      hash = incHash(hash, arr, off, len);
+      written += len;
+    }
+
+    out.close();
+
+    int hash2 = 0;
+    for (char[] buffer : out.buffers) {
+      hash2 = incHash(hash2, buffer, 0, buffer.length);
+    }
+
+    assertEquals(hash, hash2);
+  }
+
+
+  public int incHash(int hash, char[] arr, int off, int len) {
+    for (int i=off; i<off+len; i++) {
+      hash = hash * 31 + arr[i];
+    }
+    return hash;
+  }
+}
+

solr/solrj/src/java/org/apache/solr/common/util/FastOutputStream.java

@@ -57,8 +57,8 @@ public class FastOutputStream extends OutputStream implements DataOutput {
 
   public void write(byte b) throws IOException {
     if (pos >= buf.length) {
-      flush(buf, 0, buf.length);
       written += pos;
+      flush(buf, 0, buf.length);
       pos=0;
     }
     buf[pos++] = b;
@@ -66,29 +66,40 @@ public class FastOutputStream extends OutputStream implements DataOutput {
 
   @Override
   public void write(byte arr[], int off, int len) throws IOException {
-    int space = buf.length - pos;
-    if (len < space) {
-      System.arraycopy(arr, off, buf, pos, len);
-      pos += len;
-    } else if (len<buf.length) {
-      // if the data to write is small enough, buffer it.
-      System.arraycopy(arr, off, buf, pos, space);
-      flush(buf, 0, buf.length);
-      written += buf.length;
-      pos = len-space;
-      System.arraycopy(arr, off+space, buf, 0, pos);
-    } else {
-      if (pos>0) {
-        flush(buf,0,pos);  // flush
-        written += pos;
-        pos=0;
+
+    for(;;) {
+      int space = buf.length - pos;
+
+      if (len <= space) {
+        System.arraycopy(arr, off, buf, pos, len);
+        pos += len;
+        return;
+      } else if (len > buf.length) {
+        if (pos>0) {
+          flush(buf,0,pos);  // flush
+          written += pos;
+          pos=0;
+        }
+        // don't buffer, just write to sink
+        flush(arr, off, len);
+        written += len;
+        return;
       }
-      // don't buffer, just write to sink
-      flush(arr, off, len);
-      written += len;            
+
+      // buffer is too big to fit in the free space, but
+      // not big enough to warrant writing on its own.
+      // write whatever we can fit, then flush and iterate.
+
+      System.arraycopy(arr, off, buf, pos, space);
+      written += buf.length;  // important to do this first, since buf.length can change after a flush!
+      flush(buf, 0, buf.length);
+      pos = 0;
+      off += space;
+      len -= space;
     }
   }
 
+
   /** reserve at least len bytes at the end of the buffer.
    * Invalid if len > buffer.length
    * @param len
@@ -182,8 +193,8 @@ public class FastOutputStream extends OutputStream implements DataOutput {
    */
   public void flushBuffer() throws IOException {
     if (pos > 0) {
-      flush(buf, 0, pos);
       written += pos;
+      flush(buf, 0, pos);
       pos=0;
     }
   }