[Test Only][SPARK-6235][CORE]Address various 2G limits

common/network-common/src/main/java/org/apache/spark/network/buffer/AbstractReferenceCounted.java

@@ -0,0 +1,119 @@
+/*
+ * 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.spark.network.buffer;
+
+import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
+
+import io.netty.util.internal.PlatformDependent;
+
+/**
+ * Abstract base class for classes wants to implement {@link ReferenceCounted}.
+ */
+public abstract class AbstractReferenceCounted implements ReferenceCounted {
+
+  private static final AtomicIntegerFieldUpdater<AbstractReferenceCounted> refCntUpdater;
+
+  static {
+    AtomicIntegerFieldUpdater<AbstractReferenceCounted> updater =
+        PlatformDependent.newAtomicIntegerFieldUpdater(AbstractReferenceCounted.class, "refCnt");
+    if (updater == null) {
+      updater = AtomicIntegerFieldUpdater.newUpdater(AbstractReferenceCounted.class, "refCnt");
+    }
+    refCntUpdater = updater;
+  }
+
+  private volatile int refCnt = 1;
+
+  @Override
+  public int refCnt() {
+    return refCnt;
+  }
+
+  /**
+   * An unsafe operation intended for use by a subclass that sets the reference count of the buffer directly
+   */
+  protected void setRefCnt(int refCnt) {
+    this.refCnt = refCnt;
+  }
+
+  @Override
+  public ReferenceCounted retain() {
+    doRetain(1);
+    return this;
+  }
+
+  @Override
+  public ReferenceCounted retain(int increment) {
+    if (increment <= 0) {
+      throw new IllegalArgumentException("increment: " + increment + " (expected: > 0)");
+    }
+    doRetain(increment);
+    return this;
+  }
+
+  protected ReferenceCounted doRetain(int increment) {
+    for (; ; ) {
+      int refCnt = this.refCnt;
+      final int nextCnt = refCnt + increment;
+
+      // Ensure we not resurrect (which means the refCnt was 0) and also that we encountered an overflow.
+      if (nextCnt <= increment) {
+        throw new IllegalReferenceCountException(refCnt, increment);
+      }
+      if (refCntUpdater.compareAndSet(this, refCnt, nextCnt)) {
+        break;
+      }
+    }
+    return this;
+  }
+
+  @Override
+  public boolean release() {
+    return doRelease(1);
+  }
+
+  @Override
+  public boolean release(int decrement) {
+    if (decrement <= 0) {
+      throw new IllegalArgumentException("decrement: " + decrement + " (expected: > 0)");
+    }
+    return doRelease(decrement);
+  }
+
+  protected boolean doRelease(int decrement) {
+    for (; ; ) {
+      int refCnt = this.refCnt;
+      if (refCnt < decrement) {
+        throw new IllegalReferenceCountException(refCnt, -decrement);
+      }
+
+      if (refCntUpdater.compareAndSet(this, refCnt, refCnt - decrement)) {
+        if (refCnt == decrement) {
+          deallocate();
+          return true;
+        }
+        return false;
+      }
+    }
+  }
+
+  /**
+   * Called once {@link #refCnt()} is equals 0.
+   */
+  protected abstract void deallocate();
+}

common/network-common/src/main/java/org/apache/spark/network/buffer/Allocator.java

@@ -0,0 +1,24 @@
+/*
+ * 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.spark.network.buffer;
+
+import java.nio.ByteBuffer;
+
+public interface Allocator {
+  ByteBuffer allocate(int len);
+}

common/network-common/src/main/java/org/apache/spark/network/buffer/ChunkedByteBuffer.java

@@ -0,0 +1,87 @@
+/*
+ * 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.spark.network.buffer;
+
+import java.io.Externalizable;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.nio.ByteBuffer;
+
+import io.netty.buffer.ByteBuf;
+
+public interface ChunkedByteBuffer extends Externalizable, ReferenceCounted {
+
+  /**
+   * This size of this buffer, in bytes.
+   */
+  long size();
+
+  /**
+   * Write this buffer to a outputStream.
+   */
+  void writeFully(OutputStream outputStream) throws IOException;
+
+  /**
+   * Wrap this buffer to view it as a Netty ByteBuf.
+   */
+  ByteBuf toNetty();
+
+  /**
+   * Copy this buffer into a new byte array.
+   *
+   * @throws UnsupportedOperationException if this buffer's size exceeds the maximum array size.
+   */
+  byte[] toArray();
+
+  /**
+   * Copy this buffer into a new ByteBuffer.
+   *
+   * @throws UnsupportedOperationException if this buffer's size exceeds the max ByteBuffer size.
+   */
+  ByteBuffer toByteBuffer();
+
+  InputStream toInputStream();
+
+  /**
+   * Creates an input stream to read data from this ChunkedByteBuffer.
+   *
+   * @param dispose if true, [[dispose()]] will be called at the end of the stream
+   *                in order to close any memory-mapped files which back this buffer.
+   */
+  InputStream toInputStream(boolean dispose);
+
+  /**
+   * Make a copy of this ChunkedByteBuffer, copying all of the backing data into new buffers.
+   * The new buffer will share no resources with the original buffer.
+   */
+  ChunkedByteBuffer copy();
+
+  /**
+   * Get duplicates of the ByteBuffers backing this ChunkedByteBuffer.
+   */
+  ByteBuffer[] toByteBuffers();
+
+  ChunkedByteBuffer slice(long offset, long length);
+
+  ChunkedByteBuffer duplicate();
+
+  ChunkedByteBuffer retain();
+
+  ChunkedByteBuffer retain(int increment);
+}

common/network-common/src/main/java/org/apache/spark/network/buffer/ChunkedByteBufferOutputStream.java

@@ -0,0 +1,60 @@
+/*
+ * 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.spark.network.buffer;
+
+import java.io.OutputStream;
+import com.google.common.base.Preconditions;
+
+import org.apache.spark.network.buffer.netty.ChunkedByteBuffOutputStreamImpl;
+import org.apache.spark.network.buffer.nio.ChunkedByteBufferOutputStreamImpl;
+
+public abstract class ChunkedByteBufferOutputStream extends OutputStream {
+
+  protected final boolean isDirect;
+  protected final int chunkSize;
+
+  /**
+   * An OutputStream that writes to fixed-size chunks of byte arrays.
+   *
+   * @param chunkSize size of each chunk, in bytes.
+   */
+  public ChunkedByteBufferOutputStream(int chunkSize, boolean isDirect) {
+    this.chunkSize = chunkSize;
+    this.isDirect = isDirect;
+    Preconditions.checkArgument(chunkSize > 0);
+  }
+
+  public abstract long size();
+
+  public abstract ChunkedByteBuffer toChunkedByteBuffer();
+
+  public static ChunkedByteBufferOutputStream newInstance(int chunkSize, boolean isDirect) {
+    return new ChunkedByteBuffOutputStreamImpl(chunkSize, isDirect);
+  }
+
+  public static ChunkedByteBufferOutputStream newInstance(int chunkSize) {
+    return newInstance(chunkSize, false);
+  }
+
+  public static ChunkedByteBufferOutputStream newInstance(int chunkSize, Allocator alloc) {
+    return new ChunkedByteBufferOutputStreamImpl(chunkSize, false, alloc);
+  }
+
+  public static ChunkedByteBufferOutputStream newInstance() {
+    return newInstance(4 * 1024);
+  }
+}

common/network-common/src/main/java/org/apache/spark/network/buffer/ChunkedByteBufferUtil.java

@@ -0,0 +1,104 @@
+/*
+ * 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.spark.network.buffer;
+
+import java.io.DataInput;
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.MappedByteBuffer;
+
+import com.google.common.io.ByteStreams;
+import io.netty.buffer.ByteBuf;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import sun.nio.ch.DirectBuffer;
+
+import org.apache.spark.network.buffer.netty.ChunkedByteBufImpl;
+import org.apache.spark.network.buffer.nio.ChunkedByteBufferImpl;
+
+public class ChunkedByteBufferUtil {
+  private static final Logger logger = LoggerFactory.getLogger(ChunkedByteBufferUtil.class);
+
+  public static void dispose(ByteBuffer buffer) {
+    if (buffer != null && buffer instanceof MappedByteBuffer) {
+      logger.trace("Unmapping" + buffer);
+      if (buffer instanceof DirectBuffer) {
+        DirectBuffer directBuffer = (DirectBuffer) buffer;
+        if (directBuffer.cleaner() != null) directBuffer.cleaner().clean();
+      }
+    }
+  }
+
+  public static ChunkedByteBuffer wrap() {
+    return new ChunkedByteBufferImpl();
+  }
+
+  public static ChunkedByteBuffer wrap(ByteBuffer chunk) {
+    ByteBuffer[] chunks = new ByteBuffer[1];
+    chunks[0] = chunk;
+    return wrap(chunks);
+  }
+
+  public static ChunkedByteBuffer wrap(ByteBuf chunk) {
+    ByteBuf[] chunks = new ByteBuf[1];
+    chunks[0] = chunk;
+    return wrap(chunks);
+  }
+
+  public static ChunkedByteBuffer wrap(ByteBuffer[] chunks) {
+    return new ChunkedByteBufferImpl(chunks);
+  }
+
+  public static ChunkedByteBuffer wrap(ByteBuf[] chunks) {
+    return new ChunkedByteBufImpl(chunks);
+  }
+
+  public static ChunkedByteBuffer wrap(byte[] array) {
+    return wrap(array, 0, array.length);
+  }
+
+  public static ChunkedByteBuffer wrap(byte[] array, int offset, int length) {
+    return wrap(ByteBuffer.wrap(array, offset, length));
+  }
+
+  public static ChunkedByteBuffer wrap(InputStream in, int chunkSize) throws IOException {
+    ChunkedByteBufferOutputStream out = ChunkedByteBufferOutputStream.newInstance(chunkSize);
+    ByteStreams.copy(in, out);
+    out.close();
+    return out.toChunkedByteBuffer();
+  }
+
+  public static ChunkedByteBuffer wrap(
+      DataInput from, int chunkSize, long len) throws IOException {
+    ChunkedByteBufferOutputStream out = ChunkedByteBufferOutputStream.newInstance(chunkSize);
+    final int BUF_SIZE = Math.min(chunkSize, 4 * 1024);
+    byte[] buf = new byte[BUF_SIZE];
+    while (len > 0) {
+      int r = (int) Math.min(len, BUF_SIZE);
+      from.readFully(buf, 0, r);
+      out.write(buf, 0, r);
+      len -= r;
+    }
+    out.close();
+    return out.toChunkedByteBuffer();
+  }
+
+  public static ChunkedByteBuffer wrap(InputStream in) throws IOException {
+    return wrap(in, 32 * 1024);
+  }
+}